Files

.github
benchmarks
build
ci
cloud_tpu_colabs
docs
- _static
- _templates
- _tutorials
- debugging
- export
- ffi
- jep
- notebooks
- pallas
- sphinxext
- README.md
- about.md
- advanced-autodiff.md
- advanced_guide.rst
- aot.md
- api_compatibility.md
- async_dispatch.rst
- autodidax.ipynb
- autodidax.md
- autodidax.py
- autodidax2_part1.ipynb
- 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
- contributing.md
- contributor_guide.rst
- control-flow.md
- debugging.md
- default_dtypes.md
- deprecation.md
- developer.md
- device_memory_profiling.md
- distributed_data_loading.md
- errors.rst
- extensions.rst
- external-callbacks.md
- faq.rst
- ffi.ipynb
- 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.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
- quickstart.md
- random-numbers.md
- rank_promotion_warning.rst
- requirements.txt
- sharded-computation.ipynb
- sharded-computation.md
- stateful-computations.md
- transfer_guard.rst
- tutorials.rst
- type_promotion.rst
- user_guides.rst
- working-with-pytrees.md
- xla_flags.md
examples
images
jax
jax_plugins
jaxlib
tests
third_party
.bazelrc
.bazelversion
.editorconfig
.gitignore
.pre-commit-config.yaml
.readthedocs.yml
AUTHORS
BUILD.bazel
CHANGELOG.md
CITATION.bib
CONTRIBUTING.md
LICENSE
README.md
WORKSPACE
build_wheel.py
conftest.py
platform_mappings
pyproject.toml
setup.py

docs

DOC: add documentation note about default dtypes

Mar 28, 2025

dafebd0 · Mar 28, 2025

