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

Qazalbash · 2025-08-24T23:06:57Z

This PR contains the resolution of mypy errors passed by #2032, in numpyro.distributions.transforms module.

There are two cases in particular which I am unable to resolve. You can see them by running the mypy.

log_abs_det_jacobian of many transforms have unused parameters. I have typed them as union of UnusedParam and some appropriate numpy/jax type.

Many cases were unresolvable, like, __eq__ method expects bool as return type, but & operation between arrays return array of type bool, which conflicts with the return type, therefore I have added the tag to ignore them. You will find similar tags in the file.

…forms

… and constraints

juanitorduz · 2025-08-29T17:58:35Z

This looks great, thanks! There are some minor errors

numpyro/distributions/transforms.py:82: error: Missing positional argument "x" in call to "__call__" of "TransformT"  [call-arg]
Installing missing stub packages:
numpyro/distributions/transforms.py:84: error: Name "inv" already defined on line 80  [no-redef]
numpyro/distributions/transforms.py:85: error: Incompatible types in assignment (expression has type "ReferenceType[None]", variable has type "TransformT | None")  [assignment]
numpyro/distributions/transforms.py:86: error: Incompatible return value type (got "Array | Any | None", expected "TransformT")  [return-value]
numpyro/distributions/transforms.py:1550: error: Item "ndarray[tuple[Any, ...], dtype[Any]]" of "ndarray[tuple[Any, ...], dtype[Any]] | Array" has no attribute "at"  [union-attr]

Do you need some help with these :) ?

Qazalbash · 2025-08-29T18:52:02Z

numpyro/distributions/transforms.py:82: error: Missing positional argument "x" in call to "__call__" of "TransformT"  [call-arg]
Installing missing stub packages:
numpyro/distributions/transforms.py:84: error: Name "inv" already defined on line 80  [no-redef]
numpyro/distributions/transforms.py:85: error: Incompatible types in assignment (expression has type "ReferenceType[None]", variable has type "TransformT | None")  [assignment]
numpyro/distributions/transforms.py:86: error: Incompatible return value type (got "Array | Any | None", expected "TransformT")  [return-value]

These errors are from numpyro/distributions/transforms.py#L77-L86, and I am not able to understand the significance of different conditions and the weak reference. I think you can look into this matter.

I will take this one,

numpyro/distributions/transforms.py:1550: error: Item "ndarray[tuple[Any, ...], dtype[Any]]" of "ndarray[tuple[Any, ...], dtype[Any]] | Array" has no attribute "at"  [union-attr]

Thank you for offering help ❤️.

juanitorduz · 2025-08-29T19:20:47Z

ok! Sounds like a plan! I will try to look at it in the next days :)

juanitorduz · 2025-08-30T17:45:34Z

Hey @Qazalbash I gave it a try as in d57d9a6 . MyPy is happy now, maybe you can try it ? The only key change was self.inv() -> self.inv which I think makes more sense, let's see if the tests complain ;)

Qazalbash · 2025-08-30T18:19:05Z

@juanitorduz Thanks for the changes, mypy is happy now.

Do I need to remove the plugin from pyproject.toml?

juanitorduz · 2025-08-30T18:20:51Z

It's depreciated so it's safe to remove

…ng of inverse transforms Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

numpyro/distributions/transforms.py

Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

juanitorduz · 2025-09-04T14:14:29Z

ok! I think the tests are failing because a new NNX release and changes in nnx.merge, see https://github.com/google/flax/releases/tag/v0.11.2

The other tests FAILED test/test_distributions.py::test_entropy_samples I am not sure about.

juanitorduz · 2025-09-04T17:28:24Z

Here is a patch for the first errors #2067

Qazalbash · 2025-09-05T04:30:13Z

Here's another patch #2069 😸

fehiepsi · 2025-09-05T19:26:31Z

numpyro/_typing.py

+        x: Union[jax.Array, Any],
+        y: Union[jax.Array, Any],
+        intermediates: Optional[Any] = None,
+    ) -> Union[jax.Array, Any]: ...


why we need those Union, shouldn't Any be enough?

If we only use Any, then IDEs show types like this,

otherwise,

I see. You want to see some Array there. How about using ArrayLike?

Gotcha, thanks! How about using NumLike like in my comment above?

numpyro/distributions/transforms.py

fehiepsi · 2025-09-05T19:31:01Z

numpyro/_typing.py

@@ -20,6 +22,9 @@
 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]


I'm not sure why StrictArrayT is needed, could you elaborate? I think we can use ArrayLike instead

Some functions do not operate on boolean or complex; ArrayLike contains boolean and complex types, and mypy was throwing an error over it. That's why I created StrictArrayT, it is used where only arrays are required.

But we also want to support int, float at those places. Could you let me know which function has the issue?

How about using ArrayLike in most places and in specific cases, use NumLike? StrictArray can be used at those places which expect non-scalar shapes.

StrictArray = Union[np.ndarray, jax.Array] NumLike = Union[ jax.Array, np.ndarray, np.number, bool, int, float, complex, ]

(we don't need to have T trailing in those cases I think)

…lasses and add PyTree type alias

fehiepsi

Hi @Qazalbash, I think we can get around issues by using NumLike just at some specific places. It's fine to use StrictArray at those arrays with dim >= 1. NonScalarArray is a good name for it I guess.

fehiepsi · 2025-09-06T12:24:10Z

numpyro/_typing.py

@@ -20,6 +22,9 @@
 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]


How about using ArrayLike in most places and in specific cases, use NumLike? StrictArray can be used at those places which expect non-scalar shapes.

