spatial sharding added

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -21,6 +21,8 @@
 import jax.numpy as jnp
 from jax import tree_util
 from flax import nnx
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
 from ...configuration_utils import ConfigMixin
 from ..modeling_flax_utils import FlaxModelMixin, get_activation
 from ... import common_types
@@ -144,6 +146,14 @@ def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) ->
     else:
       x_padded = x
 
+    if self.mesh is not None:
+      # Shard height dimension (index 2) along 'context' axis
+      # Shape is (Batch, Time, Height, Width, Channels)
+      # We only shard if the dimension is divisible by the mesh size to avoid XLA errors
+      if x_padded.shape[2] % self.mesh.shape["context"] == 0:
+        sharding = NamedSharding(self.mesh, P(None, None, "context", None, None))
+        x_padded = jax.lax.with_sharding_constraint(x_padded, sharding)
+
     out = self.conv(x_padded)
     return out