https://github.com/google/jax
History
Tip revision: cfa94899bf56fc2c763ac8e05630ac960ce39264 authored by Tomás Longeri on 10 July 2024, 23:46:53 UTC
[Mosaic] Add "Not implemented" prefix to error message.
Tip revision: cfa9489
File Mode Size
clusters
cudnn
debugger
export
extend
image
internal_test_util
interpreters
lax
lib
nn
numpy
ops
pallas
scipy
state
third_party
__init__.py -rw-r--r-- 581 bytes
abstract_arrays.py -rw-r--r-- 2.7 KB
ad_checkpoint.py -rw-r--r-- 36.7 KB
ad_util.py -rw-r--r-- 3.7 KB
api.py -rw-r--r-- 126.5 KB
api_util.py -rw-r--r-- 26.9 KB
array.py -rw-r--r-- 46.7 KB
basearray.py -rw-r--r-- 4.5 KB
basearray.pyi -rw-r--r-- 10.3 KB
blocked_sampler.py -rw-r--r-- 5.6 KB
cache_key.py -rw-r--r-- 10.4 KB
callback.py -rw-r--r-- 18.1 KB
checkify.py -rw-r--r-- 51.8 KB
cloud_tpu_init.py -rw-r--r-- 3.2 KB
compilation_cache.py -rw-r--r-- 9.3 KB
compilation_cache_interface.py -rw-r--r-- 863 bytes
compiler.py -rw-r--r-- 24.9 KB
compute_on.py -rw-r--r-- 1.8 KB
config.py -rw-r--r-- 63.7 KB
core.py -rw-r--r-- 119.8 KB
custom_api_util.py -rw-r--r-- 876 bytes
custom_batching.py -rw-r--r-- 10.2 KB
custom_derivatives.py -rw-r--r-- 59.3 KB
custom_partitioning.py -rw-r--r-- 21.9 KB
custom_transpose.py -rw-r--r-- 8.1 KB
debugging.py -rw-r--r-- 24.4 KB
deprecations.py -rw-r--r-- 3.8 KB
dispatch.py -rw-r--r-- 21.4 KB
distributed.py -rw-r--r-- 10.5 KB
dlpack.py -rw-r--r-- 10.9 KB
dtypes.py -rw-r--r-- 29.9 KB
earray.py -rw-r--r-- 4.0 KB
effects.py -rw-r--r-- 4.6 KB
environment_info.py -rw-r--r-- 2.0 KB
errors.py -rw-r--r-- 24.2 KB
flatten_util.py -rw-r--r-- 3.5 KB
hardware_utils.py -rw-r--r-- 1.7 KB
jaxpr_util.py -rw-r--r-- 6.3 KB
lax_reference.py -rw-r--r-- 17.7 KB
layout.py -rw-r--r-- 5.2 KB
lazy_loader.py -rw-r--r-- 1.3 KB
linear_util.py -rw-r--r-- 13.5 KB
logging_config.py -rw-r--r-- 2.0 KB
lru_cache.py -rw-r--r-- 6.4 KB
maps.py -rw-r--r-- 84.7 KB
mesh.py -rw-r--r-- 11.1 KB
monitoring.py -rw-r--r-- 3.8 KB
op_shardings.py -rw-r--r-- 3.8 KB
partition_spec.py -rw-r--r-- 1.6 KB
path.py -rw-r--r-- 1.1 KB
pickle_util.py -rw-r--r-- 2.3 KB
pjit.py -rw-r--r-- 115.0 KB
pretty_printer.py -rw-r--r-- 15.1 KB
prng.py -rw-r--r-- 44.4 KB
profiler.py -rw-r--r-- 15.9 KB
public_test_util.py -rw-r--r-- 10.3 KB
random.py -rw-r--r-- 97.4 KB
shard_alike.py -rw-r--r-- 3.7 KB
sharding.py -rw-r--r-- 7.6 KB
sharding_impls.py -rw-r--r-- 59.1 KB
sharding_specs.py -rw-r--r-- 8.4 KB
source_info_util.py -rw-r--r-- 9.9 KB
sourcemap.py -rw-r--r-- 6.5 KB
stages.py -rw-r--r-- 28.3 KB
test_util.py -rw-r--r-- 68.1 KB
tpu_custom_call.py -rw-r--r-- 17.8 KB
traceback_util.py -rw-r--r-- 8.9 KB
tree.py -rw-r--r-- 9.1 KB
tree_util.py -rw-r--r-- 45.7 KB
typing.py -rw-r--r-- 3.2 KB
util.py -rw-r--r-- 19.5 KB
xla_bridge.py -rw-r--r-- 42.9 KB

back to top