Extract JAX and TF from test_ndarray.cpp. Put stub tests behind flag.

This commit extracts the functions involving JAX and TensorFlow types
from the module in `test_ndarray.cpp` into their own respective modules
and sets up the boilerplate code to turn them into independent tests.
This allows to put the corresponding stub generation tests behind a
flag and deactivate them by default.

Signed-off-by: Ingo Müller <[email protected]>

Author: ingomueller-net

tests/CMakeLists.txt

@@ -72,6 +72,8 @@ set(TEST_NAMES
   enum
   eval
   ndarray
+  jax
+  tensorflow
   exception
   make_iterator
   typing
@@ -90,7 +92,7 @@ endforeach()
 
 target_sources(test_intrusive_ext PRIVATE test_intrusive_impl.cpp)
 
-foreach (NAME functions classes ndarray stl enum typing make_iterator)
+foreach (NAME functions classes ndarray jax tensorflow stl enum typing make_iterator)
   if (NAME STREQUAL typing)
     set(EXTRA
       MARKER_FILE py.typed
@@ -138,6 +140,7 @@ target_link_libraries(test_inter_module_2_ext PRIVATE inter_module)
 
 set(TEST_FILES
   common.py
+  conftest.py
   test_callbacks.py
   test_classes.py
   test_eigen.py
@@ -154,6 +157,8 @@ set(TEST_FILES
   test_stl_bind_vector.py
   test_chrono.py
   test_ndarray.py
+  test_jax.py
+  test_tensorflow.py
   test_stubs.py
   test_typing.py
   test_thread.py
@@ -163,6 +168,8 @@ set(TEST_FILES
   test_functions_ext.pyi.ref
   test_make_iterator_ext.pyi.ref
   test_ndarray_ext.pyi.ref
+  test_jax_ext.pyi.ref
+  test_tensorflow_ext.pyi.ref
   test_stl_ext.pyi.ref
   test_enum_ext.pyi.ref
   test_typing_ext.pyi.ref

tests/conftest.py

@@ -0,0 +1,6 @@
+def pytest_addoption(parser):
+    parser.addoption('--enable-slow-tests',
+                     action='store_true',
+                     dest="enable-slow-tests",
+                     default=False,
+                     help="enable longrundecorated tests")

tests/test_jax.cpp

@@ -0,0 +1,22 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+
+namespace nb = nanobind;
+
+int destruct_count = 0;
+
+NB_MODULE(test_jax_ext, m) {
+    m.def("destruct_count", []() { return destruct_count; });
+    m.def("ret_jax", []() {
+        float *f = new float[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[] (float *) data;
+        });
+
+        return nb::ndarray<nb::jax, float, nb::shape<2, 4>>(f, 2, shape,
+                                                            deleter);
+    });
+}

tests/test_jax.py

@@ -0,0 +1,4 @@
+import jaxlib.xla_extension
+
+
+def ret_jax() -> jaxlib.xla_extension.DeviceArray[dtype=float32, shape=(2, 4)]: ...

tests/test_jax_ext.pyi.ref

@@ -299,35 +299,6 @@ NB_MODULE(test_ndarray_ext, m) {
                                                                 deleter);
     });
 
-    m.def("ret_jax", []() {
-        float *f = new float[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[] (float *) data;
-        });
-
-        return nb::ndarray<nb::jax, float, nb::shape<2, 4>>(f, 2, shape,
-                                                            deleter);
-    });
-
-    m.def("ret_tensorflow", []() {
-        struct alignas(256) Buf {
-            float f[8];
-        };
-        Buf *buf = new Buf({ 1, 2, 3, 4, 5, 6, 7, 8 });
-        size_t shape[2] = { 2, 4 };
-
-        nb::capsule deleter(buf, [](void *data) noexcept {
-           destruct_count++;
-           delete (Buf *) data;
-        });
-
-        return nb::ndarray<nb::tensorflow, float, nb::shape<2, 4>>(buf->f, 2, shape,
-                                                                   deleter);
-    });
-
     m.def("ret_array_scalar", []() {
             float* f = new float[1] { 1 };
             size_t shape[1] = {};

tests/test_ndarray.cpp

@@ -1,4 +1,6 @@
 import test_ndarray_ext as t
+import test_jax_ext as tj
+import test_tensorflow_ext as tt
 import pytest
 import warnings
 import importlib
@@ -377,26 +379,26 @@ def test18_return_pytorch():
 @needs_jax
 def test19_return_jax():
     collect()
-    dc = t.destruct_count()
-    x = t.ret_jax()
+    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 t.destruct_count() - dc == 1
+    assert tj.destruct_count() - dc == 1
 
 
 @needs_tensorflow
 def test20_return_tensorflow():
     collect()
-    dc = t.destruct_count()
-    x = t.ret_tensorflow()
+    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 t.destruct_count() - dc == 1
+    assert tt.destruct_count() - dc == 1
 
 
 @needs_numpy

tests/test_ndarray.py

@@ -1,8 +1,6 @@
 from typing import Annotated, overload
 
-import jaxlib.xla_extension
 from numpy.typing import ArrayLike
-import tensorflow.python.framework.ops
 
 
 class Cls:
@@ -175,8 +173,6 @@ def ret_infer_c() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4), or
 
 def ret_infer_f() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4), order='F')]: ...
 
