move params to tpu device

Author: eltsai

src/maxdiffusion/tests/wan_vae_test.py

@@ -556,17 +556,17 @@ 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)
 
     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)
 
-    # 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]
+    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