https://github.com/google/jax
Revision 5283d4b4a5298c3822039f35454adca4e80ea84b authored by Sergei Lebedev on 04 March 2024, 13:41:29 UTC, committed by jax authors on 04 March 2024, 13:42:03 UTC
This allows propagating the names bottom up -- from equations to the jaxpr,
instead of "discovering" them top-down by traversing (and rebuilding) the
jaxpr via core.subst_axis_names.

PiperOrigin-RevId: 612416803
1 parent 2dd5e9e
History
Tip revision: 5283d4b4a5298c3822039f35454adca4e80ea84b authored by Sergei Lebedev on 04 March 2024, 13:41:29 UTC
Axis names are now tracked via an effect
Tip revision: 5283d4b
File Mode Size
.github
benchmarks
build
cloud_tpu_colabs
docs
examples
images
jax
jax_plugins
jaxlib
tests
third_party
.bazelrc -rw-r--r-- 17.3 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-- 116.8 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-- 724 bytes
conftest.py -rw-r--r-- 2.5 KB
pyproject.toml -rw-r--r-- 5.5 KB
setup.py -rw-r--r-- 7.8 KB

README.md

back to top