From 75289c9e30aabd813ced4656480c8e309cde6c29 Mon Sep 17 00:00:00 2001 From: Francesco Muia <38326626+fmuia@users.noreply.github.com> Date: Mon, 17 Aug 2026 23:31:24 +0000 Subject: [PATCH 1/2] fix: pass unmapped params to jax callables as scalars (#1092) Non-trialwise parameters reach the vmapped single-trial callable as (1,)-shaped tensors (bambi emits those for intercept-only parameters, and PyMC adds the axis to a parameter fixed to a constant) and `in_axes=None` forwards them un-sliced. Arithmetic in the callable then yields a (1,) per-trial result, so the vmapped output is (n_obs, 1) while `LANLogpOp` declares `pt.vector()`. Nothing caught that. On the default backend the oversized array is written into a vector-typed storage cell and read back as uninitialised memory (measured: a total log-likelihood of 1.4e+307 on a 200-trial model). Under the JAX linker it broadcasts against the `ensure_positive_ndt` mask into (n_obs, n_obs), so `sampler="numpyro"` samples to completion and returns a log-likelihood of 2.4e+270; when a parameter is fixed it instead surfaces as a bare `AssertionError` in PyTensor's JAX `SpecifyShape` dispatch, which is what #1092 reports. Every jax-callable model with at least one non-trialwise parameter was affected, not only the fixed-parameter case. Squeeze unmapped inputs so every parameter arrives as a scalar --- the convention vmap already gives mapped parameters, the one the ONNX wrapper follows, and the one the docstrings already promised. Only a rank-1, length-1 input is scalarized, so an unmapped parameter of any other shape is passed through rather than silently flattened. BEHAVIOUR CHANGE for user-supplied JAX log-likelihoods: a callable that treats a non-trialwise parameter as an array --- `a[0]`, `a.shape[0]`, `len(a)`, `jnp.concatenate([a, ...])` --- must now use the parameter directly. Such callables previously worked only when HSSM happened to classify that parameter as non-trialwise; putting a formula on it already raised the same error before this change. They now fail immediately with an `IndexError`/`TypeError` at the first logp evaluation rather than running. ONNX-backed paths are unaffected (verified bit-identical logp before and after). Closes #1092 --- src/hssm/distribution_utils/jax.py | 49 +++++++++++++++ tests/distribution_utils/test_jax.py | 94 ++++++++++++++++++++++++++++ tests/test_hssm.py | 68 ++++++++++++++++++++ 3 files changed, 211 insertions(+) diff --git a/src/hssm/distribution_utils/jax.py b/src/hssm/distribution_utils/jax.py index 68da2685c..ecf5c75b9 100644 --- a/src/hssm/distribution_utils/jax.py +++ b/src/hssm/distribution_utils/jax.py @@ -291,6 +291,35 @@ def make_jax_logp_ops( return LANLogpOp(logp, logp_nojit, LANLogpVJPOp(logp_vjp, n_params), n_params) +def _scalarize_unmapped_inputs(logp: Callable, in_axes: list[int | None]) -> Callable: + """Pass unmapped (``in_axes=None``) inputs to `logp` as scalars. + + ``vmap`` slices the batch axis off mapped inputs, so a mapped parameter + already arrives as a scalar (and ``data`` as a single row). Unmapped + inputs are forwarded whole, which for a non-trialwise parameter means a + ``(1,)``-shaped array rather than a scalar. Squeezing that one axis keeps + the single-trial calling convention uniform. + + Only a rank-1, length-1 input is scalarized. A non-trialwise parameter is + always shaped ``(1,)`` at this point (PyMC adds the axis to the scalar RV), + so anything else is a shape this layer does not understand and is passed + through untouched rather than silently flattened. + """ + + def _scalarize(inp, axis): + if axis is not None: + return inp + arr = jnp.asarray(inp) + return arr.reshape(()) if arr.ndim == 1 and arr.shape[0] == 1 else arr + + def scalarized(*inputs): + return logp( + *(_scalarize(inp, axis) for inp, axis in zip(inputs, in_axes, strict=True)) + ) + + return scalarized + + @overload def make_jax_logp_funcs_from_callable( logp: Callable, @@ -324,6 +353,13 @@ def make_jax_logp_funcs_from_callable( extra_fields are optional additional fields that can be used in the likelihood computation. The `data` argument is a two-column numpy array with response time and response. + + When `vmap` is `True`, the callable is a *single-trial* function: `data` + is one row, **every** parameter arrives as a 0-d scalar -- trialwise + ones because `vmap` slices off the batch axis, non-trialwise ones + because they are squeezed from the `(1,)` shape PyMC gives a scalar + random variable -- and the return value must be a scalar for that + trial. Do not index a parameter (`a[0]`, `len(a)`); use it directly. vmap: If `True`, the function will be vectorized using JAX's vmap. If `False`, the function is assumed to be already vectorized. @@ -377,6 +413,19 @@ def make_jax_logp_funcs_from_callable( "No vmap is needed in your use case, since all parameters are scalars." ) + # Non-trialwise parameters reach the single-trial callable as + # `(1,)`-shaped tensors (bambi emits those for intercept-only + # parameters; see the broadcast note in `make_distribution`) and + # `in_axes=None` passes them through un-sliced. Any arithmetic in the + # callable then produces a `(1,)` per-trial result, so the vmapped + # output is `(n_obs, 1)` while `LANLogpOp` declares a vector. Squeeze + # them here so every parameter arrives as a scalar --- the same + # convention vmap already gives mapped parameters, and what the ONNX + # wrapper does for itself in its `not params_only` branch (see the + # `inp.squeeze()` in `make_jax_logp_funcs_from_onnx`; its `params_only` + # branch does *not* squeeze, and is not covered by this). + logp = _scalarize_unmapped_inputs(logp, in_axes) + if return_jit: return make_vmap_func( logp, diff --git a/tests/distribution_utils/test_jax.py b/tests/distribution_utils/test_jax.py index b8644d42a..dd99fe607 100644 --- a/tests/distribution_utils/test_jax.py +++ b/tests/distribution_utils/test_jax.py @@ -419,3 +419,97 @@ def jax_callable(data, v): with pytest.raises(ValueError, match="previously applied with data"): logp_op(None, v) + + +def test_callable_receives_scalar_unmapped_params(): + """Non-trialwise params reach the single-trial callable as scalars. + + Regression test for #1092. Bambi emits ``(1,)``-shaped tensors for + intercept-only parameters, and ``in_axes=None`` forwards them un-sliced. + Without squeezing, the callable returns a ``(1,)`` per-trial value, the + vmapped output is ``(n_obs, 1)`` instead of ``(n_obs,)``, and the graph + later fails in PyTensor's JAX ``SpecifyShape`` with a bare + ``AssertionError``. + """ + seen = {} + + def logp(data, v, a): + seen["v"] = jax.numpy.shape(v) + seen["a"] = jax.numpy.shape(a) + return data[0] * v + a + + logp_vec, _, _ = make_jax_logp_funcs_from_callable( + logp, vmap=True, params_is_reg=[True, False] + ) + + n_obs = 7 + data = np.ones((n_obs, 2), dtype=np.float32) + v = np.full((n_obs,), 2.0, dtype=np.float32) # trialwise -> mapped + a = np.array([3.0], dtype=np.float32) # non-trialwise -> (1,)-shaped + + out = logp_vec(data, v, a) + + assert seen["v"] == (), "mapped params should arrive as scalars" + assert seen["a"] == (), "unmapped params should be squeezed to scalars" + assert np.shape(out) == (n_obs,), ( + f"vmapped logp must be 1-D to match LANLogpOp's declared vector output, " + f"got {np.shape(out)}" + ) + np.testing.assert_allclose(np.asarray(out), np.full((n_obs,), 5.0), rtol=1e-6) + + +def test_scalarize_handles_params_only_signature(): + """`params_only=True` omits the data entry in `in_axes`; alignment must hold. + + Without a leading `data` axis the zip over `(inputs, in_axes)` is offset by + one relative to the usual case, so a non-trialwise parameter in first + position is the discriminating arrangement. + """ + seen = {} + + def cpn_logp(v, a, z, t): + seen["shapes"] = tuple(jax.numpy.shape(p) for p in (v, a, z, t)) + return v + a + z + t + + logp_vec, _, _ = make_jax_logp_funcs_from_callable( + cpn_logp, + vmap=True, + params_is_reg=[False, True, False, False], + params_only=True, + ) + + n_obs = 5 + a = np.linspace(1.0, 1.4, n_obs, dtype=np.float32) # trialwise -> mapped + out = logp_vec( + np.array([0.3], dtype=np.float32), # non-trialwise -> (1,)-shaped + a, + np.array([0.5], dtype=np.float32), + np.array([0.2], dtype=np.float32), + ) + + assert seen["shapes"] == ((), (), (), ()) + assert np.shape(out) == (n_obs,) + np.testing.assert_allclose(np.asarray(out), 0.3 + a + 0.5 + 0.2, rtol=1e-6) + + +def test_unmapped_inputs_that_are_not_length_one_pass_through(): + """Only the `(1,)` axis PyMC adds is squeezed; other shapes are untouched.""" + seen = {} + + def logp(data, v, w): + seen["w"] = jax.numpy.shape(w) + return data[0] * v + jax.numpy.sum(w) + + logp_vec, _, _ = make_jax_logp_funcs_from_callable( + logp, vmap=True, params_is_reg=[True, False] + ) + + n_obs = 4 + out = logp_vec( + np.ones((n_obs, 2), dtype=np.float32), + np.full((n_obs,), 2.0, dtype=np.float32), + np.array([1.0, 2.0, 3.0], dtype=np.float32), # rank-1 but not length-1 + ) + + assert seen["w"] == (3,), "a multi-element unmapped input must not be flattened" + assert np.shape(out) == (n_obs,) diff --git a/tests/test_hssm.py b/tests/test_hssm.py index 07840e4e8..cef77ccfb 100644 --- a/tests/test_hssm.py +++ b/tests/test_hssm.py @@ -5,6 +5,7 @@ import numpy as np import pymc as pm +import pytensor import pytest import xarray as xr from pymc.variational import Approximation @@ -690,3 +691,70 @@ def test_is_choice_only_and_deadline(data_ddm): assert len(model_with_deadline.response) == 2 assert model_with_deadline.response_c == "c(response, deadline)" assert model_with_deadline.response_str == "response,deadline" + + +@pytest.mark.parametrize("mode", [None, "JAX"], ids=["default", "jax_linker"]) +def test_jax_callable_non_trialwise_params_logp(data_ddm, mode): + """A custom JAX likelihood must see non-trialwise params as scalars (#1092). + + ``v`` is the parent (trialwise); ``a``, ``z`` and ``t`` are fixed, so PyMC + hands them to the Op as ``(1,)``-shaped tensors. Before #1092 those reached + the single-trial callable un-squeezed, so each trial produced a ``(1,)`` + value and the vmapped result was ``(n_obs, 1)`` instead of ``(n_obs,)``. + Nothing caught it: the default backend read the oversized buffer back as + uninitialised memory, and the JAX linker broadcast it to ``(n_obs, n_obs)`` + -- both silently wrong. This pins the values against a NumPy reference on + both linkers. + """ + import jax.numpy as jnp + + a_fixed, z_fixed, t_fixed = 1.3, 0.5, 0.2 + + def single_trial_logp(data, v, a, z, t): + rt = data[0] + mu = jnp.log(a) - jnp.log(v**2 + 0.25) + t + z + log_rt = jnp.log(jnp.maximum(rt, 1e-6)) + return -log_rt - 0.5 * ((log_rt - mu) / 0.45) ** 2 - 0.8 + + model = HSSM( + data=data_ddm, + model="ddm", + loglik=single_trial_logp, + loglik_kind="approx_differentiable", + model_config={"backend": "jax"}, + p_outlier=0, + a=a_fixed, + z=z_fixed, + t=t_fixed, + ) + + pymc_model = model.pymc_model + value_vars = list(pymc_model.value_vars) + assert [v.name for v in value_vars] == ["v_interval__"], ( + "only the parent should be free; the rest are fixed constants" + ) + + # Evaluate at an arbitrary point in the unconstrained space, then map it + # back through the model's own transform to build the reference. + unconstrained_v = np.array(0.3, dtype=pytensor.config.floatX) + rv = pymc_model.values_to_rvs[value_vars[0]] + transform = pymc_model.rvs_to_transforms[rv] + v = float(np.asarray(transform.backward(unconstrained_v, *rv.owner.inputs).eval())) + + rt = data_ddm["rt"].to_numpy() + mu = np.log(a_fixed) - np.log(v**2 + 0.25) + t_fixed + z_fixed + log_rt = np.log(np.maximum(rt, 1e-6)) + expected = -log_rt - 0.5 * ((log_rt - mu) / 0.45) ** 2 - 0.8 + # HSSM replaces the log-likelihood with LOGP_LB where rt <= t. + expected = np.where(rt - t_fixed <= 1e-15, -66.1, expected) + + observed_logp = pymc_model.logp(sum=False)[-1] + fn = pytensor.function( + value_vars, observed_logp, mode=mode, on_unused_input="ignore" + ) + actual = np.asarray(fn(unconstrained_v)) + + assert actual.shape == (len(data_ddm),) + # Loose rtol so the test is valid under either floatX; the pre-fix failures + # it guards against are off by ~300 orders of magnitude or a whole axis. + np.testing.assert_allclose(actual, expected, rtol=1e-4) From c7f107f40051256d515f5c91651f3fa452bbda67 Mon Sep 17 00:00:00 2001 From: Francesco Muia <38326626+fmuia@users.noreply.github.com> Date: Tue, 18 Aug 2026 00:07:58 +0000 Subject: [PATCH 2/2] fix: validate LANLogpOp output shape instead of failing opaquely (#1092) `LANLogpOp` declares a vector output but never checked what the wrapped JAX function actually returned, so a mismatch surfaced far from its cause: - default backend: the oversized array was written into a vector-typed storage cell, yielding garbage values (~1e306) rather than an error - JAX linker: a bare, message-less `AssertionError` inside PyTensor's `SpecifyShape` dispatch -- its ndim check is the only bare assert there; a size mismatch already raises a descriptive `ValueError` Check the output ndim on both paths and raise a `ValueError` naming the expected and actual shapes. The hint follows the direction of the mismatch: too many dimensions points at a per-trial value that is a length-1 array, too few at a function that reduced over trials. Results without an `ndim` attribute (a Python scalar, a list) fall back to `np.ndim` so they are checked too, rather than skipping the guard. On the JAX linker the check runs at trace time, so it costs nothing per call; in `perform` it is ~45 ns, about 0.002% of a realistic evaluation. `LANLogpVJPOp` deliberately gets no equivalent guard -- every VJP HSSM builds comes from `jax.vjp`, so cotangent shapes match the declared output types by construction. A comment records that. Also adds coverage for `extra_fields` in_axes alignment: `hssm.py` appends one entry per extra field to `params_is_trialwise`, so the zip over inputs and in_axes must consume both sequences exactly. Addresses the second half of #1092. --- src/hssm/distribution_utils/jax.py | 71 +++++++++++++++++++++-- tests/distribution_utils/test_jax.py | 84 ++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 6 deletions(-) diff --git a/src/hssm/distribution_utils/jax.py b/src/hssm/distribution_utils/jax.py index ecf5c75b9..4b9de003c 100644 --- a/src/hssm/distribution_utils/jax.py +++ b/src/hssm/distribution_utils/jax.py @@ -1,6 +1,6 @@ """Utilities for wraping JAX likelihoods in Pytensor Ops.""" -from typing import Callable, Literal, overload +from typing import Any, Callable, Literal, overload import jax.numpy as jnp import numpy as np @@ -14,6 +14,49 @@ from .func_utils import make_vjp_func, make_vmap_func +def _check_logp_ndim(result: Any, expected_ndim: int) -> None: + """Validate the ndim of a wrapped log-likelihood's output. + + ``LANLogpOp`` declares a vector output, but nothing forces the wrapped + JAX function to honour that. When it does not, the mismatch surfaces far + from its cause: PyTensor's default backend writes the oversized array into + a vector-typed storage cell (yielding garbage values), while the JAX + linker fails inside ``SpecifyShape`` with a bare, message-less + ``AssertionError``. Checking here turns both into a diagnosable error. + """ + actual_ndim = getattr(result, "ndim", None) + if actual_ndim is None: + # A plain Python scalar, or a list/tuple. `np.ndim` handles both, but + # raises on ragged input -- there is nothing meaningful to check then, + # and its "inhomogeneous shape" error would be worse than the one this + # guard exists to replace. + try: + actual_ndim = np.ndim(result) + except Exception: + return + if actual_ndim == expected_ndim: + return + + shape = getattr(result, "shape", "unknown") + if actual_ndim < expected_ndim: + hint = ( + "This usually means the function reduced over trials -- returning " + "`jnp.sum(...)` instead of one value per trial. A likelihood whose " + "value is shared across trials should `.reshape((1,))` instead." + ) + else: + hint = ( + "This usually means the single-trial function returns a length-1 " + "array instead of a scalar -- for example because a parameter " + "reached it as a `(1,)`-shaped array rather than a scalar." + ) + raise ValueError( + "The log-likelihood wrapped by LANLogpOp must return one value per " + f"trial ({expected_ndim}-D), but it returned a {actual_ndim}-D result " + f"of shape {shape}.\n" + hint + ) + + class LANLogpVJPOp(Op): # pylint: disable=W0223 """Wraps the VJP of a JAX log-likelihood function in a pytensor Op. @@ -82,6 +125,10 @@ def perform(self, node, inputs, output_storage): output_storage. There is one storage cell for each output of the Op. """ + # No `_check_logp_ndim` here, unlike `LANLogpOp.perform`: every VJP + # HSSM builds comes from `make_vjp_func`, i.e. from `jax.vjp`, so each + # cotangent already matches the shape of the input it belongs to -- + # which is exactly the type `make_node` declares for it. results = self.logp_vjp(*inputs[:-1], gz=inputs[-1]) for i, result in enumerate(results): @@ -182,8 +229,9 @@ def perform(self, node, inputs, output_storage): output_storage. There is one storage cell for each output of the Op. """ - result = self.logp(*inputs) - output_storage[0][0] = np.asarray(result, dtype=node.outputs[0].dtype) + result = np.asarray(self.logp(*inputs), dtype=node.outputs[0].dtype) + _check_logp_ndim(result, node.outputs[0].type.ndim) + output_storage[0][0] = result def pullback(self, inputs, outputs, cotangents): """Construct the graph for the VJP (reverse-mode gradient) of the Op. @@ -233,9 +281,20 @@ def pullback(self, inputs, outputs, cotangents): # Unwraps the JAX function for compilation with the JAX linker (e.g. sampling # through numpyro/blackjax, or pm.fit / pm.sample with backend="jax"). @jax_funcify.register(LANLogpOp) -def lan_logp_op_dispatch(op, **kwargs): # pylint: disable=W0613 - """Return the non-jitted forward function for the JAX linker.""" - return op.logp_nojit +def lan_logp_op_dispatch(op, node, **kwargs): # pylint: disable=W0613 + """Return the non-jitted forward function for the JAX linker. + + The output is shape-checked at trace time, so the guard costs nothing per + call but still fires before a bad shape reaches ``SpecifyShape``. + """ + expected_ndim = node.outputs[0].type.ndim + + def lan_logp_jax(*inputs): + result = op.logp_nojit(*inputs) + _check_logp_ndim(result, expected_ndim) + return result + + return lan_logp_jax # Required when PyTensor differentiates LANLogpOp symbolically (e.g. ADVI diff --git a/tests/distribution_utils/test_jax.py b/tests/distribution_utils/test_jax.py index dd99fe607..36548c256 100644 --- a/tests/distribution_utils/test_jax.py +++ b/tests/distribution_utils/test_jax.py @@ -513,3 +513,87 @@ def logp(data, v, w): assert seen["w"] == (3,), "a multi-element unmapped input must not be flattened" assert np.shape(out) == (n_obs,) + + +def test_extra_fields_do_not_break_in_axes_alignment(): + """`params_is_reg` covers extra_fields, so `in_axes` must line up exactly. + + ``hssm.py`` appends one ``True`` per extra field to + ``params_is_trialwise`` before handing it to the vmap layer, so the zip + over ``(inputs, in_axes)`` must consume both sequences without a leftover. + """ + seen = {} + + def logp(data, v, a, ef1, ef2): + seen["shapes"] = ( + jax.numpy.shape(v), + jax.numpy.shape(a), + jax.numpy.shape(ef1), + jax.numpy.shape(ef2), + ) + return data[0] * v + a + ef1 + ef2 + + logp_vec, _, _ = make_jax_logp_funcs_from_callable( + logp, vmap=True, params_is_reg=[True, False, True, True] + ) + + n_obs = 5 + out = logp_vec( + np.ones((n_obs, 2), dtype=np.float32), + np.full((n_obs,), 2.0, dtype=np.float32), # trialwise param + np.array([3.0], dtype=np.float32), # non-trialwise param + np.arange(n_obs, dtype=np.float32), # extra field + np.ones((n_obs,), dtype=np.float32), # extra field + ) + + assert seen["shapes"] == ((), (), (), ()) + assert np.shape(out) == (n_obs,) + + +@pytest.mark.parametrize("mode", [None, "JAX"], ids=["default", "jax_linker"]) +def test_logp_output_ndim_is_validated(mode): + """A wrong-ndim log-likelihood raises a diagnosable error on both linkers. + + Regression test for the second half of #1092. ``LANLogpOp`` declares a + vector output but never checked it, so a mismatch surfaced far from its + cause: garbage values on the default backend (an oversized array written + into a vector-typed storage cell), and a bare, message-less + ``AssertionError`` inside PyTensor's JAX ``SpecifyShape`` dispatch. + """ + n_obs = 6 + + def bad_logp(data, v, a): + # Pre-vectorised, but returns (n_obs, 1) rather than (n_obs,). + return (data[:, 0] * v + a)[:, None] + + logp_op = make_jax_logp_ops( + *make_jax_logp_funcs_from_callable(bad_logp, vmap=False) + ) + v = pt.vector("v") + a = pt.vector("a") + out = logp_op(pt.as_tensor_variable(np.ones((n_obs, 2), dtype=np.float32)), v, a) + + # Compiling succeeds on both linkers; the guard fires when the wrapped + # function actually runs, so keep the call as the only guarded statement. + fn = pytensor.function([v, a], out, mode=mode) + with pytest.raises(ValueError, match="must return one value per trial"): + fn(np.full((n_obs,), 2.0, dtype=np.float32), np.array([1.0], dtype=np.float32)) + + +def test_logp_output_ndim_error_names_the_right_cause(): + """The hint must match the direction of the mismatch, not assume one.""" + + def too_few_dims(data, v, a): + # Reduces over trials instead of returning one value per trial. + return jax.numpy.sum(data[:, 0] * v + a) + + logp_op = make_jax_logp_ops( + *make_jax_logp_funcs_from_callable(too_few_dims, vmap=False) + ) + v = pt.vector("v") + a = pt.vector("a") + out = logp_op(pt.as_tensor_variable(np.ones((4, 2), dtype=np.float32)), v, a) + + fn = pytensor.function([v, a], out) + with pytest.raises(ValueError, match="reduced over trials"): + fn(np.full((4,), 2.0, dtype=np.float32), np.array([1.0], dtype=np.float32))