rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00

History

Peter Hawkins b13733c13f Update JAX dependencies, extras, and documentation for plugins.

* Make jaxlib a direct dependency of jax.
* Remove mentions of monolithic CUDA installations from the JAX documentation.
* Drop the cuda12_pip extra and the cudnn version specific extras.
* Add a with_cuda extra to the jax-cuda12-plugin package, use it in jax's setup.py. This allows us to specify cuda extras in one place.
* Make a few small doc improvements.

2024-06-13 11:36:23 -04:00

_static

Add Distributed data loading doc

2024-05-16 23:12:13 +00:00

_templates

Add the JAX logo to sphinx docs

2019-08-05 12:34:29 -07:00

_tutorials

Clarified the type of the inputs to callback APIs

2024-05-16 11:29:09 -07:00

debugging

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

export

[export] Rename some API entry points

2024-06-13 06:44:13 +02:00

gpu_ops

Current status + build script fixes

2024-04-08 04:48:28 -07:00

jep

[docs] Enable doctest for .md documentation files

2024-06-03 11:27:34 +03:00

notebooks

[docs] Enable doctest for .md documentation files

2024-06-03 11:27:34 +03:00

pallas

Fix minor typo in Pallas docs.

2024-04-15 16:20:42 -07:00

sphinxext

DOC Add a noop implementation of the legacy directive to fix the build

2024-01-22 17:59:54 +00:00

advanced_guide.rst

Add Distributed data loading doc

2024-05-16 23:12:13 +00:00

aot.md

[export] Add documentation for jax.export

2024-06-12 19:44:47 +02:00

api_compatibility.md

add jax.prng to uncovered modules list in API policy

2023-07-18 14:13:25 -07:00

async_dispatch.rst

update documentation notes to new-style typed keys

2024-03-07 12:40:09 -08:00

autodidax.ipynb

Fix duplicate word occurrences

2023-12-19 06:15:30 +01:00

autodidax.md

Updated all the pre-commit hooks versions.

2024-04-08 00:59:02 +05:30

autodidax.py

Updated all the pre-commit hooks versions.

2024-04-08 00:59:02 +05:30

automatic-differentiation.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

automatic-vectorization.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

beginner_guide.rst

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

build_custom_gpu.sh

Current status + build script fixes

2024-04-08 04:48:28 -07:00

building_on_jax.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -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 Distributed data loading doc

2024-05-16 23:12:13 +00:00

contributing.md

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05:00

contributor_guide.rst

Add doc about regression investigation.

2023-09-29 09:51:49 -04:00

Custom_Operation_for_GPUs.md

fix typo in Custom_Operation_for_GPUs.md

2024-06-05 10:03:09 +03:00

Custom_Operation_for_GPUs.py

Current status + build script fixes

2024-04-08 04:48:28 -07:00

debugging.md

Update line127 error, debugging.md

2024-05-09 22:44:27 -04:00

deprecation.md

API deprecation doc: explicitly list some dates

2024-04-25 09:21:16 -07:00

developer.md

Use hermetic Python in JAX, see "Managing hermetic Python" in developer.md for details

2024-05-15 18:20:56 -07:00

device_memory_profiling.md

update documentation notes to new-style typed keys

2024-03-07 12:40:09 -08:00

distributed_data_loading.md

Add Distributed data loading doc

2024-05-16 23:12:13 +00:00

errors.rst

[key reuse] define KeyReuseError in jax.errors

2024-03-08 10:59:06 -08:00

faq.rst

Fix parenthesis in "Gradients contain NaN where using where"

2024-05-31 12:03:39 +01:00

glossary.rst

DOC: Fix :ref: format

2024-04-11 12:54:37 +02: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

Add in the tutorial the idea to test 1 process per node and 1 process per GPU.

2024-06-07 10:00:04 -04:00

index.rst

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

installation.md

Update JAX dependencies, extras, and documentation for plugins.

2024-06-13 11:36:23 -04:00

investigating_a_regression.md

Add a comment

2023-11-02 11:37:17 -07:00

jax_array_migration.md

Fixes broken examples, and (invalid) comment for PartitionSpec

2023-03-17 16:09:45 -07:00

jax_internal_api.rst

Add jax.errors submodule & error troubleshooting docs

2021-03-03 12:39:12 -08: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

Add jax.experimental.array_api interface

2023-11-16 14:21:04 -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.host_callback.rst

DOC: make API doc titles more uniform

2023-01-18 10:59:42 -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.maps.rst

Merge pull request #14248 from jakevdp:dead-code

2023-02-01 21:25:46 +00: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.pjit.rst

DOC: make API doc titles more uniform

2023-01-18 10:59:42 -08:00

jax.experimental.rst

Add doc for jax.experimental.serialize_executable module

2024-05-14 19:11:37 -04: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] Rename some API entry points

2024-06-13 06:44:13 +02:00

jax.extend.ffi.rst

Edit pycapsule docstring to provide a little bit more context

2024-06-07 13:07:03 -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

Finalize deprecations of several APIs

2024-05-14 10:40:40 -07:00

jax.lib.rst

Export Device as jax.Device.

2023-02-02 12:58:15 -05: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

Add sparse_sigmoid to jax.nn docs

2024-05-08 12:23:01 +05:30

jax.numpy.rst

Added trace alias to jnp.linalg

2024-05-29 22:28:44 +00: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

[key reuse] add jax.random.clone

2024-03-08 09:06:00 -08:00

jax.rst

[export] Add documentation for jax.export

2024-06-12 19:44:47 +02:00

jax.scipy.rst

Add RegularGridInterpolator to generated API docs

2024-05-17 11:24:35 -04:00

jax.sharding.rst

Deprecate XLACompatibleSharding in favor of jax.sharding.Sharding.

2024-06-05 09:07:27 -07:00

jax.stages.rst

Add a .cost_analysis() on lowered but uncompiled computations.

2023-02-06 12:57:57 -08:00

jax.tree_util.rst

Document jax.tree.* directly

2024-06-12 14:01:27 -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.rst

Minor doc fix to jaxpr.rst

2024-05-07 17:36:19 +01:00

jit-compilation.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

key-concepts.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

multi_process.md

Add in the tutorial the idea to test 1 process per node and 1 process per GPU.

2024-06-07 10:00:04 -04:00

notes.rst

DOC: re-organize notes section

2023-01-30 11:57:06 -08:00

persistent_compilation_cache.md

Jax persistent compilation cache user guide.

2024-04-09 11:48:33 -07:00

profiling.md

Remove outdated documentation section,

2024-03-18 18:25:07 -07:00

pytrees.md

Fix doc typos

2024-03-13 14:48:31 +05:30

quickstart.md

Update JAX dependencies, extras, and documentation for plugins.

2024-06-13 11:36:23 -04:00

random-numbers.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -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

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

sharded-computation.ipynb

Upgrade JAX Parallelism Sharded Computation 101 doc

2024-05-10 21:24:16 +00:00

sharded-computation.md

Upgrade JAX Parallelism Sharded Computation 101 doc

2024-05-10 21:24:16 +00:00

stateful-computations.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

transfer_guard.rst

Add a documentation for the JAX transfer guard

2022-03-25 19:50:12 +00:00

tutorials.rst

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

type_promotion.rst

[docs] Enable doctest for .md documentation files

2024-06-03 11:27:34 +03:00

user_guides.rst

[export] Add documentation for jax.export

2024-06-12 19:44:47 +02:00

working-with-pytrees.md

DOC: replace old tutorials with new content

2024-05-03 12:20:06 -07:00

README.md

To rebuild the documentation, see Update Documentation.