-def ret_jax() -> jaxlib.xla_extension.DeviceArray[dtype=float32, shape=(2, 4)]: ...
-
 def ret_numpy() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4))]: ...
 
 def ret_numpy_const() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4), writable=False)]: ...
@@ -189,8 +185,6 @@ def ret_numpy_half() -> Annotated[ArrayLike, dict(dtype='float16', shape=(2, 4))
 
 def ret_pytorch() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4))]: ...
 
-def ret_tensorflow() -> tensorflow.python.framework.ops.EagerTensor[dtype=float32, shape=(2, 4)]: ...
-
 def return_dlpack() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4))]: ...
 
 @overload

tests/test_ndarray_ext.pyi.ref

@@ -24,15 +24,23 @@ def remove_platform_dependent(s):
             i += 1
     return s2
 
+
 ref_paths = list(pathlib.Path(__file__).parent.glob('*.pyi.ref'))
+ref_path_ids = [
+    p.name.removeprefix('test_').removesuffix('.pyi.ref') for p in ref_paths
+]
 assert len(ref_paths) > 0, "Stub reference files not found!"
 
 @skip_on_unsupported
-@pytest.mark.parametrize('p_ref', ref_paths)
-def test01_check_stub_refs(p_ref):
+@pytest.mark.parametrize('p_ref', ref_paths, ids=ref_path_ids)
+def test01_check_stub_refs(p_ref, request):
     """
     Check that generated stub files match reference input
     """
+    if not request.config.getoption('enable-slow-tests') and any(
+        (x in p_ref.name for x in ['jax', 'tensorflow'])):
+        pytest.skip("skipping because slow tests are not enabled")
+
     p_in = p_ref.with_suffix('')
     with open(p_ref, 'r') as f:
         s_ref = f.read().split('\n')

tests/test_stubs.py

@@ -0,0 +1,25 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+namespace nb = nanobind;
+
+int destruct_count = 0;
+
+
+NB_MODULE(test_tensorflow_ext, m) {
+    m.def("destruct_count", []() { return destruct_count; });
+    m.def("ret_tensorflow", []() {
+        struct alignas(256) Buf {
+            float f[8];
+        };
+        Buf *buf = new Buf({ 1, 2, 3, 4, 5, 6, 7, 8 });
+        size_t shape[2] = { 2, 4 };
+
+        nb::capsule deleter(buf, [](void *data) noexcept {
+           destruct_count++;
+           delete (Buf *) data;
+        });
+
+        return nb::ndarray<nb::tensorflow, float, nb::shape<2, 4>>(buf->f, 2, shape,
+                                                                   deleter);
+    });
+}

tests/test_tensorflow.cpp

@@ -0,0 +1,4 @@
+import tensorflow.python.framework.ops
+
+
+def ret_tensorflow() -> tensorflow.python.framework.ops.EagerTensor[dtype=float32, shape=(2, 4)]: ...