Implement faster Multinomial JAX dispatch #1316

educhesne · 2025-03-23T13:15:27Z

Description

Replace the call to numpyro.distributions.util.multinomial by a custom function written in jax implementing Sequential conditional binomial sampling

Related Issue

Closes Implement faster Multinomial JAX dispatch #868
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pytensor--1316.org.readthedocs.build/en/1316/

Replace the call to numpyro.distributions.util.multinomial by a custom function

ricardoV94 · 2025-03-24T09:40:12Z

pytensor/link/jax/dispatch/random.py

@@ -450,6 +443,42 @@ def sample_fn(rng, size, dtype, n, p):
    return sample_fn


+def _jax_multinomial(n, p, shape=None, key=None):
+    if jnp.shape(n) != jnp.shape(p)[:-1]:


Theres' some redundancy here. If size (not shape btw) is provided we should just broadcast n to size and p to size + p.shape[-1]. Only if size is not provided should we broadcast n and p[...: -1] together (basically finding the implicit size)

Yes, besides there is actually no need to broadcast p at all; jax.random.binomial broadcasts it for us.
About the name shape: I kept the signature of the numpyro function but I can change it

you might be right about p although I guess it's more readable if you broadcast explicitly

pytensor/link/jax/dispatch/random.py

ricardoV94 · 2025-03-24T13:21:22Z

pytensor/link/jax/dispatch/random.py

        return sample

    return sample_fn


+def _jax_multinomial(n, p, size=None, key=None):
+    if size is not None:
+        broadcast_shape_n = jax.lax.broadcast_shapes(jnp.shape(n), size)


you don't need to broadcast size with n, that's not allowed by the RandomVariable semantics. If size is provided n must be broadcastable to size, but not the other way. You can's say n=ones((5, 1)), size=(1, 3)

Indeed, my bad

codecov · 2025-03-24T13:41:04Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.98%. Comparing base (2e9d502) to head (81f0b12).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1316   +/-   ##
=======================================
  Coverage   81.98%   81.98%           
=======================================
  Files         188      188           
  Lines       48475    48489   +14     
  Branches     8673     8673           
=======================================
+ Hits        39740    39756   +16     
+ Misses       6583     6582    -1     
+ Partials     2152     2151    -1

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/random.py	`94.01% <100.00%> (+1.28%)`	⬆️

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

educhesne · 2025-03-24T19:32:57Z

It seems the sample_fn function is never called with size=None even though pt.random.multinomial(..., size=None) is in the test. Could it be that the None size is catched in between ? (there is a static_shape happening in jax_funcify_RandomVariable I don't fully understand)
Should I remove the part which is not reached by the coverage test ?

ricardoV94 · 2025-03-24T20:39:22Z

It can be reached in some cases when for instance you have size=x.shape, and x is an input variable without known shape.

ricardoV94 · 2025-03-24T20:40:45Z

pytensor/link/jax/dispatch/random.py

        return sample

    return sample_fn


+def _jax_multinomial(n, p, size=None, key=None):


You can inline this in the dispatch for MultinomialRV. Also key=None is not really valid but we don't need default parameters anyway

educhesne · 2025-03-25T07:55:22Z

It can be reached in some cases when for instance you have size=x.shape, and x is an input variable without known shape.

I'm sorry I don't know how to update the test to cover this case

Locally I've tried this without success:

p = np.array([[0.3, 0.7, 0.0], [0.1, 0.4, 0.5]])
size = (10_000, 2)
n = np.broadcast_to(np.array([10, 40]), size)
x = pt.matrix('x')

g = pt.random.multinomial(n, p, rng=rng, size=x.shape)
g_fn = compile_random_function([x], g, mode="JAX")
samples = g_fn(np.zeros(size))

ricardoV94 · 2025-03-25T15:22:25Z

Sorry what I sad was non-sense. Here is a test with size=None

import pytensor.tensor as pt

from tests.link.jax.test_random import compile_random_function

n = pt.vector("n")
p = pt.vector("p")

g = pt.random.multinomial(n, p, size=None)
g_fn = compile_random_function([n, p], g, mode="JAX")

size is None, and it can't be inferred from n because n doesn't have a static shape.

…y function

educhesne · 2025-03-25T15:44:34Z

Thank you for your feedbacks !

ricardoV94 · 2025-03-25T15:53:55Z

Thank you for the feature!!!

ricardoV94 · 2025-03-25T16:11:58Z

Awesome

Etienne Duchesne added 2 commits March 23, 2025 14:09

Implement multinomial JAX dispatch directly in jax

af9803c

Replace the call to numpyro.distributions.util.multinomial by a custom function

improve test coverage on multinomial jax dispatch

b551daf

ricardoV94 requested changes Mar 24, 2025

View reviewed changes

ricardoV94 reviewed Mar 24, 2025

View reviewed changes

pytensor/link/jax/dispatch/random.py Outdated Show resolved Hide resolved

educhesne and others added 3 commits March 24, 2025 13:37

Merge branch 'main' into multinomial_jax_dispatch

f373d99

broadcast explicitly the arguments in mulinomial jax dispatch

9e6ac0b

fix typo error in multinomial jax dispatch

4172752

ricardoV94 reviewed Mar 24, 2025

View reviewed changes

fix broadcast in multinomial jax dispatch

aef0d4b

ricardoV94 reviewed Mar 24, 2025

View reviewed changes

improve test coverage in multinomial jax dispatch and inline auxiliar…

81f0b12

…y function

ricardoV94 approved these changes Mar 25, 2025

View reviewed changes

ricardoV94 added jax enhancement New feature or request random variables labels Mar 25, 2025

ricardoV94 merged commit 8a7356c into pymc-devs:main Mar 25, 2025
71 checks passed

educhesne deleted the multinomial_jax_dispatch branch March 25, 2025 16:28

Implement faster Multinomial JAX dispatch #1316

Implement faster Multinomial JAX dispatch #1316

Uh oh!

Conversation

educhesne commented Mar 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

ricardoV94 Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

educhesne Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

educhesne Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

educhesne commented Mar 24, 2025

Uh oh!

ricardoV94 commented Mar 24, 2025

Uh oh!

ricardoV94 Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

educhesne commented Mar 25, 2025

Uh oh!

ricardoV94 commented Mar 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

educhesne commented Mar 25, 2025

Uh oh!

ricardoV94 commented Mar 25, 2025

Uh oh!

Uh oh!

ricardoV94 commented Mar 25, 2025

Uh oh!

Uh oh!

educhesne commented Mar 23, 2025 •

edited

Loading

educhesne Mar 24, 2025 •

edited

Loading

codecov bot commented Mar 24, 2025 •

edited

Loading

ricardoV94 commented Mar 25, 2025 •

edited

Loading