[callback] Fix io_callback for callbacks that return Python literals. #20433
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
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_619100631
branch
4 times, most recently
from
e07671d to
5187651
Compare
The internal implementation of io_callback and friends currently use .shape and .dtype on the result of the callback. This fails if the callback returns a Python literal. Fixed the checks that the callback returns values of expected shape and dtype, and added tests. Reverts 19e6156 PiperOrigin-RevId: 619156176
copybara-service
bot
force-pushed
the
test_619100631
branch
from
5187651 to
75db481
Compare
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Apr 2, 2024
… mode Previously, prior to jax-ml#20433, if the Python callback returned a Python literal (which is natively a 64-bit value), and the `result_shape_dtypes` specified a 32-bit expected returned value, we would just get garbage results. In jax-ml#20433, I introduced an error in this situation. However, when trying to port the internal code that uses host_callback to `io_callback`, I am getting many instances of this error. The common scenario is a Python callback function that returns a Python scalar: ``` def f_host(): return 42. io_callback(f_host, jax.ShapeDtypeStruct((), np.float32)) ``` However, if the `f_host` were called directly JAX would canonicalize the value `42.` to a float32 (when `jax_enable_x64` is not set). I do not think that it makes sense for `io_callback` to have stricter behaviour that a direct call. In this PR we add a canonicalization step on the returned values of Python callbacks, which would cast the values to 32-bits. In some sense this is replacing the change in jax-ml#20433 to add a canonicalization step instead of an error.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Apr 2, 2024
… mode Previously, prior to jax-ml#20433, if the Python callback returned a Python literal (which is natively a 64-bit value), and the `result_shape_dtypes` specified a 32-bit expected returned value, we would just get garbage results. In jax-ml#20433, I introduced an error in this situation. However, when trying to port the internal code that uses host_callback to `io_callback`, I am getting many instances of this error. The common scenario is a Python callback function that returns a Python scalar: ``` def f_host(): return 42. io_callback(f_host, jax.ShapeDtypeStruct((), np.float32)) ``` However, if the `f_host` were called directly JAX would canonicalize the value `42.` to a float32 (when `jax_enable_x64` is not set). I do not think that it makes sense for `io_callback` to have stricter behaviour that a direct call. In this PR we add a canonicalization step on the returned values of Python callbacks, which would cast the values to 32-bits. In some sense this is replacing the change in jax-ml#20433 to add a canonicalization step instead of an error.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Apr 3, 2024
… mode Previously, prior to jax-ml#20433, if the Python callback returned a Python literal (which is natively a 64-bit value), and the `result_shape_dtypes` specified a 32-bit expected returned value, we would just get garbage results. In jax-ml#20433, I introduced an error in this situation. However, when trying to port the internal code that uses host_callback to `io_callback`, I am getting many instances of this error. The common scenario is a Python callback function that returns a Python scalar: ``` def f_host(): return 42. io_callback(f_host, jax.ShapeDtypeStruct((), np.float32)) ``` However, if the `f_host` were called directly JAX would canonicalize the value `42.` to a float32 (when `jax_enable_x64` is not set). I do not think that it makes sense for `io_callback` to have stricter behaviour that a direct call. In this PR we add a canonicalization step on the returned values of Python callbacks, which would cast the values to 32-bits. In some sense this is replacing the change in jax-ml#20433 to add a canonicalization step instead of an error.
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.
[callback] Fix io_callback for callbacks that return Python literals.
The internal implementation of io_callback and friends currently use .shape and .dtype on the result of the callback. This fails if the callback returns a Python literal.
Fixed the checks that the callback returns values of expected shape and dtype,
and added tests.
Reverts 19e6156