added vae_spatial

Author: eltsai

src/maxdiffusion/tests/wan_vae_test.py

@@ -281,7 +281,13 @@ def test_3d_conv(self):
     )
     config = pyconfig.config
     devices_array = create_device_mesh(config)
-    mesh = Mesh(devices_array, config.mesh_axes)
+    # Add vae_spatial axis to mesh for VAE operations
+    mesh_axes = list(config.mesh_axes)
+    if "vae_spatial" not in mesh_axes:
+      mesh_axes.append("vae_spatial")
+      # Reshape devices to include vae_spatial (size 1 for test)
+      devices_array = devices_array.reshape(devices_array.shape + (1,))
+    mesh = Mesh(devices_array, mesh_axes)
 
     batch_size = 1
     in_depth, in_height, in_width = 10, 32, 32