Add tests for creating Array2D views on CPU/GPU JAX Arrays

Author: horizon-blue

genmetaballs/src/cuda/bindings.cu

@@ -84,6 +84,7 @@ NB_MODULE(_genmetaballs_bindings, m) {
     utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))");
 
     bind_array2d<float, DeviceType::CPU>(utils, "CPUFloatArray2D");
+    bind_array2d<float, DeviceType::GPU>(utils, "GPUFloatArray2D");
 
 } // NB_MODULE(_genmetaballs_bindings)
 

genmetaballs/src/genmetaballs/core/__init__.py

@@ -3,10 +3,26 @@
     TwoParameterConfidence,
     ZeroParameterConfidence,
 )
-from genmetaballs._genmetaballs_bindings.utils import CPUFloatArray2D, sigmoid
+from genmetaballs._genmetaballs_bindings.utils import CPUFloatArray2D, GPUFloatArray2D, sigmoid
+
+
+def array2d_float(data, device) -> CPUFloatArray2D | GPUFloatArray2D:
+    """Create a FloatArray2D on the specified device from an array.
+
+    Args:
+        data: A 2D array of type float32.
+        device: 'cpu' or 'gpu' to specify the target device.
+    """
+    if device == "cpu":
+        return CPUFloatArray2D.from_array(data)
+    elif device == "gpu":
+        return GPUFloatArray2D.from_array(data)
+    else:
+        raise ValueError(f"Unsupported device type: {device}")
+
 
 __all__ = [
-    "CPUFloatArray2D",
+    "array2d_float",
     "ZeroParameterConfidence",
     "TwoParameterConfidence",
     "geometry",

tests/python_tests/test_utils.py

@@ -1,8 +1,10 @@
+import jax
+import jax.numpy as jnp
 import numpy as np
 import pytest
 from scipy.special import expit
 
-from genmetaballs.core import CPUFloatArray2D, sigmoid
+from genmetaballs.core import array2d_float, sigmoid
 
 NUM_RNG_SEEDS_PER_TEST = 5
 NUM_N_VALUES_PER_TEST = 5
@@ -60,11 +62,32 @@ def test_sigmoid_edge_cases(x: float) -> None:
         assert actual <= 1.0
 
 
-def test_float_array2d_creation_and_view():
+@pytest.mark.parametrize("device", ["cpu", "gpu"])
+def test_array2d_float_creation_on_jax_devices(device: str):
     """Test creation of Array2D from a numpy array."""
     rows, cols = 4, 5
+    data = jnp.arange(rows * cols, dtype=jnp.float32).reshape((rows, cols))
+    jax_device = jax.devices(device)[0]
+    data = jax.device_put(data, device=jax_device)
+    array_2d = array2d_float(data, device=device)
+
+    assert array_2d.num_rows == rows
+    assert array_2d.num_cols == cols
+    assert array_2d.ndim == 2
+
+    # then try converting back to numpy array via view
+    data_view = array_2d.as_jax()
+    assert data_view.device == data.device
+    assert jnp.allclose(data, data_view)
+
+    # Note: we can't test writability of shared view here since JAX arrays are immutable
+
+
+def test_float_array2d_view_numpy():
+    """Test creation of Array2D from a numpy array."""
+    rows, cols = 3, 4
     data = np.arange(rows * cols, dtype=np.float32).reshape((rows, cols))
-    array_2d = CPUFloatArray2D.from_array(data)
+    array_2d = array2d_float(data, device="cpu")
 
     assert array_2d.num_rows == rows
     assert array_2d.num_cols == cols
@@ -84,4 +107,4 @@ def test_create_invalid_array2d():
     data = np.arange(12, dtype=np.float32).reshape((3, 4))
 
     with pytest.raises(TypeError):
-        CPUFloatArray2D.from_array(data.reshape((3, 4, 1)))  # not 2D
+        array2d_float(data.reshape((3, 4, 1)), device="cpu")  # not 2D