Distrax: Enable Categorical to be used with a single element.

PiperOrigin-RevId: 390374075
google-deepmind · Aug 12, 2021 · 272ccb8 · 272ccb8
1 parent 941bb83
commit 272ccb8
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/distrax/_src/distributions/categorical.py b/distrax/_src/distributions/categorical.py
@@ -50,8 +50,6 @@ def __init__(self,
     """
     super().__init__()
     chex.assert_exactly_one_is_none(probs, logits)
-    chex.if_args_not_none(chex.assert_axis_dimension_gt, probs, axis=-1, val=1)
-    chex.if_args_not_none(chex.assert_axis_dimension_gt, logits, axis=-1, val=1)
     if not (jnp.issubdtype(dtype, jnp.integer) or
             jnp.issubdtype(dtype, jnp.floating)):
       raise ValueError(