StrictArray = Union[np.ndarray, jax.Array] NumLike = Union[ jax.Array, np.ndarray, np.number, bool, int, float, complex, ]

(we don't need to have T trailing in those cases I think)

fehiepsi · 2025-09-06T12:24:31Z

numpyro/_typing.py

-    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]: ...


prefer: ArrayLike and not use Any in general

fehiepsi · 2025-09-06T12:25:24Z

numpyro/_typing.py

+        x: Union[jax.Array, Any],
+        y: Union[jax.Array, Any],
+        intermediates: Optional[Any] = None,
+    ) -> Union[jax.Array, Any]: ...


Gotcha, thanks! How about using NumLike like in my comment above?

fehiepsi · 2025-09-06T12:26:33Z

numpyro/distributions/transforms.py

            inv = self._inv()
        if inv is None:
-            inv = _InverseTransform(self)
-            self._inv = weakref.ref(inv)
+            inv = cast(TransformT, _InverseTransform(self))


this seems unnecessary to me

fehiepsi · 2025-09-06T12:26:36Z

numpyro/distributions/transforms.py

-            inv = _InverseTransform(self)
-            self._inv = weakref.ref(inv)
+            inv = cast(TransformT, _InverseTransform(self))
+            self._inv = cast(TransformT, weakref.ref(inv))


why do we need this cast, it seems incorrect to me?

fehiepsi · 2025-09-06T12:30:20Z

numpyro/distributions/transforms.py

+        x: StrictArrayT,
+        y: StrictArrayT,
+        intermediates: Optional[PyTree] = None,
+    ) -> StrictArrayT:


I believe you can use ArrayLike at those places as long as the Protocol uses the type hint: log_abs_det_jacobian(...) -> NumLike:

fehiepsi · 2025-09-06T12:31:16Z

numpyro/distributions/transforms.py

@@ -191,7 +198,7 @@ def inverse_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
    def tree_flatten(self):
        return (self._inv,), (("_inv",), dict())

-    def __eq__(self, other: TransformT) -> bool:
+    def __eq__(self, other: object) -> bool:


just curious, why we allow to compare with non TransformT?

fehiepsi · 2025-09-06T12:32:48Z

numpyro/distributions/transforms.py

@@ -268,10 +275,13 @@ def __call__(self, x: ArrayLike) -> ArrayLike:
        return self.loc + self.scale * x

    def _inverse(self, y: ArrayLike) -> ArrayLike:
-        return (y - self.loc) / self.scale
+        return (y - self.loc) / self.scale  # type: ignore[call-overload,operator]


as long as loc and scale are NumLike, this should be fine I think.

fehiepsi · 2025-09-06T12:34:41Z

numpyro/distributions/transforms.py

@@ -372,7 +385,7 @@ def log_abs_det_jacobian(
                    )
                )

-        result = 0.0
+        result = jnp.zeros(())


if output is ArrayLike, I guess we can use 0.0 here.

fehiepsi · 2025-09-06T12:35:01Z

numpyro/distributions/transforms.py

-    def call_with_intermediates(
-        self, x: ArrayLike
-    ) -> Tuple[ArrayLike, Optional[ArrayLike]]:
+    def call_with_intermediates(self, x: ArrayLike) -> Tuple[ArrayLike, Sequence]:


Sequence -> PyTree

Qazalbash added 2 commits August 13, 2025 01:16

chore: incomplete changes for type hint in numpyro.distribution.trans…

c2ca7c3

…forms

refactor: update type hints and improve type safety across transforms…

3b01836

… and constraints

juanitorduz referenced this pull request Aug 30, 2025

fix errors

d57d9a6

Qazalbash and others added 2 commits August 30, 2025 23:22

fix: improve type casting in Transform class and ensure proper handli…

35e6d1b

…ng of inverse transforms Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

chore: remove numpy.typing.mypy_plugin from mypy configuration

23407d3

Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

juanitorduz reviewed Aug 30, 2025

View reviewed changes

numpyro/distributions/transforms.py Outdated Show resolved Hide resolved

fix: update _inv type to support weak references in Transform class

300fead

Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

Qazalbash requested a review from fehiepsi September 4, 2025 12:04

fehiepsi added the awaiting review label Sep 5, 2025

fehiepsi reviewed Sep 5, 2025

View reviewed changes

Qazalbash added 2 commits September 6, 2025 12:00

Merge branch 'master' into type-hint-transforms

f166e84

fix: update intermediates type to Optional[PyTree] in Transform c…

8aa2592

…lasses and add PyTree type alias

fehiepsi reviewed Sep 6, 2025

View reviewed changes

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

Are you sure you want to change the base?

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

Uh oh!

Conversation

Qazalbash commented Aug 24, 2025

Uh oh!

juanitorduz commented Aug 29, 2025

Uh oh!

Qazalbash commented Aug 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

juanitorduz commented Aug 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

juanitorduz commented Aug 30, 2025

Uh oh!

Qazalbash commented Aug 30, 2025

Uh oh!

juanitorduz commented Aug 30, 2025

Uh oh!

Uh oh!

juanitorduz commented Sep 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

juanitorduz commented Sep 4, 2025

Uh oh!

Qazalbash commented Sep 5, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Qazalbash Sep 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fehiepsi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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

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

Qazalbash commented Aug 29, 2025 •

edited

Loading

juanitorduz commented Aug 29, 2025 •

edited

Loading

juanitorduz commented Sep 4, 2025 •

edited

Loading

Qazalbash Sep 6, 2025 •

edited

Loading