Bucketize: add tests for inf inputs and boundaries

Author: Gallinator

kernels/test/op_bucketize_test.cpp

@@ -4,8 +4,8 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
-
 #include <gtest/gtest.h>
+#include <limits>
 
 using namespace ::testing;
 using executorch::aten::Scalar;
@@ -78,6 +78,21 @@ TEST_F(OpBucketizeScalarTest, ScalarEmptyBoundaries) {
   EXPECT_TENSOR_EQ(out, expected);
 }
 
+TEST_F(OpBucketizeScalarTest, ScalarInfInput) {
+  TensorFactory<ScalarType::Long> tf_out;
+  TensorFactory<ScalarType::Int> tf_bound;
+
+  Scalar value = std::numeric_limits<float>::infinity();
+  Tensor boundaries = tf_bound.make({5}, {0, 2, 4, 6, 8});
+  Tensor expected = tf_out.make({}, {5});
+  Tensor out = tf_out.zeros({});
+
+  Tensor ret = op_bucketize_out(value, boundaries, false, true, out);
+
+  EXPECT_TENSOR_EQ(ret, expected);
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
 TEST_F(OpBucketizeScalarTest, ScalarBoundaryTypes) {
   test_bucketize_bound_types();
 }
@@ -327,6 +342,42 @@ TEST_F(OpBucketizeTest, EmptyAll) {
   EXPECT_TENSOR_EQ(out, expected);
 }
 
+TEST_F(OpBucketizeTest, InfInput) {
+  TensorFactory<ScalarType::Long> tf_out;
+  TensorFactory<ScalarType::Float> tf_dtype;
+
+  Tensor values = tf_dtype.make(
+      {2},
+      {-std::numeric_limits<float>::infinity(),
+       std::numeric_limits<float>::infinity()});
+  Tensor boundaries = tf_dtype.make({5}, {0, 3, 5, 7, 9});
+  Tensor expected = tf_out.make({2}, {0, 5});
+  Tensor out = tf_out.zeros({2});
+
+  Tensor ret = op_bucketize_out(values, boundaries, false, true, out);
+
+  EXPECT_TENSOR_EQ(ret, expected);
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
+TEST_F(OpBucketizeTest, InfBoundaries) {
+  TensorFactory<ScalarType::Long> tf_out;
+  TensorFactory<ScalarType::Float> tf_dtype;
+
+  Tensor values = tf_dtype.make({2, 2}, {1, 4, 6, 8});
+  Tensor boundaries = tf_dtype.make(
+      {2},
+      {-std::numeric_limits<float>::infinity(),
+       std::numeric_limits<float>::infinity()});
+  Tensor expected = tf_out.ones({2, 2});
+  Tensor out = tf_out.zeros({2, 2});
+
+  Tensor ret = op_bucketize_out(values, boundaries, false, true, out);
+
+  EXPECT_TENSOR_EQ(ret, expected);
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
 TEST_F(OpBucketizeTest, BoundariesNDFails) {
   TensorFactory<ScalarType::Long> tf_out;
   TensorFactory<ScalarType::Float> tf_dtype;

kernels/test/test_bucketize.py

@@ -152,6 +152,29 @@ def test_bucketize_tensor_empty_input(self):
             model, (x, bounds), "test_bucketize_tensor_empty_input.pte"
         )
 
+    def test_bucketize_tensor_inf_input(self):
+        """Test bucketize.Tensor_out: (Tensor, Tensor, bool, bool) -> Tensor."""
+        model = BucketizeModule(out_int32=False, right=False)
+        x = torch.tensor([-torch.inf, torch.inf], dtype=torch.float)
+        bounds = torch.tensor([0, 3, 5, 7, 9], dtype=torch.float)
+        self._run_and_compare(model, (x, bounds), "test_bucketize_tensor_inf_input.pte")
+
+    def test_bucketize_tensor_inf_boundary(self):
+        """Test bucketize.Tensor_out: (Tensor, Tensor, bool, bool) -> Tensor."""
+        model = BucketizeModule(out_int32=False, right=False)
+        x = torch.tensor([[1, 2, 3, 4]], dtype=torch.float)
+        bounds = torch.tensor([-torch.inf, torch.inf], dtype=torch.float)
+        self._run_and_compare(
+            model, (x, bounds), "test_bucketize_tensor_inf_boundary.pte"
+        )
+
+    def test_bucketize_scalar_inf_input(self):
+        """Test bucketize.Tensor_out: (Tensor, Tensor, bool, bool) -> Tensor."""
+        model = BucketizeModule(out_int32=False, right=False)
+        x = torch.inf
+        bounds = torch.tensor([0, 3, 5, 7, 9], dtype=torch.float)
+        self._run_and_compare(model, (x, bounds), "test_bucketize_scalar_inf_input.pte")
+
 
 if __name__ == "__main__":
     unittest.main()