Skip to content

fix: jax v0.7.2+ LiteralArray / TypedNdArray #74

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
patrick-kidger merged 15 commits into from
Oct 12, 2025
Merged

fix: jax v0.7.2+ LiteralArray / TypedNdArray #74

patrick-kidger merged 15 commits into from
Oct 12, 2025

Conversation

@nstarman
Copy link
Contributor

@nstarman nstarman commented Oct 1, 2025 •
edited
Loading

Replaces #73

@nstarman
Copy link
Contributor Author

nstarman commented Oct 1, 2025 •
edited
Loading

@patrick-kidger, the jax + jaxtyping PRs re LiteralArray fixed many of the problems, but we still have issues like:

plum.resolver.NotFoundLookupError: `convert_element_type_dispatcher(LiteralArray(0, dtype=int32))` could not be resolved.

Besides registering specific rules, is there a better way to automatically support these primitives?

@nstarman nstarman changed the title deps: bump jaxtyping v0.3.3+ fix: jax v0.7.2+ LiteralArray / TypedNdArray Oct 1, 2025
@nstarman nstarman force-pushed the deps-bump-jaxtyping branch from 825ac87 to d56084a Compare October 1, 2025 18:03
@patrick-kidger
Copy link
Owner

patrick-kidger commented Oct 3, 2025

Hmm, this surprises me. If LiteralArray is included in the ArrayLike union then I would have imagined that plum would respect that when a rule is registered with an ArrayLike annotation, and thus be able to dispatch on it.

I'd suggest tracking down why that expected flow doesn't hold.

nstarman added 11 commits October 3, 2025 12:38
@nstarman
types: annotate core
f0b751b 
Signed-off-by: nstarman <[email protected]>
@nstarman
style: remove extraneous import
aff0973 
Signed-off-by: nstarman <[email protected]>
@nstarman
refactor: consolidate value default check
3ab4b05 
Signed-off-by: nstarman <[email protected]>
@nstarman
refactor: new rule definition
726b583 
Signed-off-by: nstarman <[email protected]>
@nstarman
deps: add ipykernel to dev
b892033 
Signed-off-by: nstarman <[email protected]>
@nstarman
refactor: use _src.core import
09ab2b5 
Signed-off-by: nstarman <[email protected]>
@nstarman
types: annotate full_lower
c4dd13e 
Signed-off-by: nstarman <[email protected]>
@nstarman
types: annotate _is_value
7d1376e 
Signed-off-by: nstarman <[email protected]>
@nstarman
types: annotate tracer
a5163f9 
Signed-off-by: nstarman <[email protected]>
@nstarman
style: improve func name for error tracing
8699aeb 
Signed-off-by: nstarman <[email protected]>
@nstarman
types: annotate __get__
b9c4cc8 
Signed-off-by: nstarman <[email protected]>
@nstarman nstarman force-pushed the deps-bump-jaxtyping branch from d56084a to b9c4cc8 Compare October 3, 2025 17:10
@nstarman nstarman mentioned this pull request Oct 3, 2025
Merged
@nstarman
Copy link
Contributor Author

nstarman commented Oct 3, 2025

I'd suggest tracking down why that expected flow doesn't hold.

@patrick-kidger I think it's because of patrick-kidger/equinox#1112

@nstarman
Copy link
Contributor Author

nstarman commented Oct 3, 2025

I tested this PR with patrick-kidger/equinox#1112 and all tests passed, so hopefully that's all we need.

@patrick-kidger patrick-kidger reopened this Oct 3, 2025
nstarman
nstarman commented Oct 3, 2025
View reviewed changes
@nstarman nstarman mentioned this pull request Oct 3, 2025
Closed
@nstarman
deps: bump equinox
d32f368 
Signed-off-by: nstarman <[email protected]>
@nstarman
deps: protect jax
c509a6a 
Signed-off-by: nstarman <[email protected]>
@nstarman
Copy link
Contributor Author

nstarman commented Oct 10, 2025

Some small code cleaning was introduced in the process of debugging. :)

@nstarman nstarman marked this pull request as ready for review October 10, 2025 09:02
@nstarman
Copy link
Contributor Author

nstarman commented Oct 10, 2025

@patrick-kidger it works!

@nstarman nstarman mentioned this pull request Oct 11, 2025
Merged
patrick-kidger
patrick-kidger reviewed Oct 12, 2025
View reviewed changes
quax/_core.py
def existing_rule():
assert False
def new_rule():
raise NotImplementedError("Abstract primitive") # pragma: no cover
Copy link
Owner

@patrick-kidger patrick-kidger Oct 12, 2025

Choose a reason for hiding this comment

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

pragma?

@patrick-kidger patrick-kidger merged commit fef6a5b into patrick-kidger:main Oct 12, 2025
5 checks passed
@patrick-kidger
Copy link
Owner

patrick-kidger commented Oct 12, 2025

LGTM, merged!

@nstarman nstarman deleted the deps-bump-jaxtyping branch October 12, 2025 18:09
patrick-kidger
patrick-kidger reviewed Oct 12, 2025
View reviewed changes
pyproject.toml
"beartype>=0.20.2",
"pytest>=8.3.5",
"pytest-env>=1.1.5",
"ipykernel>=6.30.1",
Copy link
Owner

@patrick-kidger patrick-kidger Oct 12, 2025

Choose a reason for hiding this comment

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

Actually, why did this sneak in as well?

Copy link
Contributor Author

@nstarman nstarman Oct 12, 2025

Choose a reason for hiding this comment

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

I guess dev on Jupyter.

@patrick-kidger patrick-kidger mentioned this pull request Oct 12, 2025
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@patrick-kidger patrick-kidger patrick-kidger left review comments

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@nstarman @patrick-kidger