Validate static integer shapes in VirtualBrownianTree and UnsafeBrownianPath (Fixes #707)

diffrax/_brownian/path.py

@@ -8,6 +8,7 @@
 import jax.random as jr
 import jax.tree_util as jtu
 import lineax.internal as lxi
+import numpy as np
 from jaxtyping import Array, PRNGKeyArray, PyTree
 from lineax.internal import complex_to_real_dtype
 
@@ -69,7 +70,7 @@ class UnsafeBrownianPath(AbstractBrownianPath):
 
     def __init__(
         self,
-        shape: tuple[int, ...] | PyTree[jax.ShapeDtypeStruct],
+        shape: tuple[int | np.integer, ...] | PyTree[jax.ShapeDtypeStruct],
         key: PRNGKeyArray,
         levy_area: type[
             BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
@@ -94,6 +95,14 @@ def __init__(
         self.key = key
         self.levy_area = levy_area
 
+        if any(
+            not isinstance(x, jax.ShapeDtypeStruct) or not is_tuple_of_ints(x.shape)
+            for x in jtu.tree_leaves(self.shape)
+        ):
+            raise ValueError(
+                "UnsafeBrownianPath shapes must be tuples of ints or PyTrees of "
+                "`jax.ShapeDtypeStruct`s with static integer tuple shapes."
+            )
         if any(
             not jnp.issubdtype(x.dtype, jnp.inexact)
             for x in jtu.tree_leaves(self.shape)

diffrax/_brownian/tree.py

@@ -9,6 +9,7 @@
 import jax.random as jr
 import jax.tree_util as jtu
 import lineax.internal as lxi
+import numpy as np
 from jaxtyping import Array, Inexact, PRNGKeyArray, PyTree
 from lineax.internal import complex_to_real_dtype
 
@@ -247,7 +248,7 @@ def __init__(
         t0: RealScalarLike,
         t1: RealScalarLike,
         tol: RealScalarLike,
-        shape: tuple[int, ...] | PyTree[jax.ShapeDtypeStruct],
+        shape: tuple[int | np.integer, ...] | PyTree[jax.ShapeDtypeStruct],
         key: PRNGKeyArray,
         levy_area: type[
             BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
@@ -291,6 +292,14 @@ def __init__(
             if is_tuple_of_ints(shape)
             else shape
         )
+        if any(
+            not isinstance(x, jax.ShapeDtypeStruct) or not is_tuple_of_ints(x.shape)
+            for x in jtu.tree_leaves(self.shape)
+        ):
+            raise ValueError(
+                "VirtualBrownianTree shapes must be tuples of ints or PyTrees of "
+                "`jax.ShapeDtypeStruct`s with static integer tuple shapes."
+            )
         if any(
             not jnp.issubdtype(x.dtype, jnp.inexact)
             for x in jtu.tree_leaves(self.shape)

diffrax/_misc.py

@@ -134,7 +134,9 @@ def split_by_tree(key, tree, is_leaf: Callable[[PyTree], bool] | None = None):
 
 
 def is_tuple_of_ints(obj):
-    return isinstance(obj, tuple) and all(isinstance(x, int) for x in obj)
+    return isinstance(obj, tuple) and all(
+        isinstance(x, (int, np.integer)) and not isinstance(x, bool) for x in obj
+    )
 
 
 def static_select(pred: BoolScalarLike, a: ArrayLike, b: ArrayLike) -> ArrayLike:

test/test_brownian.py

@@ -8,6 +8,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
+import numpy as np
 import pytest
 import scipy.stats as stats
 
@@ -693,3 +694,41 @@ def test_levy_area_reverse_time():
     assert jtu.tree_map(
         lambda fwd, bck: jnp.allclose(fwd, -bck), fwd_increments, back_increments
     )
+
+
+@pytest.mark.parametrize(
+    "ctr", [diffrax.VirtualBrownianTree, diffrax.UnsafeBrownianPath]
+)
+def test_brownian_shape_validation(ctr):
+    import jaxtyping
+
+    key = jr.PRNGKey(0)
+
+    def _make(shape):
+        if ctr is diffrax.VirtualBrownianTree:
+            return ctr(t0=0.0, t1=1.0, tol=1e-3, shape=shape, key=key)
+        return ctr(shape=shape, key=key)
+
+    with pytest.raises((ValueError, jaxtyping.TypeCheckError)):
+        _make((jnp.int32(2),))
+
+    for bad_shape in (
+        jax.ShapeDtypeStruct((jnp.int32(2),), jnp.float32),
+        (True,),
+    ):
+        with pytest.raises(ValueError, match="shapes must be"):
+            _make(bad_shape)
+
+    def _run_once(s):
+        @jax.jit
+        def _eval(p):
+            return p.evaluate(0.0, 0.5)
+
+        return _eval(_make(s))
+
+    for good_shape in (
+        (np.int64(2),),
+        jax.ShapeDtypeStruct((np.int64(2),), jnp.float32),
+    ):
+        assert _run_once(good_shape).shape == (2,)
+        assert _run_once(good_shape).shape == (2,)