From 6d83d2357f863a1387a7bac7eb0bbfed189c33d3 Mon Sep 17 00:00:00 2001 From: Satoshi Kawasaki Date: Thu, 15 Oct 2026 17:57:10 +0900 Subject: [PATCH 1/3] Validate static integer shapes in VirtualBrownianTree and UnsafeBrownianPath (Fixes #707) --- diffrax/_misc.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/diffrax/_misc.py b/diffrax/_misc.py index 8fa07ce1..d8e3efad 100644 --- a/diffrax/_misc.py +++ b/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: From d19e17644eb490215a81f57c79672b54356cf212 Mon Sep 17 00:00:00 2001 From: Satoshi Kawasaki Date: Thu, 15 Oct 2026 18:02:52 +0900 Subject: [PATCH 2/3] Validate static integer shapes in VirtualBrownianTree and UnsafeBrownianPath (Fixes #707) --- diffrax/_brownian/path.py | 11 ++++++++++- diffrax/_brownian/tree.py | 11 ++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/diffrax/_brownian/path.py b/diffrax/_brownian/path.py index 61f49644..0489f2d5 100644 --- a/diffrax/_brownian/path.py +++ b/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) diff --git a/diffrax/_brownian/tree.py b/diffrax/_brownian/tree.py index 8a430668..c952343d 100644 --- a/diffrax/_brownian/tree.py +++ b/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) From c4f004431268cc147d377c640976bfea6b411baa Mon Sep 17 00:00:00 2001 From: Satoshi Kawasaki Date: Thu, 15 Oct 2026 18:04:38 +0900 Subject: [PATCH 3/3] Add regression tests for #707 --- test/test_brownian.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/test/test_brownian.py b/test/test_brownian.py index d33bbda8..f727478b 100644 --- a/test/test_brownian.py +++ b/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,)