Fall back to sharding the height in the WAN VAE.

The WAN VAE always sharded the width of its activations over the
vae_spatial mesh axis. This works for the default 1280x720 landscape
size, but users may also want to run inference on portrait videos
720x1280. The VAE downsamples their width to 720, 360, 180 and 90, which
isn't always divisible by vae_spatial (e.g., 90 for 4).

This change adds a _spatial_sharding(mesh, shape) helper that shards height
if width **is not divisible** by vae_spatial, but height is divisible.
Videos with a divisible width keep the current behavior.

Author: ninatu

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -48,6 +48,21 @@ def _update_cache(cache, idx, value):
   return cache[:idx] + (value,) + cache[idx + 1 :]
 
 
+def _spatial_sharding(mesh: Optional[jax.sharding.Mesh], shape: Sequence[int]) -> Optional[NamedSharding]:
+  """Returns the sharding of a (N, D, H, W, C) array over the vae_spatial axis.
+
+  Shards the width, or the height if only the height is divisible by
+  vae_spatial. If neither is, the width is still sharded and XLA pads it.
+  """
+  if mesh is None or "vae_spatial" not in mesh.shape:
+    return None
+  _, _, height, width, _ = shape
+  num_shards = mesh.shape["vae_spatial"]
+  if width % num_shards != 0 and height % num_shards == 0:
+    return NamedSharding(mesh, P("redundant", None, "vae_spatial", None, None))
+  return NamedSharding(mesh, P("redundant", None, None, "vae_spatial", None))
+
+
 # Helper to ensure kernel_size, stride, padding are tuples of 3 integers
 def _canonicalize_tuple(x: Union[int, Sequence[int]], rank: int, name: str) -> Tuple[int, ...]:
   """Canonicalizes a value to a tuple of integers."""
@@ -121,7 +136,7 @@ def __init__(
 
   def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) -> jax.Array:
     if self.mesh is not None and "vae_spatial" in self.mesh.shape:
-      spatial_sharding = NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
+      spatial_sharding = _spatial_sharding(self.mesh, x.shape)
       if spatial_sharding is not None:
         x = jax.lax.with_sharding_constraint(x, spatial_sharding)
 
@@ -1169,12 +1184,6 @@ def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache):
     iter_ = 1 + ((t - 1 + CHUNK_SIZE - 1) // CHUNK_SIZE) if t > 1 else 1
     enc_feat_map = feat_cache._enc_feat_map
 
-    spatial_sharding = (
-        NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
-        if self.mesh is not None and "vae_spatial" in self.mesh.shape
-        else None
-    )
-
     def finalize(out, enc_feat_map):
       feat_cache._enc_feat_map = enc_feat_map
       enc = self.quant_conv(out)
@@ -1184,6 +1193,7 @@ def finalize(out, enc_feat_map):
     with jax.named_scope("AutoencoderKLWan_encode_chunk_0"):
       chunk_0 = x[:, :1, ...]
       out_0, enc_feat_map, _ = self.encoder(chunk_0, feat_cache=enc_feat_map, feat_idx=0)
+      spatial_sharding = _spatial_sharding(self.mesh, out_0.shape)
       if spatial_sharding is not None:
         out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding)
 
@@ -1280,25 +1290,23 @@ def _decode(
     x = self.post_quant_conv(z)
 
     dec_feat_map = feat_cache._feat_map
-    spatial_sharding = (
-        NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
-        if self.mesh is not None and "vae_spatial" in self.mesh.shape
-        else None
-    )
+    spatial_sharding = _spatial_sharding(self.mesh, x.shape)
 
     # First chunk (i=0)
     with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):
+      chunk_in_0 = x[:, 0:1, ...]
       if spatial_sharding is not None:
-        chunk_in_0 = jax.lax.with_sharding_constraint(x[:, 0:1, ...], spatial_sharding)
+        chunk_in_0 = jax.lax.with_sharding_constraint(chunk_in_0, spatial_sharding)
       out_0, dec_feat_map, _ = self.decoder(chunk_in_0, feat_cache=dec_feat_map, feat_idx=0)
       if spatial_sharding is not None:
         out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding)
 
     if iter_ > 1:
       # Run chunk 1 outside scan to properly form the cache shape
       with jax.named_scope("AutoencoderKLWan_decode_chunk_1"):
+        chunk_in_1 = x[:, 1:2, ...]
         if spatial_sharding is not None:
-          chunk_in_1 = jax.lax.with_sharding_constraint(x[:, 1:2, ...], spatial_sharding)
+          chunk_in_1 = jax.lax.with_sharding_constraint(chunk_in_1, spatial_sharding)
         out_chunk_1, dec_feat_map, _ = self.decoder(chunk_in_1, feat_cache=dec_feat_map, feat_idx=0)
         if spatial_sharding is not None:
           out_chunk_1 = jax.lax.with_sharding_constraint(out_chunk_1, spatial_sharding)