code: def shard_axis(x, axis_size, axis_name):
      # in_shape = x.shape
      assert x.shape[0] % axis_size == 0
  
      x = x.reshape((axis_size, -1) + x.shape[1:])
  
      x = x[jax.lax.axis_index(axis_name)]
      # print("shard out", x.shape, "in", in_shape)
  
      # assert np.prod(x.shape) * axis_size == np.prod(in_shape)
  
      return x
  
imports: import jax
  import jax.numpy
  from jax.experimental.pjit import with_sharding_constraint
  from optax import AdditiveWeightDecayState, GradientTransformation, OptState
label_desc: Use image node_id 18 for calling a global function with example usage: shard_axis(x, axis_size, axis_name) and returns: x