tensorplay.func.vjp

tensorplay.func.vjp(func: Callable, *primals, has_aux: bool = False)

Evaluates func at primals and returns a function computing the vector-Jacobian product.

Parameters:

• func (Callable) – a Python function taking one or more tensor arguments.

• primals (Tensors) – positional arguments to evaluate func at. The returned function differentiates with respect to all of them.

• has_aux (bool) – whether func returns (output, aux), where aux is carried through undifferentiated.

Returns:

(output, vjp_fn), or (output, vjp_fn, aux) when has_aux. vjp_fn takes a cotangent with the same structure as output and returns the gradients with respect to primals.

Example

>>> x = tensorplay.randn(5)
>>> out, vjp_fn = vjp(tensorplay.sin, x)
>>> (grad,) = vjp_fn(tensorplay.ones_like(out))