mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/docs
History
Jakob Roth 4f8699c8a1 Update docs of callbacks 
Callback functions should not call into JAX. This information was
missing in the docs of the callbacks. This commit adds this information
to the docs.

See: #25861, #24255
2025-01-19 11:33:20 +01:00
..
_static
[Pallas] Update TPU documentation
2024-12-06 10:59:11 -08:00
_templates
Add the JAX logo to sphinx docs
2019-08-05 12:34:29 -07:00
_tutorials
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
debugging
docs: runtime debugging tweaks
2024-08-28 11:04:14 -07:00
export
[export] Expand exporting to work with AbstractMesh.
2024-12-16 10:30:46 +02:00
ffi
Update some outdated syntax in FFI tutorial.
2024-11-12 08:34:24 -08:00
gpu_ops
Current status + build script fixes
2024-04-08 04:48:28 -07:00
jep
enable partitionable threefry by default
2025-01-13 22:46:24 -08:00
notebooks
fix the doc error: module 'scipy.misc' has no attribute 'face'
2025-01-06 12:00:02 -08:00
pallas
[Pallas TPU] Add vector support to pl.debug_print
2025-01-13 13:22:21 -08:00
sphinxext
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
about.md
add about page
2024-11-07 21:47:50 -08:00
advanced_guide.rst
Add md page about XLA&XLA_FLAGS in JAX
2024-09-05 09:27:32 +00:00
advanced-autodiff.md
Update the advanced autodiff tutorial and replace some vmap with grad
2025-01-07 15:56:23 +01:00
aot.md
Jax: Stop returning a list of cost-analyses.
2025-01-15 09:53:59 -08:00
api_compatibility.md
DOC: clarify API compatibility discussion & mention EffVer
2024-12-19 11:23:43 -08:00
async_dispatch.rst
update documentation notes to new-style typed keys
2024-03-07 12:40:09 -08:00
autodidax.ipynb
Fix some typos
2024-11-19 15:04:55 +00:00
autodidax.md
Fix some typos
2024-11-19 15:04:55 +00:00
autodidax.py
Fix some typos
2024-11-19 15:04:55 +00:00
automatic-differentiation.md
pre-commit: update hooks & pin using hashes
2024-08-27 15:23:13 -07:00
automatic-vectorization.md
pre-commit: update hooks & pin using hashes
2024-08-27 15:23:13 -07:00
beginner_guide.rst
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
build_custom_gpu.sh
Current status + build script fixes
2024-04-08 04:48:28 -07:00
building_on_jax.md
docs: sentence case page titles, section headings, some content
2024-08-12 18:12:17 -07:00
changelog.md
move changelog to top level
2021-03-08 10:44:52 -08:00
concurrency.rst
Document that JAX follows the NEP-29 deprecation policy.
2021-04-21 11:12:41 -04:00
conf.py
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
contributing.md
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
contributor_guide.rst
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
control-flow.md
Add an example on logical operators to the tutorial.
2024-11-15 12:40:41 -08:00
Custom_Operation_for_GPUs.md
Fix some typos
2024-11-19 15:04:55 +00:00
Custom_Operation_for_GPUs.py
Fix typos: Change 'arugments' to 'arguments'.
2024-11-13 15:58:45 -05:00
debugging.md
docs: shorten/clarify some page titles
2024-08-28 11:22:06 -07:00
deprecation.md
Mention python 3.13 in docs & package metadata
2024-11-20 08:23:19 -08:00
developer.md
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
device_memory_profiling.md
docs: tidy up titles and headings
2024-08-13 11:53:57 -07:00
distributed_data_loading.md
Use jax.make_array_from_process_local_data API in distributed data loading doc
2024-09-23 16:03:34 -07:00
errors.rst
Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class.
2024-09-26 08:39:30 -07:00
extensions.rst
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
external-callbacks.md
Update docs of callbacks
2025-01-19 11:33:20 +01:00
faq.rst
Move Control Flow text from Sharp Bits into its own tutorial.
2024-11-14 11:07:52 -08:00
ffi.ipynb
Add a discussion of sharding to the FFI tutorial.
2025-01-09 11:24:40 -05:00
ffi.md
Add a discussion of sharding to the FFI tutorial.
2025-01-09 11:24:40 -05:00
glossary.rst
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
gpu_memory_allocation.rst
Document cudaMallocAsync as an experimental feature.
2024-12-06 13:41:35 -05:00
gpu_performance_tips.md
Merge pull request #25444 from justinjfu:gpu_docs_update
2024-12-13 11:05:31 -08:00
gradient-checkpointing.md
consolidate the code example
2024-12-19 14:47:29 -08:00
hero.html
DOC: update main landing page style
2024-11-15 13:44:31 -08:00
index.rst
DOC: update main landing page style
2024-11-15 13:44:31 -08:00
installation.md
Add quotes to pip commands in docs around option install for zsh
2025-01-16 13:30:43 -08:00
investigating_a_regression.md
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
jax_array_migration.md
Add freshness metablock to JAX OSS docs.
2024-06-21 14:50:49 -07:00
jax-primitives.md
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
jax.debug.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.distributed.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.dlpack.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.dtypes.rst
custom prng: introduce mechanism to identify key arrays by dtype
2023-07-21 12:27:32 -07:00
jax.example_libraries.optimizers.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.example_libraries.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.example_libraries.stax.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.experimental.checkify.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.experimental.compilation_cache.rst
Cleanup deprecated compilation cache APIs.
2024-01-12 22:44:48 -08:00
jax.experimental.custom_partitioning.rst
Adding @custom_partitioning to jax.experimental API doc
2023-01-20 11:34:29 -08:00
jax.experimental.jet.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.experimental.key_reuse.rst
[key reuse] define KeyReuseError in jax.errors
2024-03-08 10:59:06 -08:00
jax.experimental.mesh_utils.rst
Add JAX mesh_utils docs
2024-02-13 00:56:21 +00:00
jax.experimental.multihost_utils.rst
Fix a sphinx error.
2023-05-05 11:08:18 -07:00
jax.experimental.pallas.mosaic_gpu.rst
[pallas:mosaic_gpu] plgpu.copy_smem_to_gmem no longer transparently commits SMEM
2024-10-31 02:21:10 -07:00
jax.experimental.pallas.rst
[pallas] Added API docs for Triton and Mosaic GPU backends
2024-10-10 12:27:53 +01:00
jax.experimental.pallas.tpu.rst
[pallas] Added API docs for Triton and Mosaic GPU backends
2024-10-10 12:27:53 +01:00
jax.experimental.pallas.triton.rst
[pallas] Added API docs for Triton and Mosaic GPU backends
2024-10-10 12:27:53 +01:00
jax.experimental.pjit.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.experimental.rst
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
jax.experimental.serialize_executable.rst
Add doc for jax.experimental.serialize_executable module
2024-05-14 19:11:37 -04:00
jax.experimental.shard_map.rst
shard_map: add API docs
2024-05-13 13:04:15 -07:00
jax.experimental.sparse.rst
Add performance note at the top of sparse docs
2025-01-17 10:17:16 -08:00
jax.export.rst
[export] Improved the documentation.
2024-12-08 17:39:24 +01:00
jax.extend.core.rst
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
jax.extend.linear_util.rst
Add jax.extend.mlir.
2023-12-06 09:16:43 -08:00
jax.extend.mlir.rst
Add jax.extend.mlir.
2023-12-06 09:16:43 -08:00
jax.extend.random.rst
Add jax.extend.mlir.
2023-12-06 09:16:43 -08:00
jax.extend.rst
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
jax.ffi.rst
Add a register_custom_type_id function to the GPU plugins.
2025-01-07 07:29:38 -08:00
jax.flatten_util.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.image.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.lax.rst
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
jax.lib.rst
Finalize some deprecations in jax.core, jax.lib.xla_bridge, and jax.lib.xla_client.
2024-12-11 09:50:33 -08:00
jax.nn.initializers.rst
update package/API reference docs to new-style typed PRNG keys
2024-03-07 12:40:09 -08:00
jax.nn.rst
Update jax doc sdpa
2024-07-15 17:30:54 +00:00
jax.numpy.rst
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
jax.ops.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.profiler.rst
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.random.rst
emphasize typed over legacy RNG keys in random module docs
2024-08-11 12:41:50 -07:00
jax.rst
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
jax.scipy.rst
feat(lib): add real-valued implementation of jax.scipy.special.fresnel
2024-09-03 09:50:19 +02:00
jax.sharding.rst
Deprecate XLACompatibleSharding in favor of jax.sharding.Sharding.
2024-06-05 09:07:27 -07:00
jax.stages.rst
Add trace to stages.Wrapped and add docs for it.
2024-06-13 14:49:54 -07:00
jax.tree_util.rst
Export KeyPath and related types to jax.tree_util
2024-07-31 06:41:33 -07:00
jax.tree.rst
Add jax.tree shortcuts for .*_with_path calls, for convenience of users.
2024-12-12 15:13:32 -08:00
jax.typing.rst
Define jax.typing.DTypeLike
2023-10-10 08:46:36 -07:00
jaxpr.md
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
jit-compilation.md
Move Control Flow text from Sharp Bits into its own tutorial.
2024-11-14 11:07:52 -08:00
key-concepts.md
Consolidate material on PRNGs and add a short summary to Key Concepts.
2024-11-15 14:44:57 -08:00
multi_process.md
docs: shorten/clarify some page titles
2024-08-28 11:22:06 -07:00
notes.rst
DOC: re-organize notes section
2023-01-30 11:57:06 -08:00
persistent_compilation_cache.md
When caching is enabled, also enable XLA caching features as well
2024-11-13 10:30:04 -08:00
profiling.md
tbp nightly instructions
2025-01-02 09:49:31 +01:00
pytrees.md
Add freshness metablock to JAX OSS docs.
2024-06-21 14:50:49 -07:00
quickstart.md
Fix Typos
2024-09-25 21:26:05 +05:30
random-numbers.md
Consolidate material on PRNGs and add a short summary to Key Concepts.
2024-11-15 14:44:57 -08:00
rank_promotion_warning.rst
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
README.md
Fixes to the documentation
2019-10-09 17:45:09 +02:00
requirements.txt
fix the doc error: module 'scipy.misc' has no attribute 'face'
2025-01-06 12:00:02 -08:00
sharded-computation.ipynb
Don't create a mesh inside jit in sharded-computation tutorial
2024-12-17 16:11:00 -08:00
sharded-computation.md
Don't create a mesh inside jit in sharded-computation tutorial
2024-12-17 16:11:00 -08:00
stateful-computations.md
Move Control Flow text from Sharp Bits into its own tutorial.
2024-11-14 11:07:52 -08:00
transfer_guard.rst
Add a documentation for the JAX transfer guard
2022-03-25 19:50:12 +00:00
tutorials.rst
Move Control Flow text from Sharp Bits into its own tutorial.
2024-11-14 11:07:52 -08:00
type_promotion.rst
Typo
2024-09-05 20:47:40 +02:00
user_guides.rst
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
working-with-pytrees.md
Fix formatting in the docs for transposing pytrees
2024-12-26 01:18:10 +00:00
xla_flags.md
[Docs] Remove --xla_gpu_enable_triton_softmax_fusion from docs
2025-01-14 12:50:59 -08:00

README.md

To rebuild the documentation, see Update Documentation.