Split _ops and _decomps. (#9323)

Author: qihqi

torchax/test/test_libraries.py

@@ -6,7 +6,6 @@
 import torchax.export
 from torchax.ops import jaten
 from torchax.ops import jlibrary
-
 # Create a `mylib` library which has a basic SDPA op.
 m = Library("mylib", "DEF")
 m.define("scaled_dot_product_attention(Tensor q, Tensor k, Tensor v) -> Tensor")

torchax/test/test_ops.py

@@ -143,26 +143,26 @@ def run_export_and_compare(testcase,
           (sample_input.input, sample_input.args, sample_input.kwargs))
       with testcase.env:
         res2 = func(input2, *args2, **kwargs2)
-      res2 = pytree.tree_map_only(tensor.Tensor, lambda t: t.torch(), res2)
-      with testcase.subTest("torchax_diff:" + str(atol)):
-        if ignore_indices and isinstance(res, tuple) and len(res) == 2:
-          diff_output(
-              testcase,
-              res[0],
-              res2[0],
-              atol=atol,
-              rtol=rtol,
-              equal_nan=equal_nan,
-              check_output=check_output)
-        else:
-          diff_output(
-              testcase,
-              res,
-              res2,
-              atol=atol,
-              rtol=rtol,
-              equal_nan=equal_nan,
-              check_output=check_output)
+        res2 = pytree.tree_map_only(tensor.Tensor, lambda t: t.torch(), res2)
+        with testcase.subTest("torchax_diff:" + str(atol)):
+          if ignore_indices and isinstance(res, tuple) and len(res) == 2:
+            diff_output(
+                testcase,
+                res[0],
+                res2[0],
+                atol=atol,
+                rtol=rtol,
+                equal_nan=equal_nan,
+                check_output=check_output)
+          else:
+            diff_output(
+                testcase,
+                res,
+                res2,
+                atol=atol,
+                rtol=rtol,
+                equal_nan=equal_nan,
+                check_output=check_output)
 
 
 ops_to_test = [
@@ -188,7 +188,6 @@ def setUp(self):
     torchax.enable_accuracy_mode()
     #self.env.config.debug_accuracy_for_each_op = True
     self.env.config.debug_print_each_op = True
-    self.env.config.debug_print_each_op_operands = True
     torch.manual_seed(0)
     self.old_var = self.env.config.use_torch_native_for_cpu_tensor
     self.env.config.use_torch_native_for_cpu_tensor = False

torchax/test_dist/test_distributed.py

@@ -14,6 +14,8 @@
 # TODO(wcromar): do something useful with group name
 GROUP_NAME = "process_group"
 
+torchax.enable_globally()
+
 
 @pytest.fixture(scope="module")
 def multi_cpu():
@@ -80,8 +82,9 @@ def test_all_reduce(op, expected, multi_cpu, process_group):
   device_count = multi_cpu
 
   def f(index):
-    dist.all_reduce(index, op)
-    return index
+    with torchax.default_env():
+      dist.all_reduce(index, op)
+      return index
 
   res = torchax.distributed.spawn(f)
 

torchax/torchax/decompositions.py

@@ -131,7 +131,7 @@ def channel_shuffle(self, groups):
 
 
 def bernoulli_float(self, p=0.5):
-  return self.bernoulli_(torch.tensor(p))
+  return self.bernoulli_(p)
 
 
 _try_register(aten.bernoulli_.float, bernoulli_float)