Test abstract array

Author: NeilGirdhar

tests/test_funcs.py

@@ -7,6 +7,7 @@
 import hypothesis.extra.numpy as npst
 import numpy as np
 import pytest
+from array_api_compat import is_jax_namespace
 from hypothesis import given
 from hypothesis import strategies as st
 
@@ -470,6 +471,16 @@ def test_dims_and_classes(self, xp: ModuleType, n_dim: int, num_classes: int):
             i = tuple(i_list)
             assert float(y[(*i, j)]) == (int(x[i]) == j)
 
+    def test_empty_shape(self, xp: ModuleType):
+        if not is_jax_namespace(xp):
+            pytest.skip("backend does not support abstract arrays")
+        import jax
+        import jax.numpy as jnp
+
+        abstract_input = jax.ShapeDtypeStruct(shape=(None, 784), dtype=jnp.float32)
+        with pytest.raises(TypeError):
+            _ = one_hot(abstract_input, 3)
+
     def test_basic(self, xp: ModuleType):
         actual = one_hot(xp.asarray([0, 1, 2]), 3)
         expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])