How to use jax.lax.scan correctly?

[noeltong]

Hello, everyone! I was using jax.pmap to start my training on 4 1080Ti GPUs. I got the training step function by

step_fn = get_train_step()

and then I combined the step function with jax.lax.scan,

p_train_step = jax.pmap(functools.partial(jax.lax.scan, train_step_fn),
                                          axis_name='device',
                                          in_axes=0,
                                          out_axes=0)

The question appeared when I executed the p_train_step function. The shape of the batch was (4, 8, 128, 128, 1), and I expected the axis 0 to be mapped by pmap. However, when I executed

(_, p_train_state, p_model_state), ploss = p_train_step((next_rng, p_train_state, p_model_state), batch)

where next_rng was 4 random keys generated by jax.random.split, p_train_state and p_model_state were replicated by flax.jax_utils.replicate, the shape of the batch received by the step_fn came out to be (128, 128, 1), thus the error came from

B, H, W, C = batch.shape

The batch dimension was lost. However, the problem disappeared when I ran

p_train_step = jax.pmap(train_step_fn,
                      axis_name='device',
                      in_axes=0,
                      out_axes=0)

then everything went well. How can I solve the problem?