Skip to main content

Browse the archive

https://github.com/google/jax
08 August 2024, 06:23:17 UTC
Revision 6f7be3cf047465a0811583cfb9860ae27e4352aa authored by Peter Hawkins on 04 March 2024, 22:00:36 UTC, committed by jax authors on 04 March 2024, 22:01:31 UTC 
Define lax.Precision directly in Python, rather than inheriting from a C++ type in jaxlib.
 
Historically, we defined Precision to be an enum exported from jaxlib using pybind11, since that was the type the old XLA ComputationBuilder classes expected as input. But we build IR using StableHLO MLIR builders these days, and there's no reason for the JAX-level Precision type to match the XLA-internal one.

In a future change I plan to change the definition of Precision in jaxlib to be defined using nanobind instead of pybind11. Nanobind defines its enum classes to be final by default, which precludes this inheritance, and that's probably a good design decision by nanobind. But as discussed above, there's no good reason to inherit in the first place.

PiperOrigin-RevId: 612575404
1 parent 67b0eb3
History
Tip revision: 6f7be3cf047465a0811583cfb9860ae27e4352aa authored by Peter Hawkins on 04 March 2024, 22:00:36 UTC
Define lax.Precision directly in Python, rather than inheriting from a C++ type in jaxlib.
Tip revision: 6f7be3c
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