def METHOD_NAME():
  y, variables = init(f)(random.PRNGKey(0), jnp.zeros((8, 4)))
  spec = meta.get_partition_spec(variables)
  shardings = jax.tree_map(lambda s: sharding.NamedSharding(mesh, s), spec)
  variables = jax.lax.with_sharding_constraint(variables, shardings)
  return variables