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

111 lines
1.5 KiB
ReStructuredText
Raw 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
Added the first draft of the Jaxpr documentation. This replaces the previous Google Doc version, and is now updated with the latest changes in Jaxpr.
2020-02-10 11:40:05 +01:00
Public API: jax package
=======================
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
Subpackages
-----------
.. toctree::
:maxdepth: 1
jax.numpy
jax.scipy
Add jax.distributed and jax.dlpack to the docs. Reorder the doc modules into something closer to alphabetical order. Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs.
2022-02-17 16:05:33 -05:00
jax.config
Add debugging tools under `jax.debug` and documentation Co-authored-by: Matthew Johnson <mattjj@google.com> Co-authored-by: Lena Martens <lenamartens@google.com>
2022-07-26 14:47:36 -07:00
jax.debug
Add jax.distributed and jax.dlpack to the docs. Reorder the doc modules into something closer to alphabetical order. Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs.
2022-02-17 16:05:33 -05:00
jax.dlpack
jax.distributed
[JAX] move example libraries from `jax.experimental` into `jax.example_libraries` The `jax.experimental.stax` and `jax.experimental.optimizers` modules are standalone examples libraries. By contrast, the remaining modules in `jax.experimental` are experimental features of the JAX core system. This change moves the two example libraries, and the README that describes them, to `jax.example_libraries` to reflect this distinction. PiperOrigin-RevId: 404405186
2021-10-19 17:30:16 -07:00
jax.example_libraries
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
jax.experimental
Add jax.distributed and jax.dlpack to the docs. Reorder the doc modules into something closer to alphabetical order. Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs.
2022-02-17 16:05:33 -05:00
jax.flatten_util
Add jax.image.resize. (#3703) * Add jax.image.resize. This is a port of `tf.image.resize()` and the `ScaleAndTranslate` operator. While I don't expect this implementation to be particularly fast, it is a useful generic implementation to which we can add optimized special cases as the need arises.
2020-07-10 09:57:59 -04:00
jax.image
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
jax.lax
add docs for jax.nn
2019-08-29 17:51:15 -07:00
jax.nn
Add new functions `jax.ops.index_add` and `jax.ops.index_update` for NumPy-style indexed updates. Create a new library `jax.ops` for user-facing ops that don't exist in NumPy or SciPy. Progress on issue #101. Fixes #122. Reenable some disabled TPU indexing tests that now pass.
2019-02-22 07:55:36 -05:00
jax.ops
Add jax.distributed and jax.dlpack to the docs. Reorder the doc modules into something closer to alphabetical order. Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs.
2022-02-17 16:05:33 -05:00
jax.profiler
add jax.random entry to jax.rst
2019-02-13 19:31:41 -08:00
jax.random
Expose jax.tree_util in the JAX docs.
2019-05-14 21:00:27 -04:00
jax.tree_util
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: 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
xla_computation
make_jaxpr
eval_shape
device_put
DOC: minor documentation & formatting fixes
2021-02-23 10:31:44 -08:00
device_put_replicated
device_put_sharded
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
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
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
custom_jvp
custom_vjp
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
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 callbacks to docs
2022-08-24 14:04:47 -07:00
Callbacks
---------
.. autosummary::
:toctree: _autosummary
pure_callback
debug.callback
Reference in New Issue Copy Permalink