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

[akatsukiwork0320]

Fixes #707.

Problem & Root Cause

Both VirtualBrownianTree and UnsafeBrownianPath store shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True).

1. When a 0-d JAX array such as jnp.int32(2) is passed inside jax.ShapeDtypeStruct(shape=(jnp.int32(2),), dtype=jnp.float32), ShapeDtypeStruct.__init__ converts shape via tuple(shape) without canonicalizing elements to Python int, leaving (Array(2, dtype=int32),) inside the static PyTreeDef metadata of the Brownian path.
   • In the JIT compiling VirtualBrownianTree.evaluate() hangs indefinitely when using jnp.int32 to pass shape to ShapeDtypeStruct #707 MWE (path.evaluate called inside an outer @jax.jit function eval_path that is defined and traced twice with two separately constructed paths), the second call triggers pjit cache comparison (_infer_params_cached) between the cached PyTreeDef and the new PyTreeDef while the outer JIT trace (trace_to_jaxpr_nocache) is active. Comparing Array(2) == Array(2) under an active trace produces a traced boolean (DynamicJaxprTracer with aval=bool[], which raises TracerBoolConversionError on boolean conversion) rather than a concrete bool, causing PyTreeDef equality to fail:
     • On JAX 0.8.0 (as reported in the issue thread) and JAX 0.11.2 (tested locally), this raises ValueError: Exception raised while checking equality of metadata fields of pytree node....
     • On JAX 0.6.2 and 0.7.2 (as reported in the issue thread), this hangs during the second JIT compilation (plausibly due to re-entrant dispatch during C++ pjit cache lookup, though not verified with a C++ stack trace).
2. Additionally, passing shape=(jnp.int32(2),) or shape=(np.int64(2),) directly as a tuple fails is_tuple_of_ints(shape) (which only checked isinstance(x, int)), leaving self.shape un-wrapped as a tuple of integers and raising the misleading error ValueError: VirtualBrownianTree dtypes all have to be floating-point.

Solution & Design Choices

1. diffrax/_misc.py (is_tuple_of_ints):
   • Accept static Python int and NumPy integers (isinstance(x, (int, np.integer)) and not isinstance(x, bool)). This preserves support for jax.ShapeDtypeStruct((np.int64(2),), ...) and also allows direct tuples like shape=(np.int64(2),) to work with lxi.default_floating_dtype() (respecting jax_enable_x64), while rejecting bool, JAX Arrays, and JAX tracers.
2. VirtualBrownianTree.__init__ and UnsafeBrownianPath.__init__:
   • The shape annotation accepts tuple[int | np.integer, ...] | PyTree[jax.ShapeDtypeStruct] so that direct tuples of NumPy integers match the runtime acceptance and pass the repository's jaxtyping/beartype checks. This contract is covered by the new regression tests.
   • Validate that every leaf x in jtu.tree_leaves(self.shape) is a jax.ShapeDtypeStruct with is_tuple_of_ints(x.shape), raising a clear ValueError at construction time before any JIT compilation occurs. Existing ShapeDtypeStruct leaves are kept as-is (preserving attributes such as weak_type).
   • (Note: If duck-typed leaves with .shape and .dtype other than jax.ShapeDtypeStruct should also be accepted, the check can be relaxed to hasattr(x, "shape") and hasattr(x, "dtype") and is_tuple_of_ints(x.shape).)
3. Tests & Validation (test/test_brownian.py):
   • Added import numpy as np and test_brownian_shape_validation covering both VirtualBrownianTree and UnsafeBrownianPath for rejected jnp.int32 / bool shapes and accepted np.int64 shapes across repeated outer @jax.jit evaluations.
   • Validation on Windows with Python 3.12.6, JAX/jaxlib 0.11.2:
     • 15 selected tests from test/test_brownian.py passed with the default jaxtyping/beartype hook: test_shape_and_dtype, test_levy_area_reverse_time, and test_brownian_shape_validation (35 deselected; the remaining 35 statistical/spline/WHK interpolation tests in test/test_brownian.py — test_statistics, test_conditional_statistics, test_spline, test_whk_interpolation — were not completed locally (no completed passing result) and are left to CI).
     • The two new shape-validation cases also passed with the hook disabled using -o addopts="" (48 deselected).
     • 66 selected related tests across six other test files (test_citation.py, test_adaptive_stepsize_controller.py, test_adjoint.py, test_solver.py, test_underdamped_langevin.py, test_sde2.py) passed (77 deselected). This was a selected regression run, not the full suite.
     • Ruff 0.13.0 check and format --check passed for all four changed files.
     • Pyright 1.1.405 reported the same existing unnecessary type-ignore error in diffrax/_misc.py (reportUnnecessaryTypeIgnoreComment at _misc.py:175 on base / _misc.py:177 on patched) on both the base and patched code, with no additional diagnostics in the four checked files. The type check is therefore not a clean pass (exit_code = 1 on both base and patched).