Reenable OpInfo tests. (#8519)

pytorch · Dec 23, 2024 · c385359 · c385359
1 parent cb44d62
commit c385359
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -12,11 +12,14 @@
 
 skiplist = {
     "_segment_reduce",
+    "_unsafe_masked_index_put_accumulate",
     "bincount", # NOTE: dtype for int input torch gives float. This is weird.
     "byte",
     "cat",
     "cholesky_solve",
+    "cov",
     "diagonal_copy",
+    "gather",
     "geqrf",
     "histogram", # hard op: AssertionError: Tensor-likes are not close!
     "histogramdd", # TypeError: histogram requires ndarray or scalar arguments, got <class 'list'> at position 1.
@@ -44,6 +47,7 @@
     "normal",
     "ormqr",
     "pca_lowrank",
+    "scatter",
     "searchsorted",
     "special.airy_ai",
     "special.scaled_modified_bessel_k0",
@@ -229,7 +233,7 @@ def test_reference_eager(self, device, dtype, op):
                              ignore_indices=ignore_index)
 
 
-instantiate_device_type_tests(TestOpInfo, globals(), only_for='cpu')
+instantiate_device_type_tests(TestOpInfo, globals(), only_for={'cpu'})
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/experimental/torch_xla2/torch_xla2/ops/mappings.py b/experimental/torch_xla2/torch_xla2/ops/mappings.py
@@ -19,8 +19,7 @@ def t2j(t):
     t = t.contiguous()
 
   try:
-    dl = torchdl.to_dlpack(t)
-    res = jaxdl.from_dlpack(dl)
+    res = jaxdl.from_dlpack(t)
   except Exception:
     # https://github.com/google/jax/issues/7657
     # https://github.com/google/jax/issues/17784