Transforms#

eval(*args)

Evaluate an array or tree of array.

grad(fun[, argnums, argnames])

Returns a function which computes the gradient of fun.

value_and_grad(fun[, argnums, argnames])

Returns a function which computes the value and gradient of fun.

jvp(fun, primals, tangents)

Compute the Jacobian-vector product.

vjp(fun, primals, cotangents)

Compute the vector-Jacobian product.

vmap(fun[, in_axes, out_axes])

Returns a vectorized version of fun.