How JAX Shards a Computation Across a Mesh

I have eight devices. I want a matmul whose contracting dimension is split across some of them. So I write the matmul normally, say where the two operands live, and compile: mesh = Mesh(np.array(jax.devices()).reshape(4, 2), ('data', 'model')) xs = jax.device_put(x, NamedSharding(mesh, P('data', 'model'))) # K split on 'model' Ws = jax.device_put(W, NamedSharding(mesh, P('model', None))) ...