Fix total_count dtype in binomial-derived distributions

Author: fehiepsi

numpyro/distributions/conjugate.py

@@ -56,7 +56,7 @@ def __init__(
         self,
         concentration1: ArrayLike,
         concentration0: ArrayLike,
-        total_count: int = 1,
+        total_count: ArrayLike = 1,
         *,
         validate_args: Optional[bool] = None,
     ):
@@ -263,7 +263,7 @@ class DirichletMultinomial(Distribution):
     def __init__(
         self,
         concentration: ArrayLike,
-        total_count: int = 1,
+        total_count: ArrayLike = 1,
         *,
         total_count_max: Optional[int] = None,
         validate_args: Optional[bool] = None,
@@ -431,7 +431,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
 
 
 def NegativeBinomial(
-    total_count: int,
+    total_count: ArrayLike,
     probs: Optional[ArrayLike] = None,
     logits: Optional[ArrayLike] = None,
     *,
@@ -473,7 +473,7 @@ class NegativeBinomialProbs(GammaPoisson):
 
     def __init__(
         self,
-        total_count: int,
+        total_count: ArrayLike,
         probs: ArrayLike,
         *,
         validate_args: Optional[bool] = None,
@@ -502,7 +502,7 @@ class NegativeBinomialLogits(GammaPoisson):
 
     def __init__(
         self,
-        total_count: int,
+        total_count: ArrayLike,
         logits: ArrayLike,
         *,
         validate_args: Optional[bool] = None,

numpyro/distributions/discrete.py

@@ -188,7 +188,7 @@ class BinomialProbs(Distribution):
     def __init__(
         self,
         probs: ArrayLike,
-        total_count: int = 1,
+        total_count: ArrayLike = 1,
         *,
         validate_args: Optional[bool] = None,
     ):
@@ -266,7 +266,7 @@ class BinomialLogits(Distribution):
     def __init__(
         self,
         logits: ArrayLike,
-        total_count: int = 1,
+        total_count: ArrayLike = 1,
         *,
         validate_args: Optional[bool] = None,
     ):
@@ -317,7 +317,7 @@ def support(self) -> constraints.Constraint:
 
 
 def Binomial(
-    total_count: int = 1,
+    total_count: ArrayLike = 1,
     probs: Optional[ArrayLike] = None,
     logits: Optional[ArrayLike] = None,
     *,
@@ -575,7 +575,7 @@ class MultinomialProbs(Distribution):
     def __init__(
         self,
         probs: Array,
-        total_count: int = 1,
+        total_count: ArrayLike = 1,
         *,
         total_count_max: Optional[int] = None,
         validate_args: Optional[bool] = None,
@@ -629,7 +629,7 @@ def support(self) -> constraints.Constraint:
 
     @staticmethod
     def infer_shapes(
-        probs: Array, total_count: int
+        probs: Array, total_count: ArrayLike
     ) -> tuple[tuple[int, ...], tuple[int, ...]]:
         batch_shape = lax.broadcast_shapes(probs[:-1], total_count)
         event_shape = probs[-1:]
@@ -647,7 +647,7 @@ class MultinomialLogits(Distribution):
     def __init__(
         self,
         logits: Array,
-        total_count: int = 1,
+        total_count: ArrayLike = 1,
         *,
         total_count_max: Optional[int] = None,
         validate_args: Optional[bool] = None,
@@ -707,7 +707,7 @@ def support(self) -> constraints.Constraint:
 
     @staticmethod
     def infer_shapes(
-        logits: Array, total_count: int
+        logits: Array, total_count: ArrayLike
     ) -> tuple[tuple[int, ...], tuple[int, ...]]:
         batch_shape = lax.broadcast_shapes(logits[:-1], total_count)
         event_shape = logits[-1:]