mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/docs/jax.rst

264 lines
4.0 KiB
ReStructuredText
Raw Permalink Normal View History

Fix index of API docs The current index of the API docs page seems to have broken links: when I click on "Automatic differentiation" for example, I get sent to the "JIT" section. This change fixes the links.
2019-08-20 23:13:15 -07:00
.. currentmodule:: jax
docs: reorganize sections * Create "extension guides" section * Sort developer notes into subsections * Move examples from advanced section into user guides * Reorder some listings, adjust some titles
2024-08-13 13:16:50 -07:00
Public API: ``jax`` package
===========================
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
Subpackages
-----------
.. toctree::
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
:maxdepth: 1
jax.numpy
jax.scipy
jax.lax
jax.random
jax.sharding
jax.debug
jax.dlpack
jax.distributed
DOC: add docs for jax.dtypes module
2023-02-13 15:07:02 -08:00
jax.dtypes
Move jex.ffi to jax.ffi.
2024-12-20 11:26:04 +00:00
jax.ffi
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.flatten_util
jax.image
jax.nn
jax.ops
jax.profiler
jax.stages
Add jax.test_util to public API docs
2025-01-23 16:04:35 -08:00
jax.test_util
Add jax.tree module with aliases of jax.tree_util
2024-02-12 13:07:59 -08:00
jax.tree
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.tree_util
Expose jax.typing & update docs
2023-02-13 15:53:08 -08:00
jax.typing
[export] Add documentation for jax.export
2024-05-31 15:14:09 +03:00
jax.export
initiate `jax.extend` via docs and top-level module set-up
2023-05-15 13:52:28 -07:00
jax.extend
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
jax.example_libraries
jax.experimental
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
DOC: fix docstring and add docs.
2021-09-27 09:48:27 -07:00
.. toctree::
:hidden:
jax.lib
DOC: make API doc titles more uniform
2023-01-18 10:59:42 -08:00
Configuration
-------------
.. autosummary::
:toctree: _autosummary
config
check_tracer_leaks
checking_leaks
debug_nans
debug_infs
default_device
default_matmul_precision
default_prng_impl
enable_checks
enable_custom_prng
enable_custom_vjp_by_custom_transpose
log_compiles
numpy_rank_promotion
transfer_guard
DOC: add transformations doc to HTML & reorganize contents
2021-03-08 16:25:04 -08:00
.. _jax-jit:
Add documentation to JAX modules.
2019-07-20 14:40:31 +01:00
Just-in-time compilation (:code:`jit`)
--------------------------------------
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
.. autosummary::
DOC: move functions in jax.html to their own pages
2022-02-04 13:09:24 -08:00
:toctree: _autosummary
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
jit
disable_jit
Copybara import of the project: -- 4fcdadbfb3f4c484fd4432203cf13b88782b9311 by Matthew Johnson <mattjj@google.com>: add jax.ensure_compile_time_eval to public api aka jax.core.eval_context COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/7987 from google:issue7535 4fcdadbfb3f4c484fd4432203cf13b88782b9311 PiperOrigin-RevId: 420928687
2022-01-10 20:57:56 -08:00
ensure_compile_time_eval
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
make_jaxpr
eval_shape
Document ShapeDtypeStruct
2023-03-21 13:53:20 -07:00
ShapeDtypeStruct
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
device_put
Add documentation for jax.device_get
2021-07-04 19:55:45 +08:00
device_get
Add `jax.default_backend()` which returns the default platform name. This can be useful when you need backend specific behaviour, e.g.: if jax.default_backend() == 'gpu': dataset = double_buffer(dataset) Or if you want to assert a given backend is the default: assert jax.default_backend() == 'tpu' I am a bit conflicted by the naming, "backend" is consistent with other APIs in JAX (e.g. jit, local_devices etc) which accept a "backend" string which is used to lookup an XLA backend by platform name.
2021-02-04 11:56:41 +00:00
default_backend
Add named_call public API. Move named_call_p to core.py from lax.py. Also move the translation rule to jax/interpreters/xla.py where the core_call translation rule is.
2020-11-04 21:01:42 -08:00
named_call
Update changelog for `named_scope` and adds it to the docs
2022-06-09 10:34:25 -07:00
named_scope
DOC: move functions in jax.html to their own pages
2022-02-04 13:09:24 -08:00
block_until_ready
Add `jax.copy_to_host_async(tree)`. A relatively common pattern I've observed is the following: ```python _, metrics = some_jax_function() with profiler.Trace('compute_metrics'): jax.block_until_ready(metrics) with profiler.Trace('copy_to_host'): metrics = jax.device_get(metrics) ``` We are missing an opportunity here to more eagerly begin the h2d copy of the metrics (e.g. overlap it with closing the "compute_metrics" context manager etc. The intention of `jax.copy_to_host_async(x)` is to make it simple to begin h2d transfers as early as possible. Adapting the above code: ```python _, metrics = some_jax_function() # Begin D2H copies as early as we can. jax.copy_to_host_async(metrics) with profiler.Trace('compute_metrics'): jax.block_until_ready(metrics) with profiler.Trace('copy_to_host'): metrics = jax.device_get(metrics) ``` PiperOrigin-RevId: 731626446
2025-02-27 01:21:39 -08:00
copy_to_host_async
Add jax.make_mesh to API docs.
2024-10-09 13:32:30 -04:00
make_mesh
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
DOC: add transformations doc to HTML & reorganize contents
2021-03-08 16:25:04 -08:00
.. _jax-grad:
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
Automatic differentiation
-------------------------
.. autosummary::
DOC: move functions in jax.html to their own pages
2022-02-04 13:09:24 -08:00
:toctree: _autosummary
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
grad
value_and_grad
Improve documentation for jax.jacobian
2024-09-26 05:09:47 -07:00
jacobian
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
jacfwd
jacrev
hessian
jvp
linearize
Add jax.linear_transpose (#3398) * Add jax.linear_transpose Co-authored-by: Matthew Johnson <mattjj@google.com> * add failing test for complex numbers * Add picky dtype check for linear_transpose * Lint fix * Allow truncating dtypes to match inputs in linear_transpose * Fix typo in shape check error * improve docstring * Don't support integer inputs; better docstring * fixup * Fix doctest Co-authored-by: Matthew Johnson <mattjj@google.com>
2020-09-16 20:29:19 -07:00
linear_transpose
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
vjp
add jax.custom_gradient to readthedocs
2023-11-20 13:26:52 -08:00
custom_gradient
include closure_convert in generated docs
2021-01-25 17:42:46 -08:00
closure_convert
add remat to top-level docs (#3554)
2020-06-25 07:26:26 -07:00
checkpoint
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
Add error message when using custom_vmap with reverse-mode AD, and add docstrings. The `custom_vmap` API is discussed in https://github.com/jax-ml/jax/issues/9073, and it remains somewhat experimental and incomplete, but it is sufficiently widely used that it seemed worth adding it to the docs. One specific pain point with `custom_vmap` is that it doesn't support reverse-mode autodiff, so I also added a better error message for this case. Before this change, using `grad` with a `custom_vmap` function would fail with an `assert` deep within the JAX internals. This now fails with a `NotImplementedError` that describes the problem. PiperOrigin-RevId: 704353963
2024-12-09 11:17:08 -08:00
Customization
-------------
Document methods of custom_jvp/custom_vjp
2024-08-14 15:37:20 -07:00
``custom_jvp``
~~~~~~~~~~~~~~
.. autosummary::
:toctree: _autosummary
custom_jvp
custom_jvp.defjvp
custom_jvp.defjvps
``custom_vjp``
~~~~~~~~~~~~~~
.. autosummary::
:toctree: _autosummary
custom_vjp
custom_vjp.defvjp
Add error message when using custom_vmap with reverse-mode AD, and add docstrings. The `custom_vmap` API is discussed in https://github.com/jax-ml/jax/issues/9073, and it remains somewhat experimental and incomplete, but it is sufficiently widely used that it seemed worth adding it to the docs. One specific pain point with `custom_vmap` is that it doesn't support reverse-mode autodiff, so I also added a better error message for this case. Before this change, using `grad` with a `custom_vmap` function would fail with an `assert` deep within the JAX internals. This now fails with a `NotImplementedError` that describes the problem. PiperOrigin-RevId: 704353963
2024-12-09 11:17:08 -08:00
``custom_batching``
~~~~~~~~~~~~~~~~~~~
.. autosummary::
:toctree: _autosummary
custom_batching.custom_vmap
custom_batching.custom_vmap.def_vmap
custom_batching.sequential_vmap
Add docstrings for jax.Array APIs `make_array_from_callback` and `make_array_from_single_device_arrays`. PiperOrigin-RevId: 487929688
2022-11-11 15:20:27 -08:00
jax.Array (:code:`jax.Array`)
-----------------------------
.. autosummary::
:toctree: _autosummary
Expose jax.typing & update docs
2023-02-13 15:53:08 -08:00
Array
Add docstrings for jax.Array APIs `make_array_from_callback` and `make_array_from_single_device_arrays`. PiperOrigin-RevId: 487929688
2022-11-11 15:20:27 -08:00
make_array_from_callback
make_array_from_single_device_arrays
Add `jax.make_array_from_process_local_data` to the docs. Preview: https://jax--21444.org.readthedocs.build/en/21444/_autosummary/jax.make_array_from_process_local_data.html#jax.make_array_from_process_local_data PiperOrigin-RevId: 637660828
2024-05-27 09:40:49 -07:00
make_array_from_process_local_data
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
document jax.Array methods and attributes
2024-08-16 06:37:19 -07:00
Array properties and methods
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. autosummary::
:toctree: _autosummary
Array.addressable_shards
Array.all
Array.any
Array.argmax
Array.argmin
Array.argpartition
Array.argsort
Array.astype
Array.at
Array.choose
Array.clip
Array.compress
Make committed a public property of jax.Array. Why? Because users need to know if an array is committed or not since JAX raises errors based on committedness of a jax.Array. JAX also makes decisions about dispatching based on committedness of a jax.Array. But the placement of such arrays on devices is an internal implementation detail. PiperOrigin-RevId: 686329828
2024-10-15 19:45:25 -07:00
Array.committed
document jax.Array methods and attributes
2024-08-16 06:37:19 -07:00
Array.conj
Array.conjugate
Array.copy
Array.copy_to_host_async
Array.cumprod
Array.cumsum
Array.device
Array.diagonal
Array.dot
Array.dtype
Array.flat
Array.flatten
Array.global_shards
Array.imag
Array.is_fully_addressable
Array.is_fully_replicated
Array.item
Array.itemsize
Array.max
Array.mean
Array.min
Array.nbytes
Array.ndim
Array.nonzero
Array.prod
Array.ptp
Array.ravel
Array.real
Array.repeat
Array.reshape
Array.round
Array.searchsorted
Array.shape
Array.sharding
Array.size
Array.sort
Array.squeeze
Array.std
Array.sum
Array.swapaxes
Array.take
Array.to_device
Array.trace
Array.transpose
Array.var
Array.view
Array.T
Array.mT
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
Vectorization (:code:`vmap`)
----------------------------
.. autosummary::
DOC: move functions in jax.html to their own pages
2022-02-04 13:09:24 -08:00
:toctree: _autosummary
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
vmap
DOC: move functions in jax.html to their own pages
2022-02-04 13:09:24 -08:00
numpy.vectorize
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
Parallelization (:code:`pmap`)
------------------------------
.. autosummary::
DOC: move functions in jax.html to their own pages
2022-02-04 13:09:24 -08:00
:toctree: _autosummary
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
pmap
devices
local_devices
Replace `host_id` with `process_index` terminology, take 2. We're switching to the new terminology to avoid confusion in cases where multiple jax processes are running on a single host, and each process has a unique process_index/host_id. This keeps aliases for the old `host_id` APIs for now, but these will eventually be removed. This was originally commited in b77ef5138b631378e6a8ceb8bafc94fe91239bae, but reverted in 14acd070c2afb11c81fc91f43790577cd48cbf67 due to Google-internal test failures from renaming the local_devices argument name. This change is identical except it also adds staging for the argument name change.
2021-04-20 17:56:41 -07:00
process_index
DOC: add a table of contents for top level API docs (#2946) This makes them easier to scan.
2020-05-04 12:37:29 -07:00
device_count
local_device_count
Replace `host_id` with `process_index` terminology, take 2. We're switching to the new terminology to avoid confusion in cases where multiple jax processes are running on a single host, and each process has a unique process_index/host_id. This keeps aliases for the old `host_id` APIs for now, but these will eventually be removed. This was originally commited in b77ef5138b631378e6a8ceb8bafc94fe91239bae, but reverted in 14acd070c2afb11c81fc91f43790577cd48cbf67 due to Google-internal test failures from renaming the local_devices argument name. This change is identical except it also adds staging for the argument name change.
2021-04-20 17:56:41 -07:00
process_count
Add a `jax.process_indices` function. The `jax.host_ids` function has be long deprecated, but the suggested alternative of `list(range(jax.process_count()))` relies on the current behavior that the list of process indices is always dense. In the future we may want to allow dynamic addition and removal of processes in which case `jax.process_count` and `jax.process_indices` would need to be updated, and it is useful for users to be able to use this forward-compatible interface. PiperOrigin-RevId: 662142636
2024-08-12 10:29:15 -07:00
process_indices
Add callbacks to docs
2022-08-24 14:04:47 -07:00
Callbacks
---------
.. autosummary::
:toctree: _autosummary
pure_callback
DOC: add docstrings for callback functions
2023-04-12 07:33:09 -07:00
experimental.io_callback
Add callbacks to docs
2022-08-24 14:04:47 -07:00
debug.callback
DOC: add docstrings for callback functions
2023-04-12 07:33:09 -07:00
debug.print
Add jax.print_environment_info()
2022-09-12 15:39:33 -07:00
Miscellaneous
-------------
.. autosummary::
:toctree: _autosummary
Export Device as jax.Device. Users are writing things like jax.lib.xla_client.Device in type annotations which is not a public API. Add a supported public name for the Device type.
2023-02-02 12:58:15 -05:00
Device
Add jax.print_environment_info()
2022-09-12 15:39:33 -07:00
print_environment_info
add docs for jax.clear_caches Co-authored-by: Roy Frostig <frostig@google.com>
2023-04-07 14:42:31 -07:00
live_arrays
clear_caches
Reference in New Issue Copy Permalink