Revision 473d1c3685570e505bed58512afa05fb7d7a8935 authored by Skye Wanderman-Milne on 28 March 2023, 20:42:51 UTC, committed by Skye Wanderman-Milne on 28 March 2023, 22:28:13 UTC
I forgot that the default setting is actually in jaxlib:
https://github.com/openxla/xla/blob/fbe9a80fdb8c429e8a175962459da348cd560a50/xla/python/xla_client.py#L135

To be able to make this change as a jax-only release, I manually set
the env var on Cloud TPU if it isn't already set.
1 parent 4061bbb
Raw File
.gitignore
*.pyc
*.so
*.egg-info
*.whl
/build/lib
/build/bazel*
/dist/
.ipynb_checkpoints
/bazel-*
.jax_configure.bazelrc
/tensorflow
.DS_Store
.mypy_cache/
.pytype/
/docs/build
*_pb2.py
/docs/notebooks/.ipynb_checkpoints/
/docs/_autosummary
.idea
.vscode
.envrc
jax.iml
.bazelrc.user

# virtualenv/venv directories
/venv/
/bin/
/include/
/lib/
/share/
back to top