Add tests for geometric JAX samples

aesara-devs · Feb 21, 2023 · 311c901 · 311c901
1 parent 02bb94f
commit 311c901
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/tests/link/jax/test_random.py b/tests/link/jax/test_random.py
@@ -369,6 +369,21 @@ def test_random_bernoulli(size):
     np.testing.assert_allclose(samples.mean(axis=0), 0.5, 1)
 
 
+@pytest.mark.parametrize(
+    "p, size",
+    [
+        (0.6, ()),
+        (0.2, (4,)),
+    ],
+)
+def test_random_geometric(p, size):
+    rng = shared(np.random.RandomState(123))
+    g = at.random.geometric(p, size=(1000,) + size, rng=rng)
+    g_fn = function([], g, mode=jax_mode)
+    samples = g_fn()
+    np.testing.assert_allclose(samples.mean(), 1 / p, atol=0.1)
+
+
 def test_random_mvnormal():
     rng = shared(np.random.RandomState(123))