[ninatu]

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.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[gemini-code-assist]

Code Review

This pull request introduces a helper function _spatial_sharding to dynamically determine the sharding of a 5D array over the vae_spatial axis, replacing hardcoded sharding logic in __call__, _encode, and _decode. The review feedback recommends defensively checking that the input shape is 5D in _spatial_sharding to prevent unpacking errors, and addresses a potential NameError in _decode when spatial_sharding is None.

[gemini-code-assist]

There is a critical bug in _decode where if spatial_sharding is None (e.g., when running on a single device without a mesh), chunk_in_0 (line 1298) and chunk_in_1 (line 1307) are never defined. This will raise a NameError when they are passed to self.decoder on lines 1299 and 1308.

To fix this, you should define chunk_in_0 and chunk_in_1 unconditionally and apply the sharding constraint conditionally. For example:

# 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(chunk_in_0, spatial_sharding)
  out_0, dec_feat_map, _ = self.decoder(chunk_in_0, feat_cache=dec_feat_map, feat_idx=0)

And similarly for chunk_in_1.

[gemini-code-assist]

To prevent potential ValueError during unpacking, we should defensively check that the input shape is indeed 5D (i.e., has a length of 5) before attempting to unpack it. If it is not 5D, we can safely return None to skip applying the sharding constraint.

Suggested change

	if mesh is None or "vae_spatial" not in mesh.shape:
	return None
	_, _, height, width, _ = shape
	if mesh is None or "vae_spatial" not in mesh.shape or len(shape) != 5:
	return None
	_, _, height, width, _ = shape

[entrpn]

Thanks @ninatu! Supporting portrait 720x1280 without XLA padding on vae_spatial makes a lot of sense.

One important sharding/perf catch when tracing 720x1280 (H=1280, W=720) through the 4 spatial pyramid levels, plus a few consistency items:

1. W <-> H sharding thrashing (all-to-alls) across pyramid levels and scan_fn carries on 720x1280 (vae_spatial=4 or 8):

   • For 720x1280 (H=1280, W=720) with vae_spatial=4, levels 0–2 (1280x720, 640x360, 320x180) all satisfy width % 4 == 0, while only the latent bottleneck (160x90) has 90 % 4 != 0 and 160 % 4 == 0 (similarly, for vae_spatial=8, 1280x720 and 640x360 satisfy width % 8 == 0 while 320x180 and 160x90 flip to height).
   • Because WanCausalConv3d.__call__ calls _spatial_sharding(self.mesh, x.shape) on its local activation shape while _encode (from out_0.shape) and _decode (from post_quant_conv(z).shape) compute a single height-sharded spatial_sharding:
     • Every WanCausalConv3d at 1280x720, 640x360, and 320x180 constrains x to width-sharded,
     • Meanwhile scan_fn in _encode (lines 1248, 1252–1257) and _decode (lines 1350–1357) constrains chunk, out_chunk, and every 5D cache tensor in next_feat_map across all pyramid levels back to height-sharded, triggering repeated all-to-all transposes between H and W on every temporal scan chunk.
   • Could we choose the spatial sharding axis once from the bottleneck/latent (H_lat, W_lat) (i.e. (H // 8, W // 8) in _encode and (z.shape[2], z.shape[3]) in _decode) and apply that same axis across all pyramid levels in encoder/decoder/WanCausalConv3d so 720x1280 stays height-sharded end-to-end with zero all-to-alls?

2. AutoencoderKLWan2p2 (src/maxdiffusion/models/wan/autoencoder_kl_wan_2p2.py:1344, 1397):

   • AutoencoderKLWan2p2 imports WanCausalConv3d from autoencoder_kl_wan.py, so its internal convs will now flip to height sharding at 160x90 while its _encode/_decode still hardcode P(None, None, None, "vae_spatial", None). Worth updating autoencoder_kl_wan_2p2.py to match (noting its None vs "redundant" leading axis).

3. Pipeline divisibility warning (wan_pipeline.py:559, wan_vace_pipeline_2_1.py:471, wan_pipeline_animate.py:423):

   • These still warn whenever (width // 8) % vae_spatial != 0; updating the condition to warn only when both (width // 8) % vae_spatial != 0 and (height // 8) % vae_spatial != 0 will avoid a false warning on 720x1280.

4. Unit test (src/maxdiffusion/tests/wan/wan_vae_test.py):

   • Could we add a small portrait encode/decode test case where (width // 8) % vae_spatial != 0 and (height // 8) % vae_spatial == 0 (e.g. H=32, W=24 on 4 devices)?