Introduce signature instead of ndim_supp and ndims_params

Author: ricardoV94

pytensor/link/jax/dispatch/random.py

@@ -270,13 +270,19 @@ def sample_fn(rng, size, dtype, *parameters):
 
 
 @jax_sample_fn.register(ptr.ChoiceRV)
-def jax_funcify_choice(op):
+def jax_funcify_choice(op: ptr.ChoiceRV):
     """JAX implementation of `ChoiceRV`."""
 
+    p_none = op.p_none
+
     def sample_fn(rng, size, dtype, *parameters):
         rng_key = rng["jax_state"]
         rng_key, sampling_key = jax.random.split(rng_key, 2)
-        (a, p, replace) = parameters
+        if p_none:
+            (a, replace) = parameters
+            p = None
+        else:
+            (a, p, replace) = parameters
         smpl_value = jax.random.choice(sampling_key, a, size, replace, p)
         rng["jax_state"] = rng_key
         return (rng, smpl_value)