From 311c9010a81f7396d8244a1883a2b6bdfb442187 Mon Sep 17 00:00:00 2001
From: Smit-create <smitlunagariya.mat18@itbhu.ac.in>
Date: Mon, 20 Feb 2023 17:57:37 +0530
Subject: [PATCH] Add tests for geometric JAX samples

---
 tests/link/jax/test_random.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tests/link/jax/test_random.py b/tests/link/jax/test_random.py
index 4d2afe615b..230b66c65d 100644
--- 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))