test abstract size

Author: NeilGirdhar

tests/test_funcs.py

@@ -474,6 +474,27 @@ def test_basic(self, xp: ModuleType):
         expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
         xp_assert_equal(actual, expected)
 
+    def test_2d(self, xp: ModuleType):
+        actual = one_hot(xp.asarray([[2, 1, 0], [1, 0, 2]]), 3, axis=1)
+        expected = xp.asarray(
+            [
+                [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]],
+                [[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
+            ]
+        )
+        xp_assert_equal(actual, expected)
+
+    @pytest.mark.skip_xp_backend(
+        Backend.ARRAY_API_STRICTEST, reason="backend doesn't support Boolean indexing"
+    )
+    def test_abstract_size(self, xp: ModuleType):
+        x = xp.arange(5)
+        x = x[x > 2]
+        x = xp.astype(x, xp.int64)
+        actual = one_hot(x, 5)
+        expected = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0]])
+        xp_assert_equal(actual, expected)
+
     @pytest.mark.skip_xp_backend(
         Backend.TORCH_GPU, reason="Puts Pytorch into a bad state."
     )