Skip to content

Fix jaxtyping>=0.3.0 support for NumPy annotations#730

Merged
yukinarit merged 1 commit intoyukinarit:mainfrom
pablovela5620:codex/jaxtyping-0-3-support
Mar 18, 2026
Merged

Fix jaxtyping>=0.3.0 support for NumPy annotations#730
yukinarit merged 1 commit intoyukinarit:mainfrom
pablovela5620:codex/jaxtyping-0-3-support

Conversation

@pablovela5620
Copy link
Copy Markdown
Contributor

Fixes #643.

Summary

I used Codex to help investigate the regression, isolate the root cause, and validate the fix.

pyserde's jaxtyping support broke with jaxtyping>=0.3.0 because jaxtyping annotations no longer subclass their backing array type.

Before 0.3.0, an annotation like jaxtyping.Float[np.ndarray, "2 2"] was a subclass of np.ndarray, so pyserde detected it with issubclass(..., np.ndarray).

From 0.3.0 onward, jaxtyping annotations inherit from jaxtyping.AbstractArray instead, while still exposing array_type. That caused pyserde's jaxtyping NumPy handling to stop matching and fall through to SerdeError: Unsupported type: ndarray.

This patch keeps the fix small:

  • switch jaxtyping NumPy detection to jaxtyping.AbstractArray plus array_type is np.ndarray
  • remove the jaxtyping<0.3.0 pin
  • add a focused regression test for the detection helper
  • keep the existing jaxtyping round-trip tests as the main behavior check
  • update the standalone jaxtyping example to use concrete NumPy dtypes accepted by current NumPy

Local validation

Ran locally:

make test
make examples
uv run pytest tests/test_numpy.py -k jaxtyping -v
uv run python examples/type_numpy_jaxtyping.py
uv run ruff check serde/numpy.py tests/test_numpy.py examples/type_numpy_jaxtyping.py

@codecov
Copy link
Copy Markdown

codecov Bot commented Mar 13, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.81%. Comparing base (f09d503) to head (9b0a9c4).
⚠️ Report is 65 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #730      +/-   ##
==========================================
- Coverage   90.83%   90.81%   -0.03%     
==========================================
  Files          13       13              
  Lines        2040     2307     +267     
  Branches      370      449      +79     
==========================================
+ Hits         1853     2095     +242     
- Misses        124      141      +17     
- Partials       63       71       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Copy Markdown
Owner

@yukinarit yukinarit left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, Thanks!

@yukinarit yukinarit added the build Build related change label Mar 18, 2026
@yukinarit yukinarit merged commit 9cd6338 into yukinarit:main Mar 18, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

build Build related change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

jaxtyping support is broken with jaxtyping>=0.3.0

2 participants