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
Peter Hawkins 46f0a3eee7 Clone RandomAlgorithm into lax.py, instead of using the version from XLA. 
Change in preparation for removing HLO ops from the XLA Python bindings.

In passing, also:
* improve how the documentation of FftType renders.
* remove some stale references to xla_client
* remove the standard_translate rule, which is unused.

PiperOrigin-RevId: 684892102
2024-10-11 11:03:15 -07:00
..
_static
[Pallas] Add block-sparse kernel tutorial
2024-09-19 12:23:03 -07: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
cuda_custom_call
Replace disable_backends with enable_backends on jax_multiplatform_test.
2024-09-27 06:15:31 -07:00
debugging
docs: runtime debugging tweaks
2024-08-28 11:04:14 -07:00
export
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
ffi
Add ffi_call tutorial
2024-08-01 15:36:32 -04:00
gpu_ops
Current status + build script fixes
2024-04-08 04:48:28 -07:00
jep
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
notebooks
[Pallas] [Docs] Replace full urls with label-based cross references
2024-10-03 14:35:51 -07:00
pallas
Added jax.experimental.pallas.mosaic_gpu
2024-10-07 04:05:08 -07: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
advanced_guide.rst
Add md page about XLA&XLA_FLAGS in JAX
2024-09-05 09:27:32 +00:00
advanced-autodiff.md
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
aot.md
Add freshness metablock to JAX OSS docs.
2024-06-21 14:50:49 -07:00
api_compatibility.md
doc: remove references to submodules that no longer exist
2024-06-28 12:39:14 -07:00
async_dispatch.rst
update documentation notes to new-style typed keys
2024-03-07 12:40:09 -08:00
autodidax.ipynb
Port autodidax to use StableHLO instead of classic HLO.
2024-10-11 08:25:05 -04:00
autodidax.md
Port autodidax to use StableHLO instead of classic HLO.
2024-10-11 08:25:05 -04:00
autodidax.py
Port autodidax to use StableHLO instead of classic HLO.
2024-10-11 08:25:05 -04: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
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +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
docs: reorganize sections
2024-08-13 21:33:45 -07:00
Custom_Operation_for_GPUs.md
docs: update Custom_Operation_for_GPUs.md
2024-10-02 12:57:45 +09:00
Custom_Operation_for_GPUs.py
update README and several docs to typed RNG keys
2024-08-11 08:09:47 -07:00
debugging.md
docs: shorten/clarify some page titles
2024-08-28 11:22:06 -07:00
deprecation.md
doc: update examples of deprecation timelines
2024-09-03 16:30:43 -07:00
developer.md
Add --use_cuda_nvcc flag to enable or disable compilation of CUDA code using NVCC.
2024-09-24 11:37:00 -07: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
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
faq.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
ffi.ipynb
Deprecate the vectorized argument to pure_callback and ffi_call.
2024-10-02 11:33:51 -04:00
ffi.md
Deprecate the vectorized argument to pure_callback and ffi_call.
2024-10-02 11:33:51 -04:00
glossary.rst
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
gpu_memory_allocation.rst
Change default memory allocation to 75% instead of 90%
2023-07-13 14:44:38 -07:00
gpu_performance_tips.md
Update flags gpu_performance_tips.md
2024-06-25 08:53:43 -07:00
gradient-checkpointing.md
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
index.rst
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
installation.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
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_internal_api.rst
docs: reorganize sections
2024-08-13 21:33:45 -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.array_api.rst
[array api] Finalize array API in jax.numpy & deprecate jax.experimental.array_api
2024-08-01 11:19:17 -07: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] GPUBlockSpec no longer accepts swizzle
2024-10-11 05:11:26 -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
[pallas] Added API docs for Triton and Mosaic GPU backends
2024-10-10 12:27:53 +01: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
DOC: document jax.experimental.sparse.linalg
2023-04-28 14:18:50 -07:00
jax.export.rst
[export] Add documentation for debugging and for ensuring compatibility.
2024-06-28 08:36:55 +03:00
jax.extend.ffi.rst
Export jaxlib.xla_client.register_custom_call_target as jax.extend.ffi.register_ffi_target.
2024-07-19 08:12:25 -04: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
Edit pycapsule docstring to provide a little bit more context
2024-06-07 13:07:03 -04: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
Clone RandomAlgorithm into lax.py, instead of using the version from XLA.
2024-10-11 11:03:15 -07:00
jax.lib.rst
Add register_custom_call_target to xla_client API docs.
2024-06-27 14:40:36 -04: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
Implement jax.numpy.spacing
2024-10-03 10:40:39 -07: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
Add jax.make_mesh to API docs.
2024-10-09 13:55:43 -04: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 module with aliases of jax.tree_util
2024-02-12 13:07:59 -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
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00:00
key-concepts.md
Fix Typos
2024-09-25 21:26:05 +05:30
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
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
profiling.md
docs: tidy up titles and headings
2024-08-13 11:53:57 -07: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
pre-commit: update hooks & pin using hashes
2024-08-27 15:23:13 -07: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
CI: pin sphinx to avoid build errors on 8.0
2024-08-06 09:16:41 -07:00
sharded-computation.ipynb
Merge pull request #23643 from TomAugspurger:fix/doc-typos
2024-09-16 07:07:26 -07:00
sharded-computation.md
Merge pull request #23643 from TomAugspurger:fix/doc-typos
2024-09-16 07:07:26 -07:00
stateful-computations.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
transfer_guard.rst
Add a documentation for the JAX transfer guard
2022-03-25 19:50:12 +00:00
tutorials.rst
Add/update JAX Advanced Tutorials docs, ToC structure
2024-09-20 23:06:54 +00: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
pre-commit: update hooks & pin using hashes
2024-08-27 15:23:13 -07:00
xla_flags.md
Add md page about XLA&XLA_FLAGS in JAX
2024-09-05 09:27:32 +00:00

README.md

To rebuild the documentation, see Update Documentation.