https://github.com/google/jax
Revision 75db481299cf774cefbd271539fbccb91a93d5d4 authored by George Necula on 26 March 2024, 12:32:00 UTC, committed by jax authors on 26 March 2024, 12:32:41 UTC
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 19e6156ccec0df7a900471df7840bc421da2898b

PiperOrigin-RevId: 619156176
1 parent 33cf53c
History
Tip revision: 75db481299cf774cefbd271539fbccb91a93d5d4 authored by George Necula on 26 March 2024, 12:32:00 UTC
[callback] Fix io_callback for callbacks that return Python literals.
Tip revision: 75db481
File Mode Size
.github
benchmarks
build
cloud_tpu_colabs
docs
examples
images
jax
jax_plugins
jaxlib
tests
third_party
.bazelrc -rw-r--r-- 18.0 KB
.bazelversion -rw-r--r-- 6 bytes
.editorconfig -rw-r--r-- 292 bytes
.gitignore -rw-r--r-- 355 bytes
.pre-commit-config.yaml -rw-r--r-- 1.1 KB
.readthedocs.yml -rw-r--r-- 569 bytes
AUTHORS -rw-r--r-- 313 bytes
CHANGELOG.md -rw-r--r-- 118.1 KB
CITATION.bib -rw-r--r-- 408 bytes
CONTRIBUTING.md -rw-r--r-- 150 bytes
LICENSE -rw-r--r-- 11.1 KB
README.md -rw-r--r-- 19.8 KB
WORKSPACE -rw-r--r-- 568 bytes
conftest.py -rw-r--r-- 2.5 KB
pyproject.toml -rw-r--r-- 5.7 KB
setup.py -rw-r--r-- 7.0 KB

README.md

back to top