tensorplay.func.jvp

tensorplay.func.jvp(func: Callable, primals: Any, tangents: Any, *, strict: bool = False, has_aux: bool = False)

Evaluates func at primals together with its directional derivative along tangents.

Parameters:

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

• primals (Tensors) – a tuple of positional arguments to evaluate at.

• tangents (Tensors) – the direction to differentiate along. Must have the same python structure, shapes and dtypes as primals.

• strict (bool) – raise instead of returning zeros when the output turns out to be independent of the inputs.

• has_aux (bool) – whether func returns (output, aux).

Returns:

(output, jvp_out), or (output, jvp_out, aux) when has_aux.

Example

>>> x = tensorplay.randn(5)
>>> out, tangent_out = jvp(tensorplay.sin, (x,), (tensorplay.ones(5),))