rocm_jax

_static

[Pallas] Add block-sparse kernel tutorial

2024-09-19 12:23:03 -07:00

_templates

…

_tutorials

Add/update JAX Advanced Tutorials docs, ToC structure

2024-09-20 23:06:54 +00:00

cuda_custom_call

Merge pull request #24345 from phu0ngng:cuda_custom_call

2024-10-17 13:57:15 -07:00

debugging

docs: runtime debugging tweaks

2024-08-28 11:04:14 -07:00

export

[export] Fix github links in the export documentation

2024-10-17 08:30:28 +01: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

Remove mesh_utils.create_device_mesh from docs

2024-10-19 15:48:42 -07:00

notebooks

Remove mesh_utils.create_device_mesh from docs

2024-10-19 15:48:42 -07:00

pallas

Update Pallas distributed tutorials with jax.make_mesh

2024-10-21 12:49:56 -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

Don't add mhlo.layout_mode = "default" since that is the default even in PJRT and will help reduce cruft in the IR

2024-10-11 14:54:32 -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

…

concurrency.rst

…

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

[ROCm] Fix README.md to update AMD JAX installation instructions

2024-10-16 17:15:32 -05: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

Make ffi_call return a callable

2024-10-21 12:16:57 -04:00

ffi.md

Make ffi_call return a callable

2024-10-21 12:16:57 -04:00

glossary.rst

Add/update JAX Advanced Tutorials docs, ToC structure

2024-09-20 23:06:54 +00:00

gpu_memory_allocation.rst

…

gpu_performance_tips.md

Update flags gpu_performance_tips.md

2024-06-25 08:53:43 -07:00

gradient-checkpointing.md

Update docs/gradient-checkpointing.md

2024-10-11 12:33:10 -04: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

…

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

[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

…

jax.experimental.compilation_cache.rst

Cleanup deprecated compilation cache APIs.

2024-01-12 22:44:48 -08:00

jax.experimental.custom_partitioning.rst

…

jax.experimental.jet.rst

…

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

…

jax.experimental.pallas.mosaic_gpu.rst

[pallas:mosaic_gpu] Added layout_cast

2024-10-17 08:08:05 -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

…

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

…

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

…

jax.image.rst

…

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

…

jax.profiler.rst

…

jax.random.rst

emphasize typed over legacy RNG keys in random module docs

2024-08-11 12:41:50 -07:00

jax.rst

Make committed a public property of jax.Array.

2024-10-15 19:46:10 -07: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

…

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

…

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

…

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

…

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