jax.core.InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (8, 8, 3600) and (8, 10, 8, 3600)

[buttercutter]

How to modify the following reshape() code to solve the dimension runtime error ?
Do I need padding for the batch input array ?

# add a first dimension over gradient_accumulation_steps for minibatch slices
batch = jax.tree_map(
    lambda x: x.reshape(
        training_args.per_device_gradient_accumulation_steps, training_args.per_device_train_batch_size, *x.shape[1::]
    ),
    batch,
)

Note: In my code, training_args.per_device_gradient_accumulation_steps = 10 , and training_args.per_device_train_batch_size = 8 and batch has shape of (8, 3600)

Reference : huggingface/transformers#20855 (comment)

Traceback (most recent call last):
  File "run_summarization_flax.py", line 1338, in <module>
    main()
  File "run_summarization_flax.py", line 1264, in main
    state, train_metric = p_train_step(state, batch)
  File "/home/moe/.local/lib/python3.8/site-packages/chex/_src/fake.py", line 175, in wrapped_fn
    output = vmapped_fn(*call_args)
  File "run_summarization_flax.py", line 1173, in train_step
    batch = jax.tree_map(
  File "run_summarization_flax.py", line 1174, in <lambda>
    lambda x: x.reshape(
  File "/home/moe/.local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 793, in _reshape
    return lax.reshape(a, newshape, None)
jax.core.InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (8, 8, 3600) and (8, 10, 8, 3600)

[buttercutter]

I proceed with my own python script with the following slight modification to get it past the dimension runtime error.

Note that the -1 in the reshape operation means that the size of the last dimension will be inferred from the size of x and the other dimensions. Hence the following modification will reshape batch to have shape (8, 10, 3600)

# add a first dimension over gradient_accumulation_steps for minibatch slices
batch = jax.tree_map(
    lambda x: x.reshape(
        training_args.per_device_train_batch_size, training_args.per_device_gradient_accumulation_steps, -1 #*x.shape[1::]
    ),
    batch,
)