Docs: please clarify how to vmap nnx.Module over batch dimension about flax HOT 3 OPEN

Hey! In the NNX Transforms section I think we make it clear that you must use the nnx.* version of the transform when using NNX objects. Is there something we should improve here?

from flax.

I think currently you want this:

(loss_values, (y_preds, the_energies, accuracies)), grads = nnx.vmap(
  loss_value_and_grad_fn,
  in_axes=(None, 0, 0),
  state_axes={...: None},  # <<== add this
)

But in the near future of #3963 , you will be able to do

(loss_values, (y_preds, the_energies, accuracies)), grads = nnx.experimental.vmap(
  loss_value_and_grad_fn,
  in_axes=(None, 0, 0), # <<== just works
)

and hopefully experimental becomes the norm

from flax.

yes, you're right, the docs are clear to use the nnx transforms

I hit an issue where, I just want to map a function over each item in a batch, and when using the vmap transform, the first time I call the transformed function, it receives a single slice of the batch, which is what I was going for, but then when it gets called again, it receives the whole batch at once, which caused it to crash due to a shape issue in a call to "concatenate"

I reckon it's using an XLA version of the same function on the second invocation

For what to improve, could the docs have a section dedicated to advice about how to work with batches of data?

I was looking for more focused troubleshooting docs to deal with shape bugs in handling batches with vmap, or advice on how to apply a function over slices of a batch independently without having to change the function to handle the whole batch at once

(it's an EBM / energy based model, so I don't want to inappropriately mix energy gradients across the members of the same batch, and it breaks if I try to apply the single-item function to the whole batch)

Might just be my general inexperience with jax and a vmap issue and not a nnx issue; maybe i am missing something about how vmap works

from flax.