xmap properly compiles but then raises an axis size error

[cmunna0052]

I am not sure if this is a bug or simply a misunderstanding on my part. I am trying to shard a feed-forward network with arbitrary hidden layers using xmap. My goal is the following: Each weight matrix is split up into 10 groups of columns (corresponding to 10 devices). At each step, the current x-vector is sent in full to all 10 devices. On each device it is multiplied into 1/10th of the weight matrix columns, added to 1/10th of bias vector, and then put through the activation to get 1/10th of the next x-vector. Then, jax.lax.all_gather is called to recombine the sharded x-vectors. This repeats until the end. Here is the code:

def init_params(layer_neurons):
    params = {}
    newKey = jax.random.PRNGKey(42)
    in_axes = {}
    axis_resources = {}
    i=0
    for n_in, n_out in zip(layer_neurons[:-1],layer_neurons[1:]):
        newKey,subKey = jax.random.split(newKey)
        params['linear'+str(i)] = dict(w = jax.random.normal(key=subKey,shape=(n_in,n_out))*jnp.sqrt(2/n_in),
                           b = jnp.ones(shape = (n_out)))
        in_axes['linear'+str(i)] = {'b':{0:'hidden'+str(i)},'w':{1:'hidden'+str(i)}}
        axis_resources['hidden'+str(i)]='x'
        i+=1
    return params,in_axes,axis_resources

def forward(params,x):
    for i in range(len(params)-1): 
        x = jax.nn.sigmoid(x@(params['linear'+str(i)]['w']) + (params['linear'+str(i)]['b']))
        x = jax.lax.all_gather(x,'hidden'+str(i)) 
    curr = params['linear'+str(i+1)]
    x = x@curr['w'] + curr['b']
    return jax.lax.all_gather(x,'hidden'+str(i+1))

params, param_axes, axis_resources = init_params([784,100,10])
#param_axes = {'linear0': {'b': {0: 'hidden0'}, 'w': {1: 'hidden0'}}, 'linear1': {'b': {0: 'hidden1'}, 'w': {1: 'hidden1'}}} 
#axis_resources = {'hidden0': 'x', 'hidden1': 'x'}
device_array = np.array(jax.devices())
in_axes = (param_axes,{0:'batch'})
#batch_size = 20; inputs = normal data from MNIST

with jax.sharding.Mesh(device_array,('x',)):
    result = jax.experimental.maps.xmap(
        forward,in_axes=in_axes,out_axes={0:'batch'},axis_resources=axis_resources
    )(params,inputs)

When I run this code, it completes the compilation of forward (I can tell by adding print statements throughout), but then in the actual calculation it fails at
assert axis_size == frame_size, "axis size doesn't match"
Does anyone see what is going on here? Is there a better/easier way to do this? I can see one solution where I split up the parameter matrices manually, but I thought it would be possible to have xmap handle that for me.