Fix jax.random.categorical axis parameter for JAX 0.10.1 compatibility#1314

jmr wants to merge 1 commit into

sotetsuk:mainsotetsuk/pgx:mainfrom

jmr:jax-randomjmr/pgx:jax-randomCopy head branch name to clipboard