Where is the batch dimension with jax.pmap?

I was using jax.pmap to start my training on 4 1080Ti GPUs. I got the training step function

step_fn = get_train_step()

and then I used pmap,

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

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). The batch dimension was lost.

How can I solve this problem?