module

func

9 members

lucid.func

`lucid.func` — functional transforms (reference-framework func-compatible API).

Provides composable, higher-order transformations over Lucid functions:

vmap — vectorise over a batch dimension (C++ vmapped dispatch)
grad — gradient transform (returns a function)
grad_and_value — gradient + primal output
vjp — vector-Jacobian product (functional style)
jvp — Jacobian-vector product (exact via double-backward)
jacrev — reverse-mode Jacobian
jacfwd — forward-mode Jacobian via jvp
hessian — second-order derivative matrix
linearize — linearisation: (primals_out, tangent_fn)

All transforms compose: vmap(grad(fn)), grad(vmap(fn)), etc.

Notes

vmap does not loop over batch elements in Python. It moves the batch dimension to the front and calls the function once with the full batched tensor. The underlying C++ engine ops then run on the complete batch:

GPU (Metal/MLX): a single Metal dispatch handles all batch elements in parallel via MLX's native batched kernels.
CPU (Accelerate): BLAS calls (SGEMM, cblas_sgemm_batch, etc.) operate on the full batched tensor with SIMD parallelism.

This gives real vectorisation — not a sequential Python loop.

Reduction ops inside the vmapped function must use an explicit dim argument. An unqualified .sum() / .mean() reduces over the batch dimension as well, which is almost never desired.
In-place ops inside a vmapped function are not supported.
randomness='error' (default) raises if the function calls any random op.

Functions

fnvmap→ Callable

Vectorise func over a batch axis to produce a batched function.

fngrad→ Callable

Build a function returning the gradient of func.

fngrad_and_value→ Callable

Return a function computing both the gradient and the primal value.

fnvjp→ tuple

Compute the vector-Jacobian product of func at primals.

fnjvp→ tuple

Compute the Jacobian-vector product of func at primals.

fnlinearize→ tuple

Linearise func around primals for reuse across many tangents.

fnjacrev→ Callable

Build a function returning the reverse-mode Jacobian of func.

fnjacfwd→ Callable

Build a function returning the forward-mode Jacobian of func.

fnhessian→ Callable

Build a function returning the Hessian of a scalar-valued func.