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
File | Mode | Size |
---|---|---|
jax_cpp | ||
__init__.py | -rw-r--r-- | 581 bytes |
advi.py | -rw-r--r-- | 4.8 KB |
datasets.py | -rw-r--r-- | 3.1 KB |
differentially_private_sgd.py | -rw-r--r-- | 8.9 KB |
examples_test.py | -rw-r--r-- | 2.0 KB |
gaussian_process_regression.py | -rw-r--r-- | 4.5 KB |
kernel_lsq.py | -rw-r--r-- | 2.5 KB |
mnist_classifier.py | -rw-r--r-- | 3.1 KB |
mnist_classifier_fromscratch.py | -rw-r--r-- | 3.0 KB |
mnist_vae.py | -rw-r--r-- | 4.9 KB |
onnx2xla.py | -rw-r--r-- | 4.7 KB |
spmd_mnist_classifier_fromscratch.py | -rw-r--r-- | 4.9 KB |
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...