|
..
|
|
_static
|
|
_templates
|
|
debugging
|
|
jax-101
|
|
jep
|
|
notebooks
|
|
pallas
|
|
sphinxext
|
|
tutorials
|
|
autodidax.ipynb
|
|
Custom_Operation_for_GPUs.md
|
|
README.md
|
|
advanced_guide.rst
|
|
aot.md
|
|
api_compatibility.md
|
|
async_dispatch.rst
|
|
autodidax.md
|
|
autodidax.py
|
|
beginner_guide.rst
|
|
building_on_jax.md
|
|
changelog.md
|
|
concurrency.rst
|
|
conf.py
|
|
contributing.md
|
|
contributor_guide.rst
|
|
deprecation.md
|
|
developer.md
|
|
device_memory_profiling.md
|
|
errors.rst
|
|
faq.rst
|
|
glossary.rst
|
|
gpu_memory_allocation.rst
|
|
gpu_performance_tips.md
|
|
index.rst
|
|
installation.md
|
|
investigating_a_regression.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.array_api.rst
|
|
jax.experimental.checkify.rst
|
|
jax.experimental.compilation_cache.rst
|
|
jax.experimental.custom_partitioning.rst
|
|
jax.experimental.host_callback.rst
|
|
jax.experimental.jet.rst
|
|
jax.experimental.key_reuse.rst
|
|
jax.experimental.maps.rst
|
|
jax.experimental.multihost_utils.rst
|
|
jax.experimental.pjit.rst
|
|
jax.experimental.rst
|
|
jax.experimental.sparse.rst
|
|
jax.extend.linear_util.rst
|
|
jax.extend.mlir.rst
|
|
jax.extend.random.rst
|
|
jax.extend.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.rst
|
|
jax.scipy.rst
|
|
jax.sharding.rst
|
|
jax.stages.rst
|
|
jax.tree_util.rst
|
|
jax.typing.rst
|
|
jax_array_migration.md
|
|
jax_internal_api.rst
|
|
jaxpr.rst
|
|
multi_process.md
|
|
notes.rst
|
|
profiling.md
|
|
pytrees.md
|
|
rank_promotion_warning.rst
|
|
requirements.txt
|
|
transfer_guard.rst
|
|
type_promotion.rst
|
|
user_guides.rst
|