WIP log_normal

pytorch · Oct 9, 2024 · ed95937 · ed95937
1 parent 07d0823
commit ed95937
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -48,7 +48,6 @@
     "linalg.tensorsolve",
     "linalg.vector_norm",
     "linspace",
-    "log_normal",
     "logspace",
     "lu",
     "lu_solve",
@@ -260,6 +259,7 @@ def test_reference_eager(self, device, dtype, op):
       if isinstance(t, torch.Tensor) and t.is_sparse:
         continue
       check_output = op.name not in random_ops
+      print("[DEBUG] sample_input: ", sample_input)
 
       if op.name == "special.polygamma":
         # The polygamma function is inaccurate for values < 1.

diff --git a/experimental/torch_xla2/torch_xla2/decompositions.py b/experimental/torch_xla2/torch_xla2/decompositions.py
@@ -297,4 +297,5 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tens
     torch.ops.aten.nll_loss2d_backward,
     torch.ops.aten.bernoulli_.Tensor,
     torch.ops.aten.bernoulli_.float,
+    torch.ops.aten.log_normal,
 ])
diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -47,6 +47,7 @@
   torch.ops.aten.logical_not_: torch.ops.aten.logical_not,
   torch.ops.aten.unsqueeze_: torch.ops.aten.unsqueeze,
   torch.ops.aten.transpose_: torch.ops.aten.transpose,
+  torch.ops.aten.log_normal_: torch.ops.aten.log_normal,
 }