Added argument check to all primitives. #3197
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
* Added argument check to all primitives. The issue that inspired this is that `lax.tie_in` is easy to misuse if the first argument is not a JAX type, then it silently disappears. This means that `lax.tie_in((x, x), const)` is the same as `const` even though `x` is a tracer. This error would be caught previously if core.skip_checks == False because then `bind` checks its arguments. I have essentially added an unconditional argument check to `bind`. In case this is considered too inefficient, we can add argument checking to individual primivites, e.g., tie_in. For most primitives if a non-JAX array is passed, the `impl` rule would fire and `numpy` would report the error somehow, perhaps. * Merged find_top_trace with check_args This was previously merged as jax-ml#2948 but reverted awaiting the fixes in some user code.
mattjj
added a commit
that referenced
this pull request
May 31, 2020
The previous version was written and tested for performance; the revised version caused at least a 25% slowdown in the dispatch time of `lax.add(1, 2)` (and so likely a much bigger slowdown for the find_top_trace timing alone). Instead, we can just change the error message in xla.abstractify, since invalid types lead to abstractification errors when we apply primitive impls.
mattjj
added a commit
that referenced
this pull request
May 31, 2020
The previous version was written and tested for performance; the revised version caused at least a 25% slowdown in the dispatch time of `lax.add(1, 2)` (and so likely a much bigger slowdown for the find_top_trace timing alone). Instead, we can just change the error message in xla.abstractify, since invalid types lead to abstractification errors when we apply primitive impls.
Merged
mattjj
added a commit
that referenced
this pull request
Jun 1, 2020
revert find_top_trace change from #3197 The previous version was written and tested for performance; the revised version caused at least a 25% slowdown in the dispatch time of `lax.add(1, 2)` (and so likely a much bigger slowdown for the find_top_trace timing alone). Instead, we can just change the error message in xla.abstractify, since invalid types lead to abstractification errors when we apply primitive impls.
NeilGirdhar
pushed a commit
to NeilGirdhar/jax
that referenced
this pull request
Jun 11, 2020
* Added argument check to all primitives. The issue that inspired this is that `lax.tie_in` is easy to misuse if the first argument is not a JAX type, then it silently disappears. This means that `lax.tie_in((x, x), const)` is the same as `const` even though `x` is a tracer. This error would be caught previously if core.skip_checks == False because then `bind` checks its arguments. I have essentially added an unconditional argument check to `bind`. In case this is considered too inefficient, we can add argument checking to individual primivites, e.g., tie_in. For most primitives if a non-JAX array is passed, the `impl` rule would fire and `numpy` would report the error somehow, perhaps. * Merged find_top_trace with check_args This was previously merged as jax-ml#2948 but reverted awaiting the fixes in some user code.
NeilGirdhar
pushed a commit
to NeilGirdhar/jax
that referenced
this pull request
Jun 11, 2020
revert find_top_trace change from jax-ml#3197 The previous version was written and tested for performance; the revised version caused at least a 25% slowdown in the dispatch time of `lax.add(1, 2)` (and so likely a much bigger slowdown for the find_top_trace timing alone). Instead, we can just change the error message in xla.abstractify, since invalid types lead to abstractification errors when we apply primitive impls.
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.
The issue that inspired this is that
lax.tie_iniseasy to misuse if the first argument is not a JAX type, then
it silently disappears. This means that
lax.tie_in((x, x), const)is the same as
consteven thoughxis a tracer.This error would be caught previously if core.skip_checks == False
because then
bindchecks its arguments. I have essentially addedan unconditional argument check to
bind.In case this is considered too inefficient, we can add argument
checking to individual primivites, e.g., tie_in. For most primitives
if a non-JAX array is passed, the
implrule would fire andnumpywould report the error somehow, perhaps.
This was previously merged as #2948 but reverted awaiting the fixes
in some user code.