Broadcast icdf search state against distribution parameters

The scan carrying the (lo, hi) search state started at the shape of the
quantile value, but logcdf(k) broadcasts k against the distribution
parameters. When the parameters had a larger broadcast shape than the
value, the state changed shape inside scan, raising a ValueError (scalar
value, vector parameters) or silently misbroadcasting (value and
parameters broadcasting to a larger common shape).

Probe logcdf once to obtain the full broadcast shape and start the
search state there. Only the shape of the probe is used, so the probe
computation is pruned from the compiled graph.

Author: jacobpstein

pymc/distributions/dist_math.py

@@ -99,10 +99,9 @@ def discrete_icdf_via_search(logcdf, value, lower, upper=None, *, start=None, ma
     whose quantile function has no closed form (e.g. ``Poisson``, ``Binomial``,
     ``NegativeBinomial``).
 
-    The search uses fixed-length ``scan`` loops (rather than a data-dependent
-    ``while``) so that the resulting graph remains convertible to the C, Numba
-    and JAX backends. Each step is idempotent once an element has converged, so
-    running the full ``max_iters`` iterations does not change the result.
+    The search uses fixed-length ``scan`` loops rather than a data-dependent
+    ``until`` condition. Each step is idempotent once an element has converged,
+    so running the full ``max_iters`` iterations does not change the result.
 
     Parameters
     ----------
@@ -131,12 +130,21 @@ def discrete_icdf_via_search(logcdf, value, lower, upper=None, *, start=None, ma
         of integers exactly representable as float64.
     """
     log_value = pt.log(value)
-    lower = pt.zeros_like(value) + pt.as_tensor_variable(lower, dtype="float64")
+    lower = pt.as_tensor_variable(lower).astype("float64")
+
+    # The scan below carries (lo, hi) as recurrent state, whose shape must stay
+    # fixed across iterations. Since logcdf(k) broadcasts k against the
+    # distribution parameters, the state must start at the full broadcast shape
+    # of value and the parameters. Probing logcdf gives that shape; only the
+    # shape is used, so the probe computation itself is pruned from the graph.
+    result_shape = logcdf(pt.broadcast_arrays(lower, log_value)[0]).shape
+    log_value = pt.broadcast_to(log_value, result_shape)
+    lower = pt.broadcast_to(lower, result_shape)
 
     if upper is None:
         hi_init = lower + 1.0
         if start is not None:
-            start = pt.zeros_like(value) + pt.as_tensor_variable(start, dtype="float64")
+            start = pt.broadcast_to(pt.as_tensor_variable(start).astype("float64"), result_shape)
             hi_init = pt.maximum(pt.floor(start), hi_init)
 
         # Phase 1: geometrically expand an upper bracket until cdf(hi) >= value.
@@ -153,7 +161,7 @@ def expand(lo, hi):
         lo, hi = lo[-1], hi[-1]
     else:
         lo = lower
-        hi = pt.zeros_like(value) + pt.as_tensor_variable(upper, dtype="float64")
+        hi = pt.broadcast_to(pt.as_tensor_variable(upper).astype("float64"), result_shape)
 
     # Phase 2: bisect [lo, hi] keeping the invariant cdf(hi) >= value > cdf(lo - 1).
     def bisect(lo, hi):

tests/distributions/test_discrete.py

@@ -399,6 +399,27 @@ def test_poisson(self):
             st.poisson.ppf(q, 1e4),
         )
 
+    def test_discrete_icdf_batched_params(self):
+        # The search state must be broadcast against the distribution
+        # parameters, not just the quantile, or scan fails / misbroadcasts
+        mu = np.array([2.0, 5.0, 20.0])
+        np.testing.assert_array_equal(
+            icdf(pm.Poisson.dist(mu=mu), 0.5).eval(),
+            st.poisson.ppf(0.5, mu),
+        )
+
+        q = np.array([[0.1], [0.5], [0.9]])
+        n = np.array([10, 40])
+        p = np.array([0.3, 0.6])  # broadcasts with q to shape (3, 2)
+        np.testing.assert_array_equal(
+            icdf(pm.Binomial.dist(n=n, p=p), q).eval(),
+            st.binom.ppf(q, n, p),
+        )
+        np.testing.assert_array_equal(
+            icdf(pm.NegativeBinomial.dist(n=n, p=p), q).eval(),
+            st.nbinom.ppf(q, n, p),
+        )
+
     @pytest.mark.parametrize("n", [2, 3, 4])
     def test_categorical(self, n):
         domain = Domain(range(n), dtype="int64", edges=(0, n))