fix test

Author: tengyifei

torchprime/sharding/shard_model.py

@@ -186,7 +186,7 @@ def shard_torchax_model_from_config(
   """
   import jax
   from jax.sharding import NamedSharding, PartitionSpec
-  from torchax.interop import torch_view
+  from torchax.interop import jax_view, torch_view
 
   jax_mark_sharding = torch_view(jax.lax.with_sharding_constraint)
 
@@ -197,7 +197,7 @@ def shard_param(tensor, spec: tuple[str, ...]):
     # and models are usually constructed eagerly in torchax.
     return torch_view(
       jax.make_array_from_callback(
-        tensor.shape, sharding, lambda slice_index: tensor[slice_index]
+        tensor.shape, sharding, lambda slice_index: jax_view(tensor[slice_index])
       )
     )
 

torchprime/sharding/tests/test_shard_model.py

@@ -239,7 +239,8 @@ def test_shard_model_from_config_torchax():
     devices = mesh_utils.create_device_mesh((jax.device_count(),))
     mesh = Mesh(devices, ("fsdp",))
 
-    model = shard_torchax_model_from_config(model, config, mesh)
+    with torchax.default_env():
+      model = shard_torchax_model_from_config(model, config, mesh)
 
     # In order to shard activations, corresponding modules are
     # wrapped with ShardedModule.