test

pytorch · Jan 24, 2025 · 3f7c67a · 3f7c67a
1 parent 6822925
commit 3f7c67a
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/.github/workflows/torch_xla2.yml b/.github/workflows/torch_xla2.yml
@@ -44,5 +44,5 @@ jobs:
         working-directory: torchax
         shell: bash
         run: |
-          pytest test/
+          pytest --forked test/
           XLA_FLAGS=--xla_force_host_platform_device_count=4 pytest -n 0 test_dist/
diff --git a/torchax/test-requirements.txt b/torchax/test-requirements.txt
@@ -3,6 +3,7 @@ absl-py
 immutabledict
 pytest
 pytest-xdist
+pytest-forked
 sentencepiece
 expecttest
 optax
diff --git a/torchax/torchax/tensor.py b/torchax/torchax/tensor.py
@@ -167,7 +167,7 @@ def tolist(self):
     return self._elem.tolist()
 
   def shard_(self, sharding):
-    self.apply_(jax.lax.with_sharding_constraint, sharding)
+    self.apply_jax_(jax.lax.with_sharding_constraint, sharding)
 
 
 def debug_accuracy(func, args, kwargs, current_output):