Add GumbelRV JAX implementation

Author: rlouf

aesara/link/jax/dispatch/random.py

@@ -110,6 +110,7 @@ def sample_fn(rng, size, dtype, *parameters):
 
 
 @jax_sample_fn.register(aer.CauchyRV)
+@jax_sample_fn.register(aer.GumbelRV)
 @jax_sample_fn.register(aer.LaplaceRV)
 @jax_sample_fn.register(aer.LogisticRV)
 @jax_sample_fn.register(aer.NormalRV)

tests/link/jax/test_random.py

@@ -123,6 +123,22 @@ def test_random_updates(rng_ctor):
             "gamma",
             lambda a, b: (a, 0.0, b),
         ),
+        (
+            aer.gumbel,
+            [
+                set_test_value(
+                    at.lvector(),
+                    np.array([1, 2], dtype=np.int64),
+                ),
+                set_test_value(
+                    at.dscalar(),
+                    np.array(1.0, dtype=np.float64),
+                ),
+            ],
+            (2,),
+            "gumbel_r",
+            lambda *args: args,
+        ),
         (
             aer.laplace,
             [