Revision 634f58c7d5bae8aad68a881280385f58ca2d0239 authored by Peter Hawkins on 03 May 2022, 01:56:50 UTC, committed by jax authors on 03 May 2022, 01:57:27 UTC
In particular, pjit/xmap work on CPU these days.

PiperOrigin-RevId: 446085110
1 parent b7293d5
History
File Mode Size
jax_cpp
__init__.py -rw-r--r-- 575 bytes
advi.py -rw-r--r-- 4.8 KB
control.py -rw-r--r-- 5.7 KB
control_test.py -rw-r--r-- 7.4 KB
datasets.py -rw-r--r-- 3.1 KB
differentially_private_sgd.py -rw-r--r-- 8.8 KB
examples_test.py -rw-r--r-- 3.5 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
resnet50.py -rw-r--r-- 4.3 KB
spmd_mnist_classifier_fromscratch.py -rw-r--r-- 4.9 KB

back to top