Name	Name	Last commit message	Last commit date
parent directory ..
_static	_static	Add navigation breadcrumbs to docs.	Mar 12, 2025
_templates	_templates	Add the JAX logo to sphinx docs	Aug 5, 2019
_tutorials	_tutorials	Remove the tutorial about legacy custom calls.	Feb 7, 2025
debugging	debugging	Fixed printing order of results in jax.debug.print documentation.	Feb 28, 2025
export	export	Rename some internal APIs (set_abstract_mesh -> use_abstract_mesh and…	Mar 13, 2025
ffi	ffi	Fix typos	Feb 18, 2025
jep	jep	Rename xla_extension_version to jaxlib_extension_version to reflect i…	Mar 26, 2025
notebooks	notebooks	Reverts `e0c0933`	Mar 20, 2025
pallas	pallas	Fix error in pallas tutorial	Mar 17, 2025
sphinxext	sphinxext	Update references to the GitHub url in JAX codebase to reflect move f…	Sep 20, 2024
README.md	README.md	Fixes to the documentation	Oct 9, 2019
about.md	about.md	Update references to JAX's GitHub repo	Mar 5, 2025
advanced-autodiff.md	advanced-autodiff.md	Update the advanced autodiff tutorial and replace some vmap with grad	Jan 7, 2025
advanced_guide.rst	advanced_guide.rst	Explicit sharding docs	Mar 14, 2025
aot.md	aot.md	[better_errors] Cleanup use of DebugInfo.arg_names and result_paths	Feb 23, 2025
api_compatibility.md	api_compatibility.md	DOC: clarify API compatibility discussion & mention EffVer	Dec 19, 2024
async_dispatch.rst	async_dispatch.rst	update documentation notes to new-style typed keys	Mar 7, 2024
autodidax.ipynb	autodidax.ipynb	documentation: fix function name	Feb 4, 2025
autodidax.md	autodidax.md	documentation: fix function name	Feb 4, 2025
autodidax.py	autodidax.py	documentation: fix function name	Feb 4, 2025
autodidax2_part1.ipynb	autodidax2_part1.ipynb	Part 1 of a new autodidax based on "stackless"	Feb 12, 2025
autodidax2_part1.md	autodidax2_part1.md	Part 1 of a new autodidax based on "stackless"	Feb 12, 2025
autodidax2_part1.py	autodidax2_part1.py	Part 1 of a new autodidax based on "stackless"	Feb 12, 2025
automatic-differentiation.md	automatic-differentiation.md	pre-commit: update hooks & pin using hashes	Aug 27, 2024
automatic-vectorization.md	automatic-vectorization.md	pre-commit: update hooks & pin using hashes	Aug 27, 2024
beginner_guide.rst	beginner_guide.rst	Update references to the GitHub url in JAX codebase to reflect move f…	Sep 20, 2024
build_custom_gpu.sh	build_custom_gpu.sh	Current status + build script fixes	Apr 8, 2024
building_on_jax.md	building_on_jax.md	docs: sentence case page titles, section headings, some content	Aug 13, 2024
changelog.md	changelog.md	move changelog to top level	Mar 8, 2021
concurrency.rst	concurrency.rst	Document that JAX follows the NEP-29 deprecation policy.	Apr 21, 2021
conf.py	conf.py	Explicit sharding docs	Mar 14, 2025
contributing.md	contributing.md	Update references to the GitHub url in JAX codebase to reflect move f…	Sep 20, 2024
contributor_guide.rst	contributor_guide.rst	Part 1 of a new autodidax based on "stackless"	Feb 12, 2025
control-flow.md	control-flow.md	Add an example on logical operators to the tutorial.	Nov 15, 2024
debugging.md	debugging.md	docs: shorten/clarify some page titles	Aug 28, 2024
default_dtypes.md	default_dtypes.md	DOC: add documentation note about default dtypes	Mar 28, 2025
deprecation.md	deprecation.md	DOC: update numpy/scipy versions in deprecation doc	Feb 6, 2025
developer.md	developer.md	Update instructions for usage of `:build_jaxlib=false` flag.	Mar 27, 2025
device_memory_profiling.md	device_memory_profiling.md	docs: tidy up titles and headings	Aug 13, 2024
distributed_data_loading.md	distributed_data_loading.md	Use `jax.make_array_from_process_local_data` API in distributed data …	Sep 23, 2024
errors.rst	errors.rst	Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeEr…	Sep 26, 2024
extensions.rst	extensions.rst	Remove the tutorial about legacy custom calls.	Feb 7, 2025
external-callbacks.md	external-callbacks.md	Revert #25982 since callbacks can now use JAX functions.	Jan 29, 2025
faq.rst	faq.rst	Move Control Flow text from Sharp Bits into its own tutorial.	Nov 14, 2024
ffi.ipynb	ffi.ipynb	Add a discussion of sharding to the FFI tutorial.	Jan 9, 2025
ffi.md	ffi.md	Add a discussion of sharding to the FFI tutorial.	Jan 9, 2025
glossary.rst	glossary.rst	Add/update JAX Advanced Tutorials docs, ToC structure	Sep 20, 2024
gpu_memory_allocation.rst	gpu_memory_allocation.rst	Document cudaMallocAsync as an experimental feature.	Dec 6, 2024
gpu_performance_tips.md	gpu_performance_tips.md	Move PGLE documentation to JAX docs.	Mar 24, 2025
gradient-checkpointing.md	gradient-checkpointing.md	consolidate the code example	Dec 19, 2024
hero.html	hero.html	DOC: update main landing page style	Nov 15, 2024
index.rst	index.rst	DOC: update main landing page style	Nov 15, 2024
installation.md	installation.md	Change documentation to recommend libtpu from pypi instead of GCS.	Feb 24, 2025
investigating_a_regression.md	investigating_a_regression.md	Update references to the GitHub url in JAX codebase to reflect move f…	Sep 20, 2024
jax-primitives.md	jax-primitives.md	Add/update JAX Advanced Tutorials docs, ToC structure	Sep 20, 2024
jax.debug.rst	jax.debug.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.distributed.rst	jax.distributed.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.dlpack.rst	jax.dlpack.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.dtypes.rst	jax.dtypes.rst	custom prng: introduce mechanism to identify key arrays by dtype	Jul 21, 2023
jax.example_libraries.optimizers.rst	jax.example_libraries.optimizers.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.example_libraries.rst	jax.example_libraries.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.example_libraries.stax.rst	jax.example_libraries.stax.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.experimental.checkify.rst	jax.experimental.checkify.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.experimental.compilation_cache.rst	jax.experimental.compilation_cache.rst	Cleanup deprecated compilation cache APIs.	Jan 13, 2024
jax.experimental.custom_dce.rst	jax.experimental.custom_dce.rst	Add custom_dce to changelogs and API docs.	Jan 27, 2025
jax.experimental.custom_partitioning.rst	jax.experimental.custom_partitioning.rst	Adding @custom_partitioning to jax.experimental API doc	Jan 20, 2023
jax.experimental.jet.rst	jax.experimental.jet.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.experimental.key_reuse.rst	jax.experimental.key_reuse.rst	[key reuse] define KeyReuseError in jax.errors	Mar 8, 2024
jax.experimental.mesh_utils.rst	jax.experimental.mesh_utils.rst	Add JAX mesh_utils docs	Feb 13, 2024
jax.experimental.multihost_utils.rst	jax.experimental.multihost_utils.rst	Fix a sphinx error.	May 5, 2023
jax.experimental.pallas.mosaic_gpu.rst	jax.experimental.pallas.mosaic_gpu.rst	[pallas:mosaic_gpu] `plgpu.copy_smem_to_gmem` no longer transparently…	Oct 31, 2024
jax.experimental.pallas.rst	jax.experimental.pallas.rst	[pallas] Added API docs for Triton and Mosaic GPU backends	Oct 10, 2024
jax.experimental.pallas.tpu.rst	jax.experimental.pallas.tpu.rst	[pallas] Added API docs for Triton and Mosaic GPU backends	Oct 10, 2024
jax.experimental.pallas.triton.rst	jax.experimental.pallas.triton.rst	[pallas] Added API docs for Triton and Mosaic GPU backends	Oct 10, 2024
jax.experimental.pjit.rst	jax.experimental.pjit.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.experimental.rst	jax.experimental.rst	Add custom_dce to changelogs and API docs.	Jan 27, 2025
jax.experimental.serialize_executable.rst	jax.experimental.serialize_executable.rst	Add doc for jax.experimental.serialize_executable module	May 14, 2024
jax.experimental.shard_map.rst	jax.experimental.shard_map.rst	shard_map: add API docs	May 13, 2024
jax.experimental.sparse.rst	jax.experimental.sparse.rst	Add performance note at the top of sparse docs	Jan 17, 2025
jax.export.rst	jax.export.rst	[export] Improved the documentation.	Dec 8, 2024
jax.extend.core.rst	jax.extend.core.rst	jax.core: deprecate a number of APIs	Dec 10, 2024
jax.extend.linear_util.rst	jax.extend.linear_util.rst	Add jax.extend.mlir.	Dec 6, 2023
jax.extend.mlir.rst	jax.extend.mlir.rst	Add jax.extend.mlir.	Dec 6, 2023
jax.extend.random.rst	jax.extend.random.rst	Add jax.extend.mlir.	Dec 6, 2023
jax.extend.rst	jax.extend.rst	Move jex.ffi to jax.ffi.	Dec 29, 2024
jax.ffi.rst	jax.ffi.rst	Add a register_custom_type_id function to the GPU plugins.	Jan 7, 2025
jax.flatten_util.rst	jax.flatten_util.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.image.rst	jax.image.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.lax.rst	jax.lax.rst	Fix syntax error and typos for composite primitive docstring.	Mar 11, 2025
jax.lib.rst	jax.lib.rst	Finalize some deprecations in jax.core, jax.lib.xla_bridge, and jax.l…	Dec 11, 2024
jax.nn.initializers.rst	jax.nn.initializers.rst	update package/API reference docs to new-style typed PRNG keys	Mar 7, 2024
jax.nn.rst	jax.nn.rst	Add identity activation	Mar 23, 2025
jax.numpy.rst	jax.numpy.rst	Remove deprecated jax.experimental.array_api	Jan 6, 2025
jax.ops.rst	jax.ops.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.profiler.rst	jax.profiler.rst	DOC: make API doc titles more uniform	Jan 18, 2023
jax.random.rst	jax.random.rst	Add jax.random.multinomial.	Mar 12, 2025
jax.rst	jax.rst	Add `jax.copy_to_host_async(tree)`.	Feb 27, 2025
jax.scipy.rst	jax.scipy.rst	add pascal matrix	Mar 26, 2025
jax.sharding.rst	jax.sharding.rst	Deprecate `XLACompatibleSharding` in favor of `jax.sharding.Sharding`.	Jun 5, 2024
jax.stages.rst	jax.stages.rst	docstrings and API reference doc listing for the traced AOT stage	Feb 12, 2025
jax.test_util.rst	jax.test_util.rst	add docstrings for check_vjp and check_jvp	Jan 24, 2025
jax.tree.rst	jax.tree.rst	Add jax.tree shortcuts for .*_with_path calls, for convenience of users.	Dec 12, 2024
jax.tree_util.rst	jax.tree_util.rst	Export `KeyPath` and related types to `jax.tree_util`	Jul 31, 2024
jax.typing.rst	jax.typing.rst	Define jax.typing.DTypeLike	Oct 10, 2023
jax_array_migration.md	jax_array_migration.md	Remove the jax Array migration guide from the TOC tree but keep the d…	Mar 19, 2025
jaxpr.md	jaxpr.md	Add/update JAX Advanced Tutorials docs, ToC structure	Sep 20, 2024
jit-compilation.md	jit-compilation.md	Move Control Flow text from Sharp Bits into its own tutorial.	Nov 14, 2024
key-concepts.md	key-concepts.md	Consolidate material on PRNGs and add a short summary to Key Concepts.	Nov 15, 2024
multi_process.md	multi_process.md	docs: shorten/clarify some page titles	Aug 28, 2024
notes.rst	notes.rst	DOC: add documentation note about default dtypes	Mar 28, 2025
persistent_compilation_cache.md	persistent_compilation_cache.md	Add docs	Dec 13, 2024
profiling.md	profiling.md	tbp nightly instructions	Jan 2, 2025
pytrees.md	pytrees.md	Add freshness metablock to JAX OSS docs.	Jun 21, 2024
quickstart.md	quickstart.md	Fix Typos	Sep 25, 2024
random-numbers.md	random-numbers.md	Consolidate material on PRNGs and add a short summary to Key Concepts.	Nov 15, 2024
rank_promotion_warning.rst	rank_promotion_warning.rst	Avoid 'from jax import config' imports	Apr 11, 2024
requirements.txt	requirements.txt	fix the doc error: module 'scipy.misc' has no attribute 'face'	Jan 6, 2025
sharded-computation.ipynb	sharded-computation.ipynb	Merge pull request #26279 from MichaelHudgins:tsan-resultstore	Feb 6, 2025
sharded-computation.md	sharded-computation.md	Merge pull request #26279 from MichaelHudgins:tsan-resultstore	Feb 6, 2025
stateful-computations.md	stateful-computations.md	Update stateful-computations.md	Mar 10, 2025
transfer_guard.rst	transfer_guard.rst	Add a documentation for the JAX transfer guard	Mar 25, 2022
tutorials.rst	tutorials.rst	Move Control Flow text from Sharp Bits into its own tutorial.	Nov 14, 2024
type_promotion.rst	type_promotion.rst	Typo	Sep 5, 2024
user_guides.rst	user_guides.rst	DOC: add documentation note about default dtypes	Mar 28, 2025
working-with-pytrees.md	working-with-pytrees.md	Fix formatting in the docs for transposing pytrees	Dec 26, 2024
xla_flags.md	xla_flags.md	[Docs] Remove `--xla_gpu_enable_triton_softmax_fusion` from docs	Jan 14, 2025

README.md

To rebuild the documentation, see Update Documentation.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Files

docs

docs

README.md

Files

docs

Directory actions

More options

Directory actions

More options

Latest commit

History

docs

Folders and files

parent directory

README.md