chore: incomplete changes for type hint in numpyro.distribution.transforms

Author: Qazalbash

numpyro/_typing.py

@@ -7,9 +7,11 @@
 from typing import Any, Protocol, runtime_checkable
 
 try:
-    from typing import ParamSpec, TypeAlias
+    from typing import ParamSpec, TypeAlias, Union
 except ImportError:
-    from typing_extensions import ParamSpec, TypeAlias
+    from typing_extensions import ParamSpec, TypeAlias, Union
+
+import numpy as np
 
 import jax
 from jax.typing import ArrayLike
@@ -20,6 +22,9 @@
 Message: TypeAlias = dict[str, Any]
 TraceT: TypeAlias = OrderedDict[str, Message]
 
+# ArrayLike type has StaticScalar, StrictArrayT has everything except StaticScalars
+StrictArrayT: TypeAlias = Union[np.ndarray, jax.Array]
+
 
 @runtime_checkable
 class ConstraintT(Protocol):
@@ -87,20 +92,20 @@ def is_discrete(self) -> bool: ...
 
 @runtime_checkable
 class TransformT(Protocol):
-    domain = ConstraintT
-    codomain = ConstraintT
+    domain: ConstraintT = ...
+    codomain: ConstraintT = ...
     _inv: "TransformT" = None
 
-    def __call__(self, x: ArrayLike) -> ArrayLike: ...
-    def _inverse(self, y: ArrayLike) -> ArrayLike: ...
+    def __call__(self, x: jax.Array) -> jax.Array: ...
+    def _inverse(self, y: jax.Array) -> jax.Array: ...
     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: jax.Array, y: jax.Array, intermediates=None
+    ) -> jax.Array: ...
+    def call_with_intermediates(self, x: jax.Array) -> tuple[jax.Array, None]: ...
     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) -> jax.Array: ...

numpyro/distributions/constraints.py

@@ -805,7 +805,7 @@ def tree_flatten(self):
 greater_than_eq: ConstraintT = _GreaterThanEq
 less_than: ConstraintT = _LessThan
 less_than_eq: ConstraintT = _LessThanEq
-independent: ConstraintT = _IndependentConstraint
+independent = _IndependentConstraint
 integer_interval: ConstraintT = _IntegerInterval
 integer_greater_than: ConstraintT = _IntegerGreaterThan
 interval: ConstraintT = _Interval