Improve ndarray testing (wjakob#1168)

Author: hpkfft

tests/test_jax.py

@@ -0,0 +1,64 @@
+import test_ndarray_ext as t
+import test_jax_ext as tj
+import pytest
+import warnings
+import importlib
+from common import collect
+
+try:
+    import jax.numpy as jnp
+    def needs_jax(x):
+        return x
+except:
+    needs_jax = pytest.mark.skip(reason="JAX is required")
+
+
+@needs_jax
+def test01_constrain_order_jax():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            c = jnp.zeros((3, 5))
+        except:
+            pytest.skip('jax is missing')
+
+    z = jnp.zeros((3, 5, 4, 6))
+    assert t.check_order(z) == 'C'
+
+
+@needs_jax
+def test02_implicit_conversion_jax():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            c = jnp.zeros((3, 5))
+        except:
+            pytest.skip('jax is missing')
+
+    t.implicit(jnp.zeros((2, 2), dtype=jnp.int32))
+    t.implicit(jnp.zeros((2, 2, 10), dtype=jnp.float32)[:, :, 4])
+    t.implicit(jnp.zeros((2, 2, 10), dtype=jnp.int32)[:, :, 4])
+    t.implicit(jnp.zeros((2, 2, 10), dtype=jnp.bool_)[:, :, 4])
+
+    with pytest.raises(TypeError) as excinfo:
+        t.noimplicit(jnp.zeros((2, 2), dtype=jnp.int32))
+
+    with pytest.raises(TypeError) as excinfo:
+        t.noimplicit(jnp.zeros((2, 2), dtype=jnp.uint8))
+
+
+@needs_jax
+def test03_return_jax():
+    collect()
+    dc = tj.destruct_count()
+    x = tj.ret_jax()
+    assert x.shape == (2, 4)
+    assert jnp.all(x == jnp.array([[1,2,3,4], [5,6,7,8]], dtype=jnp.float32))
+    del x
+    collect()
+    assert tj.destruct_count() - dc == 1
+
+
+@needs_jax
+def test04_check_jax():
+    assert t.check(jnp.zeros((1)))

tests/test_ndarray.cpp

@@ -12,9 +12,9 @@ int destruct_count = 0;
 static float f_global[] { 1, 2, 3, 4, 5, 6, 7, 8 };
 static int i_global[] { 1, 2, 3, 4, 5, 6, 7, 8 };
 
-#if defined(__aarch64__)
+#if defined(__aarch64__) || defined(__AVX512FP16__)
 namespace nanobind::detail {
-    template <> struct dtype_traits<__fp16> {
+    template <> struct dtype_traits<_Float16> {
         static constexpr dlpack::dtype value {
             (uint8_t) dlpack::dtype_code::Float, // type code
             16, // size in bits
@@ -392,17 +392,17 @@ NB_MODULE(test_ndarray_ext, m) {
             v(i) = -v(i);
     }, "x"_a.noconvert());
 
-#if defined(__aarch64__)
+#if defined(__aarch64__) || defined(__AVX512FP16__)
     m.def("ret_numpy_half", []() {
-        __fp16 *f = new __fp16[8] { 1, 2, 3, 4, 5, 6, 7, 8 };
+        _Float16 *f = new _Float16[8] { 1, 2, 3, 4, 5, 6, 7, 8 };
         size_t shape[2] = { 2, 4 };
 
         nb::capsule deleter(f, [](void *data) noexcept {
             destruct_count++;
-            delete[] (__fp16*) data;
+            delete[] (_Float16*) data;
         });
-        return nb::ndarray<nb::numpy, __fp16, nb::shape<2, 4>>(f, 2, shape,
-                                                               deleter);
+        return nb::ndarray<nb::numpy, _Float16, nb::shape<2, 4>>(f, 2, shape,
+                                                                 deleter);
     });
 #endif
 

tests/test_ndarray.py

@@ -1,6 +1,4 @@
 import test_ndarray_ext as t
-import test_jax_ext as tj
-import test_tensorflow_ext as tt
 import pytest
 import warnings
 import importlib
@@ -20,21 +18,6 @@ def needs_torch(x):
 except:
     needs_torch = pytest.mark.skip(reason="PyTorch is required")
 
-try:
-    import tensorflow as tf
-    import tensorflow.config
-    def needs_tensorflow(x):
-        return x
-except:
-    needs_tensorflow = pytest.mark.skip(reason="TensorFlow is required")
-
-try:
-    import jax.numpy as jnp
-    def needs_jax(x):
-        return x
-except:
-    needs_jax = pytest.mark.skip(reason="JAX is required")
-
 try:
     import cupy as cp
     def needs_cupy(x):
@@ -158,19 +141,6 @@ def test05_constrain_order():
     assert t.check_order(np.zeros((3, 5, 4, 6), order='F')[:, 2, :, :]) == '?'
 
 
-@needs_jax
-def test06_constrain_order_jax():
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        try:
-            c = jnp.zeros((3, 5))
-        except:
-            pytest.skip('jax is missing')
-
-    z = jnp.zeros((3, 5, 4, 6))
-    assert t.check_order(z) == 'C'
-
-
 @needs_torch
 @pytest.mark.filterwarnings
 def test07_constrain_order_pytorch():
@@ -190,18 +160,6 @@ def test07_constrain_order_pytorch():
         assert t.check_device(torch.zeros(3, 5, device='cuda')) == 'cuda'
 
 
-@needs_tensorflow
-def test08_constrain_order_tensorflow():
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        try:
-            c = tf.zeros((3, 5))
-        except:
-            pytest.skip('tensorflow is missing')
-
-    assert t.check_order(c) == 'C'
-
-
 @needs_numpy
 def test09_write_from_cpp():
     x = np.zeros(10, dtype=np.float32)
@@ -251,48 +209,6 @@ def test11_implicit_conversion_pytorch():
         t.noimplicit(torch.zeros(2, 2, 10, dtype=torch.float32)[:, :, 4])
 
 
-@needs_tensorflow
-def test12_implicit_conversion_tensorflow():
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        try:
-            c = tf.zeros((3, 5))
-        except:
-            pytest.skip('tensorflow is missing')
-
-        t.implicit(tf.zeros((2, 2), dtype=tf.int32))
-        t.implicit(tf.zeros((2, 2, 10), dtype=tf.float32)[:, :, 4])
-        t.implicit(tf.zeros((2, 2, 10), dtype=tf.int32)[:, :, 4])
-        t.implicit(tf.zeros((2, 2, 10), dtype=tf.bool)[:, :, 4])
-
-        with pytest.raises(TypeError) as excinfo:
-            t.noimplicit(tf.zeros((2, 2), dtype=tf.int32))
-
-        with pytest.raises(TypeError) as excinfo:
-            t.noimplicit(tf.zeros((2, 2), dtype=tf.bool))
-
-
-@needs_jax
-def test13_implicit_conversion_jax():
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        try:
-            c = jnp.zeros((3, 5))
-        except:
-            pytest.skip('jax is missing')
-
-    t.implicit(jnp.zeros((2, 2), dtype=jnp.int32))
-    t.implicit(jnp.zeros((2, 2, 10), dtype=jnp.float32)[:, :, 4])
-    t.implicit(jnp.zeros((2, 2, 10), dtype=jnp.int32)[:, :, 4])
-    t.implicit(jnp.zeros((2, 2, 10), dtype=jnp.bool_)[:, :, 4])
-
-    with pytest.raises(TypeError) as excinfo:
-        t.noimplicit(jnp.zeros((2, 2), dtype=jnp.int32))
-
-    with pytest.raises(TypeError) as excinfo:
-        t.noimplicit(jnp.zeros((2, 2), dtype=jnp.uint8))
-
-
 def test14_destroy_capsule():
     collect()
     dc = t.destruct_count()
@@ -376,31 +292,6 @@ def test18_return_pytorch():
     assert t.destruct_count() - dc == 1
 
 
-@needs_jax
-def test19_return_jax():
-    collect()
-    dc = tj.destruct_count()
-    x = tj.ret_jax()
-    assert x.shape == (2, 4)
-    assert jnp.all(x == jnp.array([[1,2,3,4], [5,6,7,8]], dtype=jnp.float32))
-    del x
-    collect()
-    assert tj.destruct_count() - dc == 1
-
-
-@needs_tensorflow
-def test20_return_tensorflow():
-    collect()
-    dc = tt.destruct_count()
-    x = tt.ret_tensorflow()
-    assert x.get_shape().as_list() == [2, 4]
-    assert tf.math.reduce_all(
-               x == tf.constant([[1,2,3,4], [5,6,7,8]], dtype=tf.float32))
-    del x
-    collect()
-    assert tt.destruct_count() - dc == 1
-
-
 @needs_numpy
 def test21_return_array_scalar():
     collect()
@@ -504,16 +395,6 @@ def test28_check_torch():
     assert t.check(torch.zeros((1)))
 
 
-@needs_tensorflow
-def test29_check_tensorflow():
-    assert t.check(tf.zeros((1)))
-
-
-@needs_jax
-def test30_check_jax():
-    assert t.check(jnp.zeros((1)))
-
-
 @needs_numpy
 def test31_rv_policy():
     def p(a):
@@ -879,8 +760,6 @@ def test45_implicit_conversion_cupy():
 @needs_numpy
 def test46_implicit_conversion_contiguous_complex():
     # Test fix for issue #709
-    import numpy as np
-
     c_f32 = np.random.rand(10, 10)
     c_c64 = c_f32.astype(np.complex64)
 
@@ -907,7 +786,6 @@ def test_conv(x):
 
 @needs_numpy
 def test_47_ret_infer():
-    import numpy as np
     assert np.all(t.ret_infer_c() == [[1, 2, 3, 4], [5, 6, 7, 8]])
     assert np.all(t.ret_infer_f() == [[1, 3, 5, 7], [2, 4, 6, 8]])
 
@@ -956,13 +834,12 @@ def test50_test_matrix4f_copy():
 
 @needs_numpy
 def test51_return_from_stack():
-    import numpy as np
     assert np.all(t.ret_from_stack_1() == [1,2,3])
     assert np.all(t.ret_from_stack_2() == [1,2,3])
 
+
 @needs_numpy
 def test52_accept_np_both_true_contig():
-    import numpy as np
     a = np.zeros((2, 1), dtype=np.float32)
     assert a.flags['C_CONTIGUOUS'] and a.flags['F_CONTIGUOUS']
     t.accept_np_both_true_contig_a(a)
@@ -972,6 +849,5 @@ def test52_accept_np_both_true_contig():
 
 @needs_numpy
 def test53_issue_930():
-    import numpy as np
     wrapper = t.Wrapper(np.ones(3, dtype=np.float32))
     assert wrapper.value[0] == 1

tests/test_tensorflow.py

@@ -0,0 +1,65 @@
+import test_ndarray_ext as t
+import test_tensorflow_ext as ttf
+import pytest
+import warnings
+import importlib
+from common import collect
+
+try:
+    import tensorflow as tf
+    import tensorflow.config
+    def needs_tensorflow(x):
+        return x
+except:
+    needs_tensorflow = pytest.mark.skip(reason="TensorFlow is required")
+
+
+@needs_tensorflow
+def test01_constrain_order_tensorflow():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            c = tf.zeros((3, 5))
+        except:
+            pytest.skip('tensorflow is missing')
+
+    assert t.check_order(c) == 'C'
+
+
+@needs_tensorflow
+def test02_implicit_conversion_tensorflow():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            c = tf.zeros((3, 5))
+        except:
+            pytest.skip('tensorflow is missing')
+
+        t.implicit(tf.zeros((2, 2), dtype=tf.int32))
+        t.implicit(tf.zeros((2, 2, 10), dtype=tf.float32)[:, :, 4])
+        t.implicit(tf.zeros((2, 2, 10), dtype=tf.int32)[:, :, 4])
+        t.implicit(tf.zeros((2, 2, 10), dtype=tf.bool)[:, :, 4])
+
+        with pytest.raises(TypeError) as excinfo:
+            t.noimplicit(tf.zeros((2, 2), dtype=tf.int32))
+
+        with pytest.raises(TypeError) as excinfo:
+            t.noimplicit(tf.zeros((2, 2), dtype=tf.bool))
+
+
+@needs_tensorflow
+def test03_return_tensorflow():
+    collect()
+    dc = ttf.destruct_count()
+    x = ttf.ret_tensorflow()
+    assert x.get_shape().as_list() == [2, 4]
+    assert tf.math.reduce_all(
+               x == tf.constant([[1,2,3,4], [5,6,7,8]], dtype=tf.float32))
+    del x
+    collect()
+    assert ttf.destruct_count() - dc == 1
+
+
+@needs_tensorflow
+def test04_check_tensorflow():
+    assert t.check(tf.zeros((1)))