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.
simplify(*args)	Simplify the graph that computes the arrays.