pmap got inconsistent sizes for array axes to be mapped:

[mwitiderrick]

What am I doing wrong here?

import jax.numpy as jnp
import flaxmodels as fm

num_classes = 2
dtype = jnp.float32
model = fm.ResNet50(output='log_softmax', pretrained=None, num_classes=num_classes, dtype=dtype)

import functools
@functools.partial(jax.pmap)
def create_train_state(rng):
  """Creates initial `TrainState`."""
  return TrainState.create(apply_fn = model.apply,params = variables['params'],tx = optax.adam(0.01,0.9),batch_stats = variables['batch_stats'])

@functools.partial(jax.pmap, axis_name='ensemble')
def apply_model(state, images, labels):
  def loss_fn(params,batch_stats):
    logits,batch_stats = model.apply({'params': params,'batch_stats': batch_stats},images, train=True,rngs={'dropout': jax.random.PRNGKey(0)}, mutable=['batch_stats'])
    one_hot = jax.nn.one_hot(labels, 2)
    loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot).mean()
    return loss, (logits, batch_stats)

  (loss, (logits, batch_stats)), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params,state.batch_stats)
  probs = jax.lax.pmean(jax.nn.softmax(logits), axis_name='ensemble')

  accuracy = jnp.mean(jnp.argmax(probs, -1) == labels)
  return grads,loss, accuracy

@jax.pmap
def update_model(state, grads):
  return state.apply_gradients(grads=grads)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [55], in <cell line: 3>()
      1 start = time.time()
      2 num_epochs = 10
----> 3 state, epoch_loss, epoch_accuracy, testing_accuracy, testing_loss = train_one_epoch(state, train_loader,num_epochs)
      4 print("Total time: ", time.time() - start, "seconds")

Input In [51], in train_one_epoch(state, dataloader, num_epochs)
     17 train_loss = np.mean(epoch_loss)
     18 train_accuracy = np.mean(epoch_accuracy)
---> 20 _, test_loss, test_accuracy = jax_utils.unreplicate(apply_model(state, test_images, test_labels))
     21 testing_accuracy.append(test_accuracy)
     22 testing_loss.append(test_loss)

    [... skipping hidden 4 frame]

File ~/miniforge3/lib/python3.9/site-packages/jax/_src/api.py:1652, in _mapped_axis_size(tree, vals, dims, name, kws)
   1650 sizes = [x.shape[d] if d is not None else None for x, d in zip(vals, dims)]
   1651 sizes = tree_unflatten(tree, sizes)
-> 1652 raise ValueError(msg.format(f"the tree of axis sizes is:\n{sizes}")) from None

ValueError: pmap got inconsistent sizes for array axes to be mapped:
the tree of axis sizes is:

[cgarciae]

Hey @mwitiderrick! Its a bit hard to tell from your example but my guess is that the inputs of apply_model have different shapes for their first dimension. See jax-ml/jax#9933.

[mwitiderrick]

Thanks @cgarciae thanks, the me take a look.

Update: The issue was that I was not replicating the test data. I think this error message can be improved. cc @marcvanzee