|
..
|
|
_static
|
|
_templates
|
|
_tutorials
|
|
debugging
|
|
export
|
|
ffi
|
|
internals
|
|
jep
|
|
new_docs
|
|
notebooks
|
|
pallas
|
|
sphinxext
|
|
array_refs.ipynb
|
|
autodidax.ipynb
|
|
autodidax2_part1.ipynb
|
|
ffi.ipynb
|
|
hijax_custom_derivatives.ipynb
|
|
hijax_types.ipynb
|
|
parallel.ipynb
|
|
README.md
|
|
__init__.py
|
|
about.md
|
|
advanced_autodiff.md
|
|
advanced_guides.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
|
|
benchmarking.md
|
|
buffer_donation.md
|
|
build_custom_gpu.sh
|
|
building_on_jax.md
|
|
changelog.md
|
|
complex-differentiation.md
|
|
concurrency.rst
|
|
conf.py
|
|
config_options.rst
|
|
contributing.md
|
|
contributor_guide.rst
|
|
control-flow.md
|
|
custom_pytrees.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
|
|
fault_tolerance.rst
|
|
ffi.md
|
|
ffi.py
|
|
glossary.rst
|
|
gpu_memory_allocation.rst
|
|
gpu_performance_tips.md
|
|
gradient-checkpointing.md
|
|
hero.html
|
|
higher-order.md
|
|
hijax_custom_derivatives.md
|
|
hijax_custom_derivatives.py
|
|
hijax_types.md
|
|
hijax_types.py
|
|
index.rst
|
|
installation.md
|
|
investigating_a_regression.md
|
|
jacobian-vector-products.md
|
|
jax-101.rst
|
|
jax-primitives.md
|
|
jax.ad_checkpoint.rst
|
|
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_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.random.rst
|
|
jax.experimental.rst
|
|
jax.experimental.serialize_executable.rst
|
|
jax.experimental.sparse.rst
|
|
jax.experimental.xla_metadata.rst
|
|
jax.export.rst
|
|
jax.extend.backend.rst
|
|
jax.extend.core.rst
|
|
jax.extend.linear_util.rst
|
|
jax.extend.lowering.rst
|
|
jax.extend.mlir.rst
|
|
jax.extend.pallas.rst
|
|
jax.extend.random.rst
|
|
jax.extend.rst
|
|
jax.extend.xla.rst
|
|
jax.ffi.rst
|
|
jax.flatten_util.rst
|
|
jax.image.rst
|
|
jax.lax.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
|
|
migrate_pmap.md
|
|
multi_process.md
|
|
notes.rst
|
|
parallel.md
|
|
parallel.py
|
|
persistent_compilation_cache.md
|
|
profiling.md
|
|
pytrees.md
|
|
random-numbers.md
|
|
rank_promotion_warning.rst
|
|
requirements.txt
|
|
shardy_jax_migration.md
|
|
stateful-computations.md
|
|
the-training-cookbook.py
|
|
the-training-cookbook.rst
|
|
tracing.md
|
|
transfer_guard.rst
|
|
type_promotion.rst
|
|
xla_flags.md
|