Add StudentTRV JAX implementation

Author: rlouf

pytensor/link/jax/dispatch/random.py

@@ -208,6 +208,24 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
+@jax_sample_fn.register(aer.StudentTRV)
+def jax_sample_fn_t(op):
+    """JAX implementation of `StudentTRV`."""
+
+    def sample_fn(rng, size, dtype, *parameters):
+        rng_key = rng["jax_state"]
+        (
+            df,
+            loc,
+            scale,
+        ) = parameters
+        sample = loc + jax.random.t(rng_key, df, size, dtype) * scale
+        rng["jax_state"] = jax.random.split(rng_key, num=1)[0]
+        return (rng, sample)
+
+    return sample_fn
+
+
 @jax_sample_fn.register(aer.ChoiceRV)
 def jax_funcify_choice(op):
     """JAX implementation of `ChoiceRV`."""

tests/link/jax/test_random.py

@@ -205,6 +205,26 @@ def test_random_updates(rng_ctor):
             "randint",
             lambda *args: args,
         ),
+        (
+            aer.t,
+            [
+                set_test_value(
+                    at.dscalar(),
+                    np.array(2.0, dtype=np.float64),
+                ),
+                set_test_value(
+                    at.dvector(),
+                    np.array([1.0, 2.0], dtype=np.float64),
+                ),
+                set_test_value(
+                    at.dscalar(),
+                    np.array(1.0, dtype=np.float64),
+                ),
+            ],
+            (2,),
+            "t",
+            lambda *args: args,
+        ),
         (
             aer.uniform,
             [