Transforms

eval(*args[, retain_graph])	Evaluate an array or tree of array.
grad(fun[, argnums])	Returns a function which computes the gradient of fun.
value_and_grad(fun[, argnums])	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.