Fix Auto-sharding Bug for Residual blocks.

[gaga1313]

Fix: Prevent sharding conflicts between non-primary axes in auto_shard.py

The current heuristic in auto_shard.py only checks for sharding conflicts against the single dominant model_axis. It does not account for conflicts that can arise between newly assigned non-primary axes.

Problem: When testing with a Residual Block, the existing code incorrectly caused the weight parameter of the linear3 layer to be sharded on both dimensions using the model_axis.

Solution: To fix this bug, I have introduced a dim_sharded flag. This flag tracks whether a dimension has already been sharded while assigning the sharding axis to each parameter dimension, preventing erroneous double-sharding.

class ResidualBlockModel(nnx.Module):

  def __init__(self, rngs: nnx.Rngs):
    self.linear1 = nnx.Linear(128, 64, rngs=rngs)
    self.linear2 = nnx.Linear(64, 32, rngs=rngs)
    self.linear3 = nnx.Linear(96, 8, rngs=rngs)

  def __call__(self, x):
    z1 = self.linear1(x)
    z2 = self.linear2(z1)
    z1_z2 = jnp.concatenate([z2, z1], axis=-1)
    z3 = self.linear3(z1_z2)
    return z3

Reproduction script:

import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'

import jax
import jax.numpy as jnp
from parallax.sharding import auto_shard
from flax import nnx

# Create a mesh for setup
mesh = jax.make_mesh(
    (4, 2),
    ('data', 'model'),
    axis_types=(jax.sharding.AxisType.Auto,) * 2,
)

with jax.set_mesh(mesh):
  def model_fn(params, x):
    w1, b1, w2, b2, w3, b3 = params
    z1 = jnp.dot(x, w1) + b1  # (32, 64)
    z2 = jnp.dot(z1, w2) + b2  # (32, 32)
    z1_z2 = jnp.concatenate([z2, z1], axis=-1)  # (32, 96)
    z3 = jnp.dot(z1_z2, w3) + b3  # (32, 8)
    return z3

  # Define shapes
  x = jnp.ones((32, 128))
  w1 = jnp.ones((128, 64))
  b1 = jnp.ones((64,))
  w2 = jnp.ones((64, 32))
  b2 = jnp.ones((32,))
  w3 = jnp.ones((96, 8))
  b3 = jnp.ones((8,))

  params = (w1, b1, w2, b2, w3, b3)
  params_name = ("W1", "b1", "W2", "b2", "W3", "b3")


  # Let's run get_shardings and see final assignments
  (params_shd, x_shd), out_shd = auto_shard.get_shardings(model_fn, params, x)
  print("\n--- SHARDING ASSIGNMENTS ---")
  print("Params Sharding Assignments:")
  for i, p_shd in enumerate(params_shd):
    print(f"  {params_name[i]}: {p_shd}")
  print(f"Input x Sharding Assignment: {x_shd}")
  print(f"Output z3 Sharding Assignment: {out_shd}")