fix event dim of some arg constraints (#1967)

Makefile

@@ -10,8 +10,8 @@ license: FORCE
 	python scripts/update_headers.py
 
 format: license FORCE
-	ruff check --fix .
 	ruff format .
+	ruff check --fix .
 
 install: FORCE
 	pip install -e '.[dev,doc,test,examples]'

numpyro/distributions/continuous.py

@@ -2426,7 +2426,10 @@ def entropy(self):
 
 
 class Uniform(Distribution):
-    arg_constraints = {"low": constraints.dependent, "high": constraints.dependent}
+    arg_constraints = {
+        "low": constraints.dependent(is_discrete=False, event_dim=0),
+        "high": constraints.dependent(is_discrete=False, event_dim=0),
+    }
     reparametrized_params = ["low", "high"]
     pytree_data_fields = ("low", "high", "_support")
 
@@ -2727,7 +2730,7 @@ class Wishart(TransformedDistribution):
     """
 
     arg_constraints = {
-        "concentration": constraints.dependent(is_discrete=False),
+        "concentration": constraints.dependent(is_discrete=False, event_dim=0),
         "scale_matrix": constraints.positive_definite,
         "rate_matrix": constraints.positive_definite,
         "scale_tril": constraints.lower_cholesky,
@@ -2820,7 +2823,7 @@ class WishartCholesky(Distribution):
     """
 
     arg_constraints = {
-        "concentration": constraints.dependent(is_discrete=False),
+        "concentration": constraints.dependent(is_discrete=False, event_dim=0),
         "scale_matrix": constraints.positive_definite,
         "rate_matrix": constraints.positive_definite,
         "scale_tril": constraints.lower_cholesky,

numpyro/distributions/discrete.py

@@ -421,7 +421,10 @@ def Categorical(probs=None, logits=None, *, validate_args=None):
 
 
 class DiscreteUniform(Distribution):
-    arg_constraints = {"low": constraints.dependent, "high": constraints.dependent}
+    arg_constraints = {
+        "low": constraints.dependent(is_discrete=True, event_dim=0),
+        "high": constraints.dependent(is_discrete=True, event_dim=0),
+    }
     has_enumerate_support = True
     pytree_data_fields = ("low", "high", "_support")
 

numpyro/distributions/truncated.py

@@ -179,8 +179,8 @@ def var(self):
 
 class TwoSidedTruncatedDistribution(Distribution):
     arg_constraints = {
-        "low": constraints.dependent,
-        "high": constraints.dependent,
+        "low": constraints.dependent(is_discrete=False, event_dim=0),
+        "high": constraints.dependent(is_discrete=False, event_dim=0),
     }
     reparametrized_params = ["low", "high"]
     supported_types = (Cauchy, Laplace, Logistic, Normal, SoftLaplace, StudentT)

test/test_distributions.py

@@ -1327,6 +1327,17 @@ def test_has_rsample(jax_dist, sp_dist, params):
                 transf_dist.rsample(random.PRNGKey(0))
 
 
+@pytest.mark.parametrize(
+    "jax_dist_cls, sp_dist, params", CONTINUOUS + DISCRETE + DIRECTIONAL
+)
+def test_args_attributes(jax_dist_cls, sp_dist, params):
+    jax_dist = jax_dist_cls(*params)
+    for constraint in jax_dist.arg_constraints.values():
+        if jax_dist_cls != dist.Delta:
+            constraint.event_dim
+        constraint.is_discrete
+
+
 @pytest.mark.parametrize("batch_shape", [(), (4,), (3, 2)])
 def test_unit(batch_shape):
     log_factor = random.normal(random.PRNGKey(0), batch_shape)