Add StandardNormalRV JAX implementation

Author: rlouf

aesara/link/jax/dispatch/random.py

@@ -114,6 +114,7 @@ def sample_fn(rng, size, dtype, *parameters):
 @jax_sample_fn.register(aer.LaplaceRV)
 @jax_sample_fn.register(aer.LogisticRV)
 @jax_sample_fn.register(aer.NormalRV)
+@jax_sample_fn.register(aer.StandardNormalRV)
 def jax_sample_fn_loc_scale(op):
     """JAX implementation of random variables in the loc-scale families.
 

tests/link/jax/test_random.py

@@ -221,6 +221,13 @@ def test_random_updates(rng_ctor):
             "randint",
             lambda *args: args,
         ),
+        (
+            aer.standard_normal,
+            [],
+            (2,),
+            "norm",
+            lambda *args: args,
+        ),
         (
             aer.t,
             [