flax
flax copied to clipboard
nnx.pmap documentation seems to be for vmap
This is correct. Sorry it was adapted very quickly from vmap for an experiment.
Will update soon.
No problem, thank you @cgarciae
will there be any updates soon...?