Reverse-mode AD via mlx_vjp: given a function fn(primals) -> outputs and a
cotangent per output, returns the cotangent w.r.t. EVERY primal. This is the
natural primitive for segmented backward (segment forward + dh -> [dh_in,
...dLoRA]) — no surrogate scalar and no value_and_grad closure. Unlike
ValueAndGrad it returns a vjp for ALL primals (not a selected argnum subset).
fn must build its output graph from the input arrays in argument order and
return them (most often a single array). It must not throw; a throw is
captured and surfaced from apply(). The MlxArrays handed to fn are owned by
mlx — do NOT dispose them. Lifecycle/house-rules mirror ValueAndGrad.
Run the vjp on primals with cotangents (one per output of fn).
Returns the outputs (fn(primals)) and one cotangent per primal (in primal
order). Caller owns all returned arrays.
Reverse-mode AD via mlx_vjp: given a function
fn(primals) -> outputsand a cotangent per output, returns the cotangent w.r.t. EVERY primal. This is the natural primitive for segmented backward (segment forward + dh -> [dh_in, ...dLoRA]) — no surrogate scalar and no value_and_grad closure. Unlike ValueAndGrad it returns a vjp for ALL primals (not a selected argnum subset).fnmust build its output graph from the input arrays in argument order and return them (most often a single array). It must not throw; a throw is captured and surfaced from apply(). The MlxArrays handed tofnare owned by mlx — do NOT dispose them. Lifecycle/house-rules mirror ValueAndGrad.