tensorplay.autograd API

Compute the sum of gradients of given tensors with respect to graph leaves.

The gradients are accumulated in the leaves’ .grad. If any of tensors is non-scalar and requires gradient, grad_tensors gives the “vector” of the vector-Jacobian product for each of them (None for scalars and tensors that do not need one).

Parameters:

• tensors (Sequence[Tensor] or Tensor) – Tensors whose derivative is computed.

• grad_tensors (Sequence[Tensor or None] or Tensor, optional) – The “vector” in the vector-Jacobian product, usually gradients w.r.t. each element of corresponding tensors.

• retain_graph (bool, optional) – If False, the graph used to compute the grad will be freed. Defaults to the value of create_graph.

• create_graph (bool, optional) – If True, the graph of the derivative is constructed, allowing higher order derivative products.

• grad_variables (Sequence[Tensor or None] or Tensor, optional) – Deprecated spelling of grad_tensors.

• inputs (Sequence[Tensor] or Tensor or dict, optional) – Tensors w.r.t. which the gradient will be accumulated into .grad. All other tensors are ignored, and only the part of the graph that leads to these runs. A non-leaf named here keeps its gradient as if it had called retain_grad(). A dict (e.g. dict(model.named_parameters())) names its values.

Compute and return the sum of gradients of outputs with respect to the inputs.

grad_outputs should be a sequence of length matching output containing the “vector” in vector-Jacobian product, usually the pre-computed gradients w.r.t. each of the outputs. If an output doesn’t require_grad, then the gradient can be None).

Parameters:

• outputs (sequence of Tensor or GradientEdge) – outputs of the differentiated function.

• inputs (sequence of Tensor or GradientEdge) – Inputs w.r.t. which the gradient will be returned (and not accumulated into .grad).

• grad_outputs (sequence of Tensor) – The “vector” in the vector-Jacobian product. Usually gradients w.r.t. each output. None values can be specified for scalar Tensors or ones that don’t require grad. If a None value would be acceptable for all grad_tensors, then this argument is optional. Default: None.

• retain_graph (bool, optional) – If False, the graph used to compute the grad will be freed. Note that in nearly all cases setting this option to True is not needed and often can be worked around in a much more efficient way. Defaults to the value of create_graph.

• create_graph (bool, optional) – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Default: False.

• only_inputs (bool, optional) – Deprecated and ignored; gradients are only ever returned for inputs.

• allow_unused (Optional[bool], optional) – If False, specifying inputs that were not used when computing outputs (and therefore their grad is always zero) is an error. Defaults to the value of materialize_grads.

• is_grads_batched (bool, optional) – If True, the first dimension of each tensor in grad_outputs is a batch dimension: one vector-Jacobian product is computed per entry, in a single tensorplay.vmap() over the backward pass. Default: False.

• materialize_grads (bool, optional) – If True, the gradient of an unused input is returned as zeros instead of None. Default: False.

Returns a tuple with one gradient per input, or a dict keyed like inputs when inputs is a dict.

Compute the dot product between the Jacobian of the given function at the point given by the inputs and a vector v.

Parameters:

• func (function) – a Python function that takes Tensor inputs and returns a tuple of Tensors or a Tensor.

• inputs (tuple of Tensors or Tensor) – inputs to the function func.

• v (tuple of Tensors or Tensor) – The vector for which the Jacobian vector product is computed. Must be the same size as the input of func. This argument is optional when the input to func contains a single element and (if it is not provided) will be set as a Tensor containing a single 1.

• create_graph (bool, optional) – If True, both the output and result will be computed in a differentiable way. Note that when strict is False, the result can not require gradients or be disconnected from the inputs. Defaults to False.

• strict (bool, optional) – If True, an error will be raised when we detect that there exists an input such that all the outputs are independent of it. If False, we return a Tensor of zeros as the jvp for said inputs, which is the expected mathematical value. Defaults to False.

Returns:

tuple with:

func_output (tuple of Tensors or Tensor): output of func(inputs)

jvp (tuple of Tensors or Tensor): result of the dot product with the same shape as the output.

Return type:

output (tuple)

Example

>>> def exp_reducer(x):
...     return x.exp().sum(dim=1)
>>> inputs = tensorplay.rand(4, 4)
>>> v = tensorplay.ones(4, 4)
>>> jvp(exp_reducer, inputs, v)
(tensor([6.3090, 4.6742, 7.9114, 8.2106]),
 tensor([6.3090, 4.6742, 7.9114, 8.2106]))

>>> def adder(x, y):
...     return 2 * x + 3 * y
>>> inputs = (tensorplay.rand(2), tensorplay.rand(2))
>>> v = (tensorplay.ones(2), tensorplay.ones(2))
>>> jvp(adder, inputs, v)
(tensor([2.2399, 2.5005]),
 tensor([5., 5.]))

mode (str, optional): “forward” runs func once with the dual-tensor

engine and propagates tangents op by op (the default; func must be written with operators supported by forward-mode). “reversed” computes the jvp via the double backwards trick and works for any twice-differentiable func.