Skip to content

PartialWithParameters

svetlanna.PartialWithParameters 

PartialWithParameters(
    function: Callable[
        Concatenate[_Input, _Params], _Output
    ],
    *args: args,
    **kwargs: kwargs
)

Bases: Module, Generic[_Input, _Output]

Wrap an arbitrary function with trainable keyword arguments.

This behaves like functools.partial, but only keyword arguments are supported. Use this wrapper when you want keyword arguments to be registered as trainable parameters or as buffers (for tensor-valued constants). This is especially useful for multi-device workflows, since parameters and buffers move with the module.

Parameters:

  • function (Callable[Concatenate[_Input, _Params], _Output]) –

    Arbitrary function with parameters.

  • *args (args, default: () ) –

    Positional arguments (not supported; must be empty).

  • **kwargs (kwargs, default: {} ) –

    Keyword arguments for the function. Values are registered as parameters, buffers, or plain attributes depending on their type.

Examples:

Suppose you have a function that describes a nonlinear response and has trainable parameters. See the example in NonlinearElement.