fix

Author: susanbao

src/maxdiffusion/tests/wan_transformer_test.py

@@ -38,7 +38,6 @@
 from maxdiffusion.pyconfig import HyperParameters
 from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline
 import qwix
-import numpy as np
 
 RealQtRule = qwix.QtRule
 
@@ -69,9 +68,17 @@ def test_nnx_pixart_alpha_text_projection(self):
     key = jax.random.key(0)
     rngs = nnx.Rngs(key)
     dummy_caption = jnp.ones((1, 512, 4096))
-    num_devices = len(jax.devices())
-    device_mesh = np.array(jax.devices()).reshape((1, num_devices))
-    mesh = Mesh(device_mesh, axis_names=('embed', 'mlp'))
+    pyconfig.initialize(
+        [
+            None,
+            os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
+        ],
+        unittest=True,
+    )
+    config = pyconfig.config
+    devices_array = create_device_mesh(config)
+    mesh = Mesh(devices_array, config.mesh_axes)
+
     with mesh:
       layer = NNXPixArtAlphaTextProjection(rngs=rngs, in_features=4096, hidden_size=5120)
     dummy_output = layer(dummy_caption)
@@ -80,9 +87,20 @@ def test_nnx_pixart_alpha_text_projection(self):
   def test_nnx_timestep_embedding(self):
     key = jax.random.key(0)
     rngs = nnx.Rngs(key)
+    pyconfig.initialize(
+        [
+            None,
+            os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
+        ],
+        unittest=True,
+    )
+    config = pyconfig.config
+    devices_array = create_device_mesh(config)
+    mesh = Mesh(devices_array, config.mesh_axes)
 
     dummy_sample = jnp.ones((1, 256))
-    layer = NNXTimestepEmbedding(rngs=rngs, in_channels=256, time_embed_dim=5120)
+    with mesh:
+      layer = NNXTimestepEmbedding(rngs=rngs, in_channels=256, time_embed_dim=5120)
     dummy_output = layer(dummy_sample)
     assert dummy_output.shape == (1, 5120)