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

numpyro/_typing.py

@@ -4,14 +4,15 @@
 
 from collections import OrderedDict
 from collections.abc import Callable
-from typing import Any, Optional, Protocol, Union, runtime_checkable
+from typing import Any, Optional, Protocol, TypeVar, Union, runtime_checkable
 import weakref
 
 try:
     from typing import ParamSpec, TypeAlias
 except ImportError:
     from typing_extensions import ParamSpec, TypeAlias
 
+
 import numpy as np
 
 import jax
@@ -36,6 +37,9 @@
 """A generic type for a pytree, i.e. a nested structure of lists, tuples, dicts, and arrays."""
 
 
+NumLikeT = TypeVar("NumLikeT", bound=NumLike)
+
+
 @runtime_checkable
 class ConstraintT(Protocol):
     """A protocol for typing constraints."""
@@ -45,10 +49,10 @@ def is_discrete(self) -> bool: ...
     @property
     def event_dim(self) -> int: ...
 
-    def __call__(self, x: ArrayLike) -> ArrayLike: ...
+    def __call__(self, x: NumLike) -> ArrayLike: ...
     def __repr__(self) -> str: ...
-    def check(self, value: ArrayLike) -> ArrayLike: ...
-    def feasible_like(self, prototype: ArrayLike) -> ArrayLike: ...
+    def check(self, value: NumLike) -> ArrayLike: ...
+    def feasible_like(self, prototype: NumLike) -> NumLike: ...
 
 
 @runtime_checkable