JAX: Implement geometric RV sampling

aesara/link/jax/dispatch/random.py

@@ -389,3 +389,20 @@ def sample_fn(rng, size, dtype, *parameters):
  return (rng, samples)
 
  return sample_fn
+
+
+@jax_sample_fn.register(aer.GeometricRV)
+def jax_sample_fn_geometric(op):
+ """JAX implementation of `GeometricRV`."""
+
+ def sample_fn(rng, size, dtype, *parameters):
+ rng_key = rng["jax_state"]
+ rng_key, sampling_key = jax.random.split(rng_key, 2)
+ p = parameters[0]
+ sample_num = jax.numpy.log(jax.random.uniform(sampling_key, size))
+ sample = sample_num / jax.numpy.log1p(-p)
+ sample_ceil = jax.numpy.ceil(sample)
+ rng["jax_state"] = rng_key
+ return (rng, sample_ceil)
+
+ return sample_fn

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))