Implement faster Multinomial JAX dispatch#1316

Merged

ricardoV94 merged 7 commits intopymc-devs:mainfrom

educhesne:multinomial_jax_dispatch

Mar 25, 2025

Commits on Mar 23, 2025

Implement multinomial JAX dispatch directly in jax
Etienne Duchesne
committed
improve test coverage on multinomial jax dispatch
Etienne Duchesne
committed

Commits on Mar 24, 2025

Merge branch 'main' into multinomial_jax_dispatch
educhesne
authored
broadcast explicitly the arguments in mulinomial jax dispatch
Etienne Duchesne
committed
fix typo error in multinomial jax dispatch
Etienne Duchesne
committed
fix broadcast in multinomial jax dispatch
Etienne Duchesne
committed

Commits on Mar 25, 2025

improve test coverage in multinomial jax dispatch and inline auxiliary function
Etienne Duchesne
committed