tensorplay.func.grad

tensorplay.func.grad(func: Callable, argnums: int | tuple[int, ...] = 0, has_aux: bool = False) → Callable

Returns a function computing the gradient of func.

func must return a scalar tensor; the returned function has the same signature and returns the gradient with respect to argnums. Because it is again an ordinary function of the same inputs, grad(grad(f)) is the second derivative.

Parameters:

• func (Callable) – a function returning a single-element tensor.

• argnums (int or Tuple[int]) – which positional arguments to differentiate with respect to. Default: 0.

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

Example

>>> x = tensorplay.randn([])
>>> grad(tensorplay.sin)(x)