Fix Dirichlet.arg_constraints event_dim

pytorch · Jan 29, 2021 · 6d15dd3 · 6d15dd3
1 parent dbfaf96
commit 6d15dd3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/test/distributions/test_distributions.py b/test/distributions/test_distributions.py
@@ -3926,6 +3926,11 @@ def test_params_constraints(self):
                     except KeyError:
                         continue  # ignore optional parameters
 
+                    # Check param shape is compatible with distribution shape.
+                    self.assertGreaterEqual(value.dim(), constraint.event_dim)
+                    value_batch_shape = value.shape[:value.dim() - constraint.event_dim]
+                    torch.broadcast_shapes(dist.batch_shape, value_batch_shape)
+
                     if is_dependent(constraint):
                         continue
 

diff --git a/torch/distributions/dirichlet.py b/torch/distributions/dirichlet.py
@@ -40,7 +40,7 @@ class Dirichlet(ExponentialFamily):
         concentration (Tensor): concentration parameter of the distribution
             (often referred to as alpha)
     """
-    arg_constraints = {'concentration': constraints.positive}
+    arg_constraints = {'concentration': constraints.independent(constraints.positive, 1)}
     support = constraints.simplex
     has_rsample = True