[callback] Allow external callbacks to return 64-bit values in 32-bit mode #20534
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
gnecula
requested review from
sharadmv and
superbobry
and removed request for
superbobry and
sharadmv
superbobry
reviewed
Apr 3, 2024
| @@ -225,7 +225,7 @@ def _check_shape_dtype(shape_dtype): | |||
| dt = np.dtype(shape_dtype.dtype) | |||
| if dtypes.canonicalize_dtype(dt) != dt: | |||
There was a problem hiding this comment.
Should we check that jax_enable_x64 is disabled as well or does it follow from canonicalize_dtype(dt) != dt?
There was a problem hiding this comment.
We are using the same pattern elsewhere in the code, to check that x64 is not enabler and the type is a 64-bit type.
superbobry
approved these changes
Apr 3, 2024
google-ml-butler
bot
added
kokoro:force-run
pull ready
… 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.
Previously, prior to #20433, if the Python callback returned a Python literal (which is natively a 64-bit value), and the
result_shape_dtypesspecified a 32-bit expected returned value, we would just get garbage results. In #20433, I introduced an error in this situation. However, when trying to port the internal code that uses host_callback toio_callback, I am getting many instances of this error. The common scenario is a Python callback function that returns a Python scalar:However, if the
f_hostwere called directly JAX would canonicalize the value42.to afloat32(whenjax_enable_x64is not set). I do not think that it makes sense forio_callbackto have stricter behavior that a direct call.In this PR we add a canonicalization step on the returned values of Python callbacks, which casts the values to 32-bits if JAX is running in 32-bit mode.
Note that the above example should return an error in 64-bit mode, because the actual returned value is a 64-bit value but the declared expected value is
np.float32. To avoid the error in both 64-bit and 32-bit mode, the python callback should returnnp.float32(42.).In some sense this is replacing the change in #20433 to add a canonicalization step instead of an error.