Skip to content

ENH: Test tools for jax.jit and dask - #115

Merged
lucascolley merged 1 commit into
data-apis:mainfrom
crusaderky:test_jit
Jan 24, 2025
Merged

ENH: Test tools for jax.jit and dask#115
lucascolley merged 1 commit into
data-apis:mainfrom
crusaderky:test_jit

Conversation

@crusaderky

@crusaderkycrusaderky commented Jan 21, 2025•
edited
Loading

Copy link
Copy Markdown
Contributor

Porting of scipy/scipy#22308, with additional support for Dask.

@lucascolleylucascolley added the enhancement New feature or request label Jan 21, 2025
@lucascolleylucascolley changed the title ENH Test tools for jax.jit and daskENH: Test tools for jax.jit and daskJan 21, 2025
crusaderky
crusaderky commented Jan 21, 2025
View reviewed changes
Comment threadtests/test_funcs.py
crusaderky
crusaderky commented Jan 21, 2025
View reviewed changes
Comment threadtests/test_funcs.py
lazy_xp_function(kron, static_argnames="xp")
lazy_xp_function(nunique, static_argnames="xp")
lazy_xp_function(pad, static_argnames=("pad_width", "mode", "constant_values", "xp"))
# FIXME calls in1d which calls xp.unique_values without size

@crusaderkycrusaderkyJan 21, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

This is non-trivial to fix. Will look into it in a follow-up.

@lucascolley
lucascolley self-requested a review January 21, 2025 10:42
@crusaderkycrusaderky mentioned this pull request Jan 21, 2025
Open
@crusaderky
crusaderkyforce-pushed the test_jit branch 5 times, most recently from 57e6cd0 to cd84a8bCompareJanuary 21, 2025 18:06
@crusaderkycrusaderky mentioned this pull request Jan 22, 2025
Merged
@crusaderky
crusaderkyforce-pushed the test_jit branch 2 times, most recently from 64ba88b to 47ef8b5CompareJanuary 22, 2025 16:08
crusaderky
crusaderky commented Jan 22, 2025
View reviewed changes
Comment threadsrc/array_api_extra/testing.py
This was referenced Jan 22, 2025
Closed
Merged
Closed
lucascolley
lucascolley approved these changes Jan 23, 2025
View reviewed changes

@lucascolleylucascolley left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

thanks @crusaderky, LGTM! I don't have much to say on the details of how this interfaces with JAX and Dask, but it all sounds reasonable to me, and the code looks maintainable.

Comment threaddocs/api-reference.md Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
Comment threadsrc/array_api_extra/testing.py Outdated
@lucascolleylucascolley added this to the 0.6.1 milestone Jan 24, 2025
@crusaderky

crusaderky commented Jan 24, 2025

Copy link
Copy Markdown
ContributorAuthor

@lucascolley all comments addressed

@lucascolley
lucascolley self-requested a review January 24, 2025 16:31
@lucascolley
lucascolley merged commit 8fa3fd2 into data-apis:mainJan 24, 2025
@crusaderky
crusaderky deleted the test_jit branch January 24, 2025 17:58
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@lucascolleylucascolleylucascolley approved these changes

Assignees

No one assigned

Labels

enhancementNew feature or requestlazy arraysxpx.testing

Projects

None yet

Milestone

0.6.1

Development

Successfully merging this pull request may close these issues.

2 participants

@crusaderky@lucascolley