[CherryPick] Lower isneginf(). (#8926)

Author: ysiraichi

codegen/xla_native_functions.yaml

@@ -60,6 +60,7 @@ full_codegen:
   - hardswish_backward
   - inverse
   - isnan
+  - isneginf
   - leaky_relu
   - le.Scalar
   - le.Tensor

test/test_operations.py

@@ -35,6 +35,7 @@
     all_types_and_complex_and,
     all_types_and,
 )
+import torch.utils._pytree as pytree
 import torch_xla
 import torch_xla.core.xla_builder as xb
 import torch_xla.core.xla_op_registry as xor
@@ -2384,24 +2385,35 @@ def test_cummax_0_sized_dimension(self):
 
     self.assertEqual(actual, expected)
 
-  def test_conj(self):
-    # Leave the factory out of the fallback count.
-    tensor = torch.rand(2, 2, dtype=torch.complex64)
-
+  def _test_no_fallback(self, runf, args):
     met.clear_all()
 
     def run(device):
-      return torch.conj(tensor.to(device))
+      args_ = pytree.tree_map_only(torch.Tensor,
+                                   lambda t: t.clone().detach().to(device),
+                                   args)
+      return runf(*args_)
 
     actual = run("cpu")
     expected = run(xm.xla_device())
 
-    self.assertEqual(
-        met.executed_fallback_ops(), [],
-        message="expected no fallback operations.")
+    self.assertFalse(
+        met.executed_fallback_ops(), msg="expected no fallback operations.")
     self.assertEqual(
         actual, expected.cpu(), message="XLA results should match CPU results.")
 
+  def test_conj_no_fallback(self):
+    tensor = torch.rand(2, 2, dtype=torch.complex64)
+    self._test_no_fallback(torch.conj, (tensor,))
+
+  def test_isneginf_no_fallback(self):
+    t = torch.rand(10)
+    # Scale the tensor elements.
+    t = t * 100_000
+    # Convert to a lower precision data-type so as to get a few infs.
+    t = t.to(torch.float16)
+    self._test_no_fallback(torch.isneginf, (t,))
+
 
 class MNISTComparator(nn.Module):
 

test/test_ops.py

@@ -127,6 +127,7 @@ def get_allowed_ops_map(
     AllowedOpInfoEntry('imag'),
     AllowedOpInfoEntry('inverse'),
     AllowedOpInfoEntry('isin'),
+    AllowedOpInfoEntry('isneginf'),
     AllowedOpInfoEntry('le'),
     AllowedOpInfoEntry('linalg.cholesky'),
     AllowedOpInfoEntry('linalg.cholesky_ex'),

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -13,6 +13,7 @@
 #include "torch_xla/csrc/xla_lower_util.h"
 #include "xla/client/lib/math.h"
 #include "xla/client/lib/matrix.h"
+#include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/lib/logdet.h"
 
 namespace torch_xla {
@@ -531,6 +532,13 @@ torch_xla::XlaOpVector Isnan::Lower(LoweringContext* loctx) const {
   return ReturnOp(xla::IsNan(xla_input), loctx);
 }
 
+torch_xla::XlaOpVector Isneginf::Lower(LoweringContext* loctx) const {
+  xla::XlaOp input = loctx->GetOutputOp(operand(0));
+  return ReturnOp(xla::Eq(input, xla::MinValue(input.builder(),
+                                               XlaHelpers::TypeOfXlaOp(input))),
+                  loctx);
+}
+
 torch_xla::XlaOpVector LeakyRelu::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   xla::XlaOp negative_slope = loctx->GetOutputOp(operand(1));

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -590,6 +590,12 @@ xla::Shape IsnanOutputShape(const torch::lazy::Value& input) {
   return isnan_shape;
 }
 
+xla::Shape IsneginfOutputShape(const torch::lazy::Value& input) {
+  xla::Shape shape(GetXlaShape(input));
+  shape.set_element_type(xla::PRED);
+  return shape;
+}
+
 xla::Shape LeakyReluOutputShape(const torch::lazy::Value& input,
                                 const torch::lazy::Value& negative_slope) {
   auto lower_for_shape_fn =

torch_xla/csrc/ops/ops_xla_shape_fn.h

@@ -174,6 +174,8 @@ xla::Shape InverseOutputShape(const torch::lazy::Value& input);
 
 xla::Shape IsnanOutputShape(const torch::lazy::Value& input);
 
+xla::Shape IsneginfOutputShape(const torch::lazy::Value& input);
+
 xla::Shape LeakyReluOutputShape(const torch::lazy::Value& input,
                                 const torch::lazy::Value& negative_slope);