Function Gradient

python/tvm/relax/transform/transform.py:55–229

Reverse-mode automatic differentiation. This pass will differentiate one function in the IRModule. Now the input function must have only one dataflow block (ConvertToDataflow may need to be called first). For a given function specified by `func_name`, it generates a new function with t

(
    func_name: str, require_grads: Var | list[Var] | None = None, target_index: int = 0
)

def Gradient(
    func_name: str, require_grads: Var | list[Var] | None = None, target_index: int = 0
) -> tvm.ir.transform.Pass:
    """Reverse-mode automatic differentiation.

    This pass will differentiate one function in the IRModule. Now the input function must have only
    one dataflow block (ConvertToDataflow may need to be called first).

    For a given function specified by `func_name`, it generates a new function with the name
    `func_name + "_adjoint"`. The new function computes the gradient of the **differentiation
    target** with respect to the arguments specified by `require_grads` of the original function.

    If the function has only one return value, the return value will be specified as target. If the
    function has more than one return values, the target will be specified as the target_index-th
    return value. The target must be a scalar (0-dim tensor).

    The new function will be like:

    .. code-block:: python

        @R.function
        def main_adjoint(original_parameters):
            with R.dataflow():
                # the bindings of the original function
                ...
                # calculating the gradients
                ...
                R.output(original_outputs, grad_1, grad_2, ...)
            return (original_return_value, (grad_1, grad_2, ...))

    This AD pass also supports checkpointing as described in
    "Training deep nets with sublinear memory cost." - Chen, Tianqi, et al. (2016).
    See tvm.relax.testing.nn.checkpoint for more details.

    Parameters
    ----------
    func_name : str
        The name of the specific function.

    require_grads : Optional[Union[relax.Var, List[relax.Var]]]
        The relax variables whose adjoints is needed. Must be parameters of the given function and
        should not be duplicate. If it is not specified, adjoints of all parameters would be
        computed.

    target_index : int
        If the specified function has more than one return values, specify the index of the return
        value as the target. If it is not specified, the first return value will be the target.

    Returns
    -------
    ret : tvm.ir.transform.Pass
        The Pass.

    Examples
    --------
    The following code shows how to use this pass:

    .. code-block:: python