API Documentation

Batched maps

adv_jax_math.batch_jacfwd(fun[, argnums, ...])	Jacobian of fun evaluated column-by-column using forward-mode AD.
adv_jax_math.batch_jacrev(fun[, argnums, ...])	Jacobian of fun evaluated row-by-row using reverse-mode AD.
adv_jax_math.batch_map(fun, fun_input, /[, ...])	Compute chunk_reduction(fun(fun_input)) in batches.
adv_jax_math.batch_vectorize(pyfunc, *[, ...])	Define a vectorized function with broadcasting and batching.
adv_jax_math.batch_vmap(f, /[, in_axes, ...])	Behaves like vmap but uses scan to chunk the computations in smaller chunks.

Sparse differentiation

adv_jax_math.sparse_pullback(fn, y, /[, ...])	Compute chunk_reduction(fn(fun_input)) in batches with sparse pullbacks.
adv_jax_math.sparse_pullback_map(fn, y, *[, ...])	Wrapper for sparsity exploiting pullback.