Lower isneginf(). (#8912)

Author: ysiraichi

codegen/xla_native_functions.yaml

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

test/test_operations.py

@@ -2416,6 +2416,14 @@ def test_bitwise_right_shift_no_fallback(self):
     t2 = torch.randint(0, 10, (2,))
     self._test_no_fallback(torch.bitwise_right_shift, (t1, t2))
 
+  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 {
@@ -564,6 +565,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

@@ -606,6 +606,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

@@ -180,6 +180,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);