omnistaging on by default #4038
Merged
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
mattjj
force-pushed
the
omnistaging-on-by-default
branch
4 times, most recently
from
e7b74ca to
241267c
Compare
|
|
mattjj
force-pushed
the
omnistaging-on-by-default
branch
9 times, most recently
from
1bce9c3 to
12916cf
Compare
gnecula
reviewed
Aug 15, 2020
mattjj
force-pushed
the
omnistaging-on-by-default
branch
2 times, most recently
from
cb7461c to
3d7bbb0
Compare
This was referenced Sep 8, 2020
mattjj
force-pushed
the
omnistaging-on-by-default
branch
4 times, most recently
from
4ed7807 to
c5f2528
Compare
mattjj
force-pushed
the
omnistaging-on-by-default
branch
from
c5f2528 to
9cb19d6
Compare
mattjj
added a commit
that referenced
this pull request
Sep 18, 2020
Previously, given this function:
```python
@jax.jit
def f(x,y):
if x > y:
return x
else:
return y
```
we'd get an error message like this (after #4038, improved to help with
omnistaging debugging):
```
...
While tracing the function f at tim.py:3, this value became a tracer due to JAX operations on these lines:
operation c:bool[] = gt a:int32[] b:int32[]
from line tim.py:5 (f)
...
```
But this message is buggy! In this case, the value is a tracer because
it has a data dependence on arguments to a jitted function.
After this change, we instead produce this error message:
```
...
While tracing the function f at tim.py:3, this concrete value was not available in Python because it depends on the value of the arguments to f at tim.py:3 at positions [0, 1], and the computation of these values is being staged out.
...
```
I'm eager to iterate with further improvements, but for now I want to
fix this buggy message.
This was referenced Oct 11, 2020
hawkinsp
added a commit
to hawkinsp/jax
that referenced
this pull request
Mar 4, 2021
This is removing the device constant part of jax-ml#1668. We can do this because after jax-ml#3370 and jax-ml#4038 omnistaging removes the need for lazy device constants in a jitted context. (They could still in principle be useful in an op-by-op context, but the power:weight isn't worthwhile anymore.) After this change, the only parts of the lazy sublanguage that remain are those to do with broadcasts and transposes. We may or may not kill those in a follow-up (it hinges on whether any benefit to op-by-op execution is worth the extra complexity). This change regresses non-omnistaging users. As one particular example, test_eval_shape_big_random_array no longer passes with omnistaging disabled.
hawkinsp
added a commit
to hawkinsp/jax
that referenced
this pull request
Mar 4, 2021
Updated version of jax-ml#4536. This is removing the device constant part of jax-ml#1668. We can do this because after jax-ml#3370 and jax-ml#4038 omnistaging removes the need for lazy device constants in a jitted context. (They could still in principle be useful in an op-by-op context, but the power:weight isn't worthwhile anymore.) After this change, the only parts of the lazy sublanguage that remain are those to do with broadcasts and transposes. We may or may not kill those in a follow-up (it hinges on whether any benefit to op-by-op execution is worth the extra complexity). This change regresses non-omnistaging users. As one particular example, test_eval_shape_big_random_array no longer passes with omnistaging disabled.
hawkinsp
added a commit
to hawkinsp/jax
that referenced
this pull request
Mar 4, 2021
Updated version of jax-ml#4536. This is removing the device constant part of jax-ml#1668. We can do this because after jax-ml#3370 and jax-ml#4038 omnistaging removes the need for lazy device constants in a jitted context. (They could still in principle be useful in an op-by-op context, but the power:weight isn't worthwhile anymore.) After this change, the only parts of the lazy sublanguage that remain are those to do with broadcasts and transposes. We may or may not kill those in a follow-up (it hinges on whether any benefit to op-by-op execution is worth the extra complexity). This change regresses non-omnistaging users. As one particular example, test_eval_shape_big_random_array no longer passes with omnistaging disabled.
NeilGirdhar
pushed a commit
to NeilGirdhar/jax
that referenced
this pull request
Mar 17, 2021
Updated version of jax-ml#4536. This is removing the device constant part of jax-ml#1668. We can do this because after jax-ml#3370 and jax-ml#4038 omnistaging removes the need for lazy device constants in a jitted context. (They could still in principle be useful in an op-by-op context, but the power:weight isn't worthwhile anymore.) After this change, the only parts of the lazy sublanguage that remain are those to do with broadcasts and transposes. We may or may not kill those in a follow-up (it hinges on whether any benefit to op-by-op execution is worth the extra complexity). This change regresses non-omnistaging users. As one particular example, test_eval_shape_big_random_array no longer passes with omnistaging disabled.
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.
fixes #4114, fixes #3397, fixes #3108,