Add ChiSquareRV JAX implementation

Author: larryshamalama

aesara/link/jax/dispatch/random.py

@@ -369,7 +369,6 @@ def jax_sample_fn_wald(op):
     def sample_fn(rng, size, dtype, *parameters):
         rng_key = rng["jax_state"]
         rng_key, sampling_key = jax.random.split(rng_key, 2)
-
         mean, scale = parameters
 
         key1, key2 = jax.random.split(sampling_key, 2)
@@ -390,6 +389,21 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
+@jax_sample_fn.register(aer.ChiSquareRV)
+def jax_sample_fn_chisquare(op):
+    """JAX implementation of `ChiSquareRV`"""
+
+    def sample_fn(rng, size, dtype, *parameters):
+        rng_key = rng["jax_state"]
+        rng_key, sampling_key = jax.random.split(rng_key, 2)
+        (df,) = parameters
+        sample = jax.random.gamma(sampling_key, df / 2, size, dtype) * 2
+        rng["jax_state"] = rng_key
+        return (rng, sample)
+
+    return sample_fn
+
+
 @jax_sample_fn.register(aer.GeometricRV)
 def jax_sample_fn_geometric(op):
     """JAX implementation of `GeometricRV`."""

tests/link/jax/test_random.py

@@ -97,6 +97,19 @@ def test_random_updates(rng_ctor):
             lambda *args: args,
             None,
         ),
+        (
+            aer.chisquare,
+            [
+                set_test_value(
+                    at.dvector(),
+                    np.array([1.0, 2.0], dtype=np.float64),
+                )
+            ],
+            (2,),
+            "chi2",
+            lambda *args: args,
+            50_000,
+        ),
         (
             aer.exponential,
             [