A reusable value-and-grad over a JS-built loss graph.
loss(primals) receives the input arrays in argument order and must
return a single SCALAR MlxArray (the loss). The graph it builds is
differentiated w.r.t. the inputs named in argIdx. apply(primals)
returns the loss value and exactly argIdx.length gradients (one per
differentiated input, in argIdx order).
The closure must not throw; if the user loss throws, the error is
captured and surfaced from apply(). The MlxArrays handed to loss
are owned by mlx — do NOT dispose them; the returned loss IS disposed
of by mlx's graph.
Run the loss + backward on primals (in argument order). Returns the
scalar loss and one gradient per differentiated input. Caller owns all
returned arrays (dispose them).
A reusable value-and-grad over a JS-built loss graph.
loss(primals)receives the input arrays in argument order and must return a single SCALAR MlxArray (the loss). The graph it builds is differentiated w.r.t. the inputs named inargIdx.apply(primals)returns the loss value and exactlyargIdx.lengthgradients (one per differentiated input, inargIdxorder).The closure must not throw; if the user
lossthrows, the error is captured and surfaced fromapply(). The MlxArrays handed tolossare owned by mlx — do NOT dispose them; the returned loss IS disposed of by mlx's graph.