fix(gh-2036): MyPy Errors in numpyro.distributions.transforms Module

numpyro/_typing.py

@@ -4,13 +4,15 @@
 
 from collections import OrderedDict
 from collections.abc import Callable
-from typing import Any, Protocol, runtime_checkable
+from typing import Any, Optional, Protocol, Union, runtime_checkable
 
 try:
     from typing import ParamSpec, TypeAlias
 except ImportError:
     from typing_extensions import ParamSpec, TypeAlias
 
+import numpy as np
+
 import jax
 from jax.typing import ArrayLike
 
@@ -20,6 +22,13 @@
 Message: TypeAlias = dict[str, Any]
 TraceT: TypeAlias = OrderedDict[str, Message]
 
+# ArrayLike type has StaticScalar, StrictArrayT has everything except StaticScalars
+StrictArrayT = Union[np.ndarray, jax.Array]
+
+
+PyTree: TypeAlias = Any
+"""A generic type for a pytree, i.e. a nested structure of lists, tuples, dicts, and arrays."""
+
 
 @runtime_checkable
 class ConstraintT(Protocol):
@@ -87,20 +96,25 @@ def is_discrete(self) -> bool: ...
 
 @runtime_checkable
 class TransformT(Protocol):
-    domain = ConstraintT
-    codomain = ConstraintT
-    _inv: "TransformT" = None
+    domain: ConstraintT = ...
+    codomain: ConstraintT = ...
+    _inv: Optional["TransformT"] = ...
 
-    def __call__(self, x: ArrayLike) -> ArrayLike: ...
-    def _inverse(self, y: ArrayLike) -> ArrayLike: ...
+    def __call__(self, x: Union[jax.Array, Any]) -> Union[jax.Array, Any]: ...
+    def _inverse(self, y: Union[jax.Array, Any]) -> Union[jax.Array, Any]: ...
     def log_abs_det_jacobian(
-        self, x: ArrayLike, y: ArrayLike, intermediates=None
-    ) -> ArrayLike: ...
-    def call_with_intermediates(self, x: ArrayLike) -> tuple[ArrayLike, None]: ...
+        self,
+        x: Union[jax.Array, Any],
+        y: Union[jax.Array, Any],
+        intermediates: Optional[PyTree] = None,
+    ) -> Union[jax.Array, Any]: ...
+    def call_with_intermediates(
+        self, x: Union[jax.Array, Optional[Any]]
+    ) -> tuple[Union[jax.Array, Any], Any]: ...
     def forward_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: ...
     def inverse_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: ...
 
     @property
     def inv(self) -> "TransformT": ...
     @property
-    def sign(self) -> ArrayLike: ...
+    def sign(self) -> Union[ArrayLike, Any]: ...

numpyro/distributions/constraints.py

@@ -801,18 +801,18 @@ def tree_flatten(self):
 corr_cholesky: ConstraintT = _CorrCholesky()
 corr_matrix: ConstraintT = _CorrMatrix()
 dependent: ConstraintT = _Dependent()
-greater_than: ConstraintT = _GreaterThan
-greater_than_eq: ConstraintT = _GreaterThanEq
-less_than: ConstraintT = _LessThan
-less_than_eq: ConstraintT = _LessThanEq
-independent: ConstraintT = _IndependentConstraint
-integer_interval: ConstraintT = _IntegerInterval
-integer_greater_than: ConstraintT = _IntegerGreaterThan
-interval: ConstraintT = _Interval
+greater_than = _GreaterThan
+greater_than_eq = _GreaterThanEq
+less_than = _LessThan
+less_than_eq = _LessThanEq
+independent = _IndependentConstraint
+integer_interval = _IntegerInterval
+integer_greater_than = _IntegerGreaterThan
+interval = _Interval
 l1_ball: ConstraintT = _L1Ball()
 lower_cholesky: ConstraintT = _LowerCholesky()
 scaled_unit_lower_cholesky: ConstraintT = _ScaledUnitLowerCholesky()
-multinomial: ConstraintT = _Multinomial
+multinomial = _Multinomial
 nonnegative: ConstraintT = _Nonnegative()
 nonnegative_integer: ConstraintT = _IntegerNonnegative()
 ordered_vector: ConstraintT = _OrderedVector()
@@ -830,5 +830,5 @@ def tree_flatten(self):
 softplus_positive: ConstraintT = _SoftplusPositive()
 sphere: ConstraintT = _Sphere()
 unit_interval: ConstraintT = _UnitInterval()
-open_interval: ConstraintT = _OpenInterval
-zero_sum: ConstraintT = _ZeroSum
+open_interval = _OpenInterval
+zero_sum = _ZeroSum