flax
allow captures in nnx.grad
#5484
Open

Loading