ENH: Test tools for jax.jit and dask #115
Open
+579
−17
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
lucascolley
changed the title
crusaderky
commented
Jan 21, 2025
|
|
||
| # some xp backends are untyped | ||
| # mypy: disable-error-code=no-untyped-usage | ||
|
|
||
| lazy_xp_function(atleast_nd, static_argnames=("ndim", "xp")) | ||
| lazy_xp_function(cov, static_argnames="xp") | ||
| # FIXME .device attribute https://github.com/data-apis/array-api-compat/pull/238 |
There was a problem hiding this comment.
data-apis/array-api-compat#238
Evidence that the numpy-compat PR works: crusaderky@d27f75d
crusaderky
commented
Jan 21, 2025
| 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 |
There was a problem hiding this comment.
This is non-trivial to fix. Will look into it in a follow-up.
crusaderky
force-pushed
the
test_jit
branch
4 times, most recently
from
a00a646 to
57e6cd0
Compare
crusaderky
commented
Jan 22, 2025
Comment on lines
+70
to
+87
| allow_dask_compute : int, optional | ||
| Number of times `func` is allowed to internally materialize the Dask graph. This | ||
| is typically triggered by `bool()`, `float()`, or `np.asarray()`. | ||
|
|
||
| Set to 1 if you are aware that `func` converts the input parameters to numpy and | ||
| want to let it do so at least for the time being, knowing that it is going to be | ||
| extremely detrimental for performance. | ||
|
|
||
| If a test needs values higher than 1 to pass, it is a canary that the conversion | ||
| to numpy/bool/float is happening multiple times, which translates to multiple | ||
| computations of the whole graph. Short of making the function fully lazy, you | ||
| should at least add explicit calls to `np.asarray()` early in the function. | ||
| *Note:* the counter of `allow_dask_compute` resets after each call to `func`, so | ||
| a test function that invokes `func` multiple times should still work with this | ||
| parameter set to 1. | ||
|
|
||
| Default: 0, meaning that `func` must be fully lazy and never materialize the | ||
| graph. |
There was a problem hiding this comment.
In a previous version this was a all-or-nothing test. scipy/scipy#22240 (comment) prompted me to change it.
This was referenced Jan 22, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Porting of scipy/scipy#22308, with additional support for Dask.