Issue 1092/fix jax callable param

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
@@ -291,6 +350,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 +412,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 +472,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,

tests/distribution_utils/test_jax.py

@@ -419,3 +419,181 @@ 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,)
+
+
+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))