types: is_array(_like) provides TypeGuard by nstarman · Pull Request #1148 · patrick-kidger/equinox

nstarman · 2025-11-19T21:37:54Z

This PR:

enhances the return type annotations for is_array and is_arraylike so that they do better type narrowing.
updates pyright to get rid of a CI warning (ruff also got updated as part of uv run pre-commit autoupdate)
simplifies a few functions I saw while addressing pyright errors

nstarman · 2025-11-30T20:50:46Z

tests/test_jit.py

    array_tree = [{"a": a, "b": b}, (c,)]
-    mlp_add = jtu.tree_map(lambda u: u + 1 if eqx.is_array(u) else u, general_tree[-1])
+    mlp_add = jtu.tree_map(
+        lambda u: jnp.asarray(u) + 1 if eqx.is_array(u) else u, general_tree[-1]


is_array checks for np.generic, but this doen't support u+1. Do you want to restrict from np.generic to np.number | np.bool | ... ?

That sounds reasonable to me. I don't think jnp.array(<some np.generic that is not a number or bool>) is going to work in general anyway.

nstarman · 2025-11-30T21:31:07Z

@patrick-kidger this improvement to the return type is uncovering a real issue: np.generic isn't always a full stand-in for np.ndarray. Also, and this isn't an issue we can solve, jax has limited static support for __jax_array__.

nstarman · 2025-11-30T21:44:23Z

numpy/numpy#30335 will eventually allow for some simplifications!

equinox/_filters.py

patrick-kidger · 2025-12-02T20:12:21Z

equinox/_filters.py

+    overload,
+    runtime_checkable,
+    TypeAlias,
+    TypeGuard,


Perhaps we should check for 3.13+ and use TypeIs in that case?

I don't think that's what we want in this case. TypeIs doesn't enable type narrowing within the union.

equinox/_filters.py

patrick-kidger · 2025-12-02T20:14:36Z

equinox/_filters.py

+    _ARRAY_TYPES += (_TypedNdArray,)

+# Type alias for type checkers. Never disappears from unions: T | Never == T
+TypedNdArray: TypeAlias = _TypedNdArray  # pyright: ignore[reportInvalidTypeForm]


Do we need both TypedNdArray and _TypedNdArray? We could just put _TypedNdArray directly into the union.

equinox/_filters.py

equinox/nn/_embedding.py

patrick-kidger · 2025-12-02T20:17:58Z

tests/test_jit.py

    array_tree = [{"a": a, "b": b}, (c,)]
-    mlp_add = jtu.tree_map(lambda u: u + 1 if eqx.is_array(u) else u, general_tree[-1])
+    mlp_add = jtu.tree_map(
+        lambda u: jnp.asarray(u) + 1 if eqx.is_array(u) else u, general_tree[-1]


That sounds reasonable to me. I don't think jnp.array(<some np.generic that is not a number or bool>) is going to work in general anyway.

Copilot

Pull request overview

This PR enhances type narrowing capabilities by adding TypeGuard return type annotations to the is_array and is_array_like functions, enabling better static type checking when using these predicates.

Key changes:

Added TypeGuard return types to is_array and is_array_like functions for improved type narrowing
Refined array type definitions to accept np.number and np.bool_ instead of the broader np.generic to exclude non-array-like types (e.g., np.object_, np.flexible)
Updated tooling versions (pyright v1.1.406→v1.1.407, ruff v0.13.0→v0.14.7)

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
`equinox/_filters.py`	Core changes: added TypeGuard annotations, refined NDArrayType definition, introduced ArrayTypes and ArrayLikeTypes type aliases with TYPE_CHECKING branches, and added HasJaxArray protocol
`equinox/nn/_mlp.py`	Added `jnp.asarray()` call to handle NumPy scalar types (`np.number`/`np.bool_`) that TypeGuard now includes in array types
`equinox/_doc_utils.py`	Simplified `doc_repr` function logic by flattening nested conditionals
`.pre-commit-config.yaml`	Updated pre-commit hook versions for ruff and pyright

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

equinox/_filters.py

equinox/nn/_mlp.py

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

nstarman · 2025-12-03T23:27:01Z

equinox/_filters.py

+# array-like.
+NDArrayType: TypeAlias = np.ndarray | np.number | np.bool_
+_NDARRAY_TYPES: Final = get_args(NDArrayType)
+_ARRAY_TYPES = _NDARRAY_TYPES + (jax.Array,)


I think the change to type narrow NDArrayType from generic to number | bool means it excludes bfloat16!
🤔
We want to include that, while still excluding np.object_ and np.flexible.

Ah, interesting! I think just adding jax.dtypes.bfloat16 should work?

It would also probably be fine to include np.generic as an overestimation, to be robust to some of the more esoteric 4-bit (etc) dtypes as well?

That looks to be necessary since jax doesn't provide a nice union of the ml_dtypes that it imports.
I wish they would! Or just subclass np.number

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

nstarman · 2025-12-09T15:27:02Z

Yeesh. Did not expect this small bit of type narrowing to be so difficult. The state of array scalars (instances of np.generic) is a big mess.

patrick-kidger · 2025-12-14T16:47:37Z

equinox/_filters.py


+# Type alias for type checkers. Never disappears from unions: T | Never == T
+if TYPE_CHECKING:
+    ArrayTypes: TypeAlias = NDArrayType | jax.Array | LiteralArray | TypedNdArray  # type: ignore[valid-type]


The type: ignore here is the kind of thing that'll also introduce maintenance overhead, from the oddities between what type checkers accept / how JAX handles TypedNdArray / etc all intersect, and how none of these things are stable in how they behave going into the future.

Let's just drop TypedNdArray from this union, since it's so rare?

patrick-kidger · 2025-12-14T16:47:42Z

equinox/_filters.py

+# e.g. `np.object_` and `np.flexible`. But `ml_dtypes` also defines dtypes
+# that inherit from `np.generic` and can't easily be listed here individually.
+@runtime_checkable
+class ArrayLikeGeneric(Protocol):


I'd rather keep this file simple than overly precise, I think. JAX will inevitably change something, and I don't want the maintenance burden of tracking things too precisely!

Let's just use np.generic and accept that this is an upper bound?

nstarman force-pushed the typeguards branch from 2dd0a54 to 7ea915f Compare November 30, 2025 05:14

nstarman changed the title ~~types: is_array(_like) provides TypeIs~~ types: is_array(_like) provides TypeGuard Nov 30, 2025

nstarman commented Nov 30, 2025

View reviewed changes

nstarman force-pushed the typeguards branch from 911b785 to bf22e0d Compare November 30, 2025 21:23

nstarman marked this pull request as ready for review November 30, 2025 21:28

patrick-kidger reviewed Dec 2, 2025

View reviewed changes

Copilot AI review requested due to automatic review settings December 3, 2025 22:49

nstarman force-pushed the typeguards branch from df83956 to 1a5011e Compare December 3, 2025 22:49

Copilot started reviewing on behalf of nstarman December 3, 2025 22:49 View session

Copilot finished reviewing on behalf of nstarman December 3, 2025 22:51

Copilot AI reviewed Dec 3, 2025

View reviewed changes

equinox/_filters.py Outdated Show resolved Hide resolved

equinox/nn/_mlp.py Outdated Show resolved Hide resolved

nstarman force-pushed the typeguards branch from 1a5011e to 36f3faf Compare December 3, 2025 22:57

nstarman added 3 commits December 3, 2025 18:03

ci: update ruff and pyright versions in pre-commit config

b3138f8

refactor: simplify doc_repr function logic

d448235

fix: ensure jnp.save handles array-like

1ea51fb

nstarman force-pushed the typeguards branch from 36f3faf to a69558b Compare December 3, 2025 23:05

refactor: update type hints and ensure compatibility with JAX arrays

e3b8c17

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

nstarman force-pushed the typeguards branch from a69558b to e3b8c17 Compare December 3, 2025 23:06

nstarman commented Dec 3, 2025

View reviewed changes

nstarman requested a review from patrick-kidger December 9, 2025 02:39

revert NDArrayType narrowing

4c48e6c

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

nstarman force-pushed the typeguards branch from 09463d0 to 4c48e6c Compare December 9, 2025 03:11

patrick-kidger reviewed Dec 14, 2025

View reviewed changes

Uh oh!

Conversation

nstarman commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nstarman Nov 30, 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!

nstarman commented Nov 30, 2025

Uh oh!

nstarman commented Nov 30, 2025

Uh oh!

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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nstarman Dec 9, 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!

nstarman commented Dec 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nstarman commented Nov 19, 2025 •

edited

Loading

nstarman Nov 30, 2025 •

edited

Loading

nstarman Dec 9, 2025 •

edited

Loading