Repository: pyro-ppl/numpyro

Issue:
Feature request: DiscreteUniform distribution
[As suggested in the forum](https://forum.pyro.ai/t/impute-missing-counts-statistical-rethinking-15h7/3589), I'm opening this up as a feature request

If welcome, I'd be interested in working on it


Discussion:

Resolving patch:
diff --git a/examples/annotation.py b/examples/annotation.py
--- a/examples/annotation.py
+++ b/examples/annotation.py
@@ -167,9 +167,7 @@ def mace(positions, annotations):
         theta = numpyro.sample("theta", dist.Beta(0.5, 0.5))
 
     with numpyro.plate("item", num_items, dim=-2):
-        # NB: using constant logits for discrete uniform prior
-        # (NumPyro does not have DiscreteUniform distribution yet)
-        c = numpyro.sample("c", dist.Categorical(logits=jnp.zeros(num_classes)))
+        c = numpyro.sample("c", dist.DiscreteUniform(0, num_classes - 1))
 
         with numpyro.plate("position", num_positions):
             s = numpyro.sample("s", dist.Bernoulli(1 - theta[positions]))
diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
--- a/numpyro/distributions/__init__.py
+++ b/numpyro/distributions/__init__.py
@@ -54,6 +54,7 @@
     Categorical,
     CategoricalLogits,
     CategoricalProbs,
+    DiscreteUniform,
     Geometric,
     GeometricLogits,
     GeometricProbs,
@@ -114,6 +115,7 @@
     "Delta",
     "Dirichlet",
     "DirichletMultinomial",
+    "DiscreteUniform",
     "Distribution",
     "Exponential",
     "ExpandedDistribution",
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -405,6 +405,65 @@ def Categorical(probs=None, logits=None, validate_args=None):
         raise ValueError("One of `probs` or `logits` must be specified.")
 
 
+class DiscreteUniform(Distribution):
+    arg_constraints = {"low": constraints.dependent, "high": constraints.dependent}
+    has_enumerate_support = True
+
+    def __init__(self, low=0, high=1, validate_args=None):
+        self.low, self.high = promote_shapes(low, high)
+        batch_shape = lax.broadcast_shapes(jnp.shape(low), jnp.shape(high))
+        self._support = constraints.integer_interval(low, high)
+        super().__init__(batch_shape, validate_args=validate_args)
+
+    @constraints.dependent_property(is_discrete=True, event_dim=0)
+    def support(self):
+        return self._support
+
+    def sample(self, key, sample_shape=()):
+        shape = sample_shape + self.batch_shape
+        return random.randint(key, shape=shape, minval=self.low, maxval=self.high + 1)
+
+    @validate_sample
+    def log_prob(self, value):
+        shape = lax.broadcast_shapes(jnp.shape(value), self.batch_shape)
+        return -jnp.broadcast_to(jnp.log(self.high + 1 - self.low), shape)
+
+    def cdf(self, value):
+        cdf = (jnp.floor(value) + 1 - self.low) / (self.high - self.low + 1)
+        return jnp.clip(cdf, a_min=0.0, a_max=1.0)
+
+    def icdf(self, value):
+        return self.low + value * (self.high - self.low + 1) - 1
+
+    @property
+    def mean(self):
+        return self.low + (self.high - self.low) / 2.0
+
+    @property
+    def variance(self):
+        return ((self.high - self.low + 1) ** 2 - 1) / 12.0
+
+    def enumerate_support(self, expand=True):
+        if not not_jax_tracer(self.high) or not not_jax_tracer(self.low):
+            raise NotImplementedError("Both `low` and `high` must not be a JAX Tracer.")
+        if np.any(np.amax(self.low) != self.low):
+            # NB: the error can't be raised if inhomogeneous issue happens when tracing
+            raise NotImplementedError(
+                "Inhomogeneous `low` not supported by `enumerate_support`."
+            )
+        if np.any(np.amax(self.high) != self.high):
+            # NB: the error can't be raised if inhomogeneous issue happens when tracing
+            raise NotImplementedError(
+                "Inhomogeneous `high` not supported by `enumerate_support`."
+            )
+        values = (self.low + jnp.arange(np.amax(self.high - self.low) + 1)).reshape(
+            (-1,) + (1,) * len(self.batch_shape)
+        )
+        if expand:
+            values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape)
+        return values
+
+
 class OrderedLogistic(CategoricalProbs):
     """
     A categorical distribution with ordered outcomes.


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -614,6 +614,8 @@ def get_sp_dist(jax_dist):
     T(dist.NegativeBinomial2, jnp.array([10.2, 7, 31]), jnp.array([10.2, 20.7, 2.1])),
     T(dist.OrderedLogistic, -2, jnp.array([-10.0, 4.0, 9.0])),
     T(dist.OrderedLogistic, jnp.array([-4, 3, 4, 5]), jnp.array([-1.5])),
+    T(dist.DiscreteUniform, -2, jnp.array([-1.0, 4.0, 9.0])),
+    T(dist.DiscreteUniform, jnp.array([-4, 3, 4, 5]), jnp.array([6])),
     T(dist.Poisson, 2.0),
     T(dist.Poisson, jnp.array([2.0, 3.0, 5.0])),
     T(SparsePoisson, 2.0),