fixing vae tests

Author: eltsai

src/maxdiffusion/tests/wan_vae_test.py

@@ -560,10 +560,13 @@ def vae_encode(video, wan_vae, vae_cache, key):
 
     p_vae_encode = functools.partial(vae_encode, wan_vae=wan_vae, vae_cache=vae_cache, key=key)
     original_video_shape = original_video.shape
-    latent = p_vae_encode(original_video)
 
-    jitted_decode = functools.partial(wan_vae.decode, feat_cache=vae_cache, return_dict=False)
-    video = jitted_decode(latent)[0]
+    # Execute VAE operations within mesh context to match sharding constraints
+    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      latent = p_vae_encode(original_video)
+
+      jitted_decode = functools.partial(wan_vae.decode, feat_cache=vae_cache, return_dict=False)
+      video = jitted_decode(latent)[0]
     video = jnp.transpose(video, (0, 4, 1, 2, 3))
     assert video.shape == original_video_shape