fixing vae tests

Author: eltsai

src/maxdiffusion/tests/wan_vae_test.py

@@ -22,11 +22,13 @@
 import jax
 import jax.numpy as jnp
 from flax import nnx
+from flax import linen as nn
 from flax.linen import partitioning as nn_partitioning
 from jax.sharding import Mesh
 from .. import pyconfig
 from ..max_utils import (
     create_device_mesh,
+    device_put_replicated,
 )
 import numpy as np
 import unittest
@@ -556,10 +558,18 @@ def vae_encode(video, wan_vae, vae_cache, key):
     # This replaces random params with the model.
     params = load_wan_vae(config.pretrained_model_name_or_path, params, "cpu")
     params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), params)
-    # Transfer params to TPU device before merging to fix device mismatch
-    tpu_device = jax.devices("tpu")[0] if jax.devices("tpu") else jax.devices()[0]
-    params = jax.device_put(params, tpu_device)
-    wan_vae = nnx.merge(graphdef, params)
+
+    logical_state_spec = nnx.get_partition_spec(state)
+    logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
+    logical_state_sharding = dict(nnx.to_flat_state(logical_state_sharding))
+
+    state_flat = dict(nnx.to_flat_state(state))
+    for path, val in flax.traverse_util.flatten_dict(params).items():
+      sharding = logical_state_sharding[path].get_value()
+      state_flat[path][...] = device_put_replicated(val, sharding)
+    state = nnx.from_flat_state(state_flat)
+
+    wan_vae = nnx.merge(graphdef, state)
 
     p_vae_encode = functools.partial(vae_encode, wan_vae=wan_vae, vae_cache=vae_cache, key=key)
     original_video_shape = original_video.shape