Can I auto-sharding a variable?

[horse6]

I have a function which will return 2 variables(z and W), and I want to shard them automatically. I know it could be sharded by using jax.lax.with_sharding_constraint or out_shardings parameter of jit. But how to shardW automatically as following example(which fails, and manual methods are commented).

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.lax import with_sharding_constraint
from jax.experimental import mesh_utils
import jax
import jax.numpy as jnp

device_mesh = mesh_utils.create_device_mesh((2, 4))
mesh = Mesh(devices=device_mesh, axis_names=('X', 'Y'))

def mesh_sharding(pspec: PartitionSpec) -> NamedSharding:
  return NamedSharding(mesh, pspec)

def f(x):
  W = jnp.ones((4, 4))
  # W = with_sharding_constraint(W, mesh_sharding(PartitionSpec('X', 'Y')))
  z = jnp.dot(x, W)
  z = with_sharding_constraint(z, mesh_sharding(PartitionSpec('X', None)))
  return z, W

x = jnp.ones((8, 4))
x_sharding = mesh_sharding(PartitionSpec('X', None))
x = jax.device_put(x, x_sharding)

#jax.debug.visualize_array_sharding(x)
# z, W = jax.jit(f, in_shardings=x_sharding, out_shardings=(None, mesh_sharding(PartitionSpec('X', 'Y'))))(x)
z, W = jax.jit(f, in_shardings=x_sharding, out_shardings=None)(x)
jax.debug.visualize_array_sharding(z)
jax.debug.visualize_array_sharding(W)