https://github.com/google/jax
History
Tip revision: a9a72fa53b4eac1212c5d33388fb140bba105047 authored by jax authors on 10 June 2024, 20:39:35 UTC
Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.
Tip revision: a9a72fa
File Mode Size
_src
example_libraries
experimental
extend
image
interpreters
lax
lib
nn
numpy
ops
scipy
tools
BUILD -rw-r--r-- 24.7 KB
__init__.py -rw-r--r-- 10.3 KB
ad_checkpoint.py -rw-r--r-- 802 bytes
api_util.py -rw-r--r-- 769 bytes
cloud_tpu_init.py -rw-r--r-- 633 bytes
collect_profile.py -rw-r--r-- 4.6 KB
core.py -rw-r--r-- 8.0 KB
custom_batching.py -rw-r--r-- 691 bytes
custom_derivatives.py -rw-r--r-- 1.3 KB
custom_transpose.py -rw-r--r-- 664 bytes
debug.py -rw-r--r-- 1.1 KB
distributed.py -rw-r--r-- 673 bytes
dlpack.py -rw-r--r-- 710 bytes
dtypes.py -rw-r--r-- 1.1 KB
errors.py -rw-r--r-- 1.2 KB
export.py -rw-r--r-- 1.3 KB
flatten_util.py -rw-r--r-- 645 bytes
monitoring.py -rw-r--r-- 1.2 KB
profiler.py -rw-r--r-- 1.1 KB
py.typed -rw-r--r-- 0 bytes
random.py -rw-r--r-- 9.6 KB
sharding.py -rw-r--r-- 1.7 KB
stages.py -rw-r--r-- 1.3 KB
test_util.py -rw-r--r-- 835 bytes
tree.py -rw-r--r-- 956 bytes
tree_util.py -rw-r--r-- 3.1 KB
typing.py -rw-r--r-- 3.3 KB
util.py -rw-r--r-- 1.1 KB
version.py -rw-r--r-- 5.7 KB

back to top