|
..
|
|
_static
|
|
_templates
|
|
_tutorials
|
|
debugging
|
|
export
|
|
ffi
|
|
internals
|
|
jep
|
|
notebooks
|
|
pallas
|
|
sphinxext
|
|
array_refs.ipynb
|
|
autodidax.ipynb
|
|
autodidax2_part1.ipynb
|
|
ffi.ipynb
|
|
sharded-computation.ipynb
|
|
README.md
|
|
about.md
|
|
advanced-autodiff.md
|
|
advanced_guide.rst
|
|
aot.md
|
|
api_compatibility.md
|
|
array_refs.md
|
|
array_refs.py
|
|
async_dispatch.rst
|
|
autodidax.md
|
|
autodidax.py
|
|
autodidax2_part1.md
|
|
autodidax2_part1.py
|
|
automatic-differentiation.md
|
|
automatic-vectorization.md
|
|
beginner_guide.rst
|
|
build_custom_gpu.sh
|
|
building_on_jax.md
|
|
changelog.md
|
|
concurrency.rst
|
|
conf.py
|
|
config_options.rst
|
|
contributing.md
|
|
contributor_guide.rst
|
|
control-flow.md
|
|
debugging.md
|
|
default_dtypes.md
|
|
deprecation.md
|
|
developer.md
|
|
device_memory_profiling.md
|
|
direct_linearize_migration.md
|
|
distributed_data_loading.md
|
|
errors.rst
|
|
extensions.rst
|
|
external-callbacks.md
|
|
faq.rst
|
|
ffi.md
|
|
glossary.rst
|
|
gpu_memory_allocation.rst
|
|
gpu_performance_tips.md
|
|
gradient-checkpointing.md
|
|
hero.html
|
|
index.rst
|
|
installation.md
|
|
investigating_a_regression.md
|
|
jax-primitives.md
|
|
jax.debug.rst
|
|
jax.distributed.rst
|
|
jax.dlpack.rst
|
|
jax.dtypes.rst
|
|
jax.example_libraries.optimizers.rst
|
|
jax.example_libraries.rst
|
|
jax.example_libraries.stax.rst
|
|
jax.experimental.checkify.rst
|
|
jax.experimental.compilation_cache.rst
|
|
jax.experimental.custom_dce.rst
|
|
jax.experimental.custom_partitioning.rst
|
|
jax.experimental.jet.rst
|
|
jax.experimental.key_reuse.rst
|
|
jax.experimental.mesh_utils.rst
|
|
jax.experimental.multihost_utils.rst
|
|
jax.experimental.pallas.mosaic_gpu.rst
|
|
jax.experimental.pallas.rst
|
|
jax.experimental.pallas.tpu.rst
|
|
jax.experimental.pallas.triton.rst
|
|
jax.experimental.pjit.rst
|
|
jax.experimental.rst
|
|
jax.experimental.serialize_executable.rst
|
|
jax.experimental.shard_map.rst
|
|
jax.experimental.sparse.rst
|
|
jax.export.rst
|
|
jax.extend.core.rst
|
|
jax.extend.linear_util.rst
|
|
jax.extend.mlir.rst
|
|
jax.extend.random.rst
|
|
jax.extend.rst
|
|
jax.ffi.rst
|
|
jax.flatten_util.rst
|
|
jax.image.rst
|
|
jax.lax.rst
|
|
jax.lib.rst
|
|
jax.nn.initializers.rst
|
|
jax.nn.rst
|
|
jax.numpy.rst
|
|
jax.ops.rst
|
|
jax.profiler.rst
|
|
jax.random.rst
|
|
jax.ref.rst
|
|
jax.rst
|
|
jax.scipy.rst
|
|
jax.sharding.rst
|
|
jax.stages.rst
|
|
jax.test_util.rst
|
|
jax.tree.rst
|
|
jax.tree_util.rst
|
|
jax.typing.rst
|
|
jax_array_migration.md
|
|
jaxpr.md
|
|
jit-compilation.md
|
|
key-concepts.md
|
|
multi_process.md
|
|
notes.rst
|
|
persistent_compilation_cache.md
|
|
profiling.md
|
|
pytrees.md
|
|
random-numbers.md
|
|
rank_promotion_warning.rst
|
|
requirements.txt
|
|
sharded-computation.md
|
|
shardy_jax_migration.md
|
|
stateful-computations.md
|
|
the-training-cookbook.py
|
|
the-training-cookbook.rst
|
|
transfer_guard.rst
|
|
tutorials.rst
|
|
type_promotion.rst
|
|
user_guides.rst
|
|
working-with-pytrees.md
|
|
xla_flags.md
|