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
_src
example_libraries
experimental
image
interpreters
lax
lib
nn
numpy
ops
scipy
tools
__init__.py -rw-r--r-- 4.9 KB
abstract_arrays.py -rw-r--r-- 752 bytes
ad_checkpoint.py -rw-r--r-- 705 bytes
api_util.py -rw-r--r-- 763 bytes
cloud_tpu_init.py -rw-r--r-- 627 bytes
config.py -rw-r--r-- 676 bytes
core.py -rw-r--r-- 86.7 KB
custom_batching.py -rw-r--r-- 651 bytes
custom_derivatives.py -rw-r--r-- 1.0 KB
custom_transpose.py -rw-r--r-- 638 bytes
distributed.py -rw-r--r-- 620 bytes
dlpack.py -rw-r--r-- 647 bytes
dtypes.py -rw-r--r-- 1016 bytes
errors.py -rw-r--r-- 966 bytes
flatten_util.py -rw-r--r-- 623 bytes
jaxpr_util.py -rw-r--r-- 6.4 KB
linear_util.py -rw-r--r-- 11.0 KB
prng.py -rw-r--r-- 862 bytes
profiler.py -rw-r--r-- 1.0 KB
py.typed -rw-r--r-- 0 bytes
random.py -rw-r--r-- 5.6 KB
stages.py -rw-r--r-- 676 bytes
test_util.py -rw-r--r-- 2.1 KB
tree_util.py -rw-r--r-- 2.5 KB
util.py -rw-r--r-- 956 bytes
version.py -rw-r--r-- 864 bytes

back to top