rocm_jax/docs/jax.rst

.. currentmodule:: jax

Public API: jax package
=======================

Subpackages
-----------

.. toctree::
   :maxdepth: 1

   jax.numpy
   jax.scipy
   jax.lax
   jax.random
   jax.sharding
   jax.debug
   jax.dlpack
   jax.distributed
   jax.flatten_util
   jax.image
   jax.nn
   jax.ops
   jax.profiler
   jax.stages
   jax.tree_util
   jax.example_libraries
   jax.experimental

.. toctree::
   :hidden:

   jax.lib

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

.. _jax-jit:

Just-in-time compilation (:code:`jit`)
--------------------------------------

.. autosummary::
  :toctree: _autosummary

    jit
    disable_jit
    ensure_compile_time_eval
    xla_computation
    make_jaxpr
    eval_shape
    device_put
    device_put_replicated
    device_put_sharded
    device_get
    default_backend
    named_call
    named_scope
    block_until_ready

.. _jax-grad:

Automatic differentiation
-------------------------

.. autosummary::
  :toctree: _autosummary

    grad
    value_and_grad
    jacfwd
    jacrev
    hessian
    jvp
    linearize
    linear_transpose
    vjp
    custom_jvp
    custom_vjp
    closure_convert
    checkpoint

jax.Array (:code:`jax.Array`)
-----------------------------

.. autosummary::
  :toctree: _autosummary

    make_array_from_callback
    make_array_from_single_device_arrays

Vectorization (:code:`vmap`)
----------------------------

.. autosummary::
  :toctree: _autosummary

    vmap
    numpy.vectorize

Parallelization (:code:`pmap`)
------------------------------

.. autosummary::
  :toctree: _autosummary

    pmap
    devices
    local_devices
    process_index
    device_count
    local_device_count
    process_count

Callbacks
---------

.. autosummary::
  :toctree: _autosummary

    pure_callback
    debug.callback

Miscellaneous
-------------

.. autosummary::
  :toctree: _autosummary

    print_environment_info
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::`
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`
			`jax.flatten_util`
			`jax.image`
			`jax.nn`
			`jax.ops`
			`jax.profiler`
			`jax.stages`
			`jax.tree_util`
			`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			`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
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`

			`make_array_from_callback`
			`make_array_from_single_device_arrays`
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`
Add jax.print_environment_info() 2022-09-12 15:39:33 -07:00
			`Miscellaneous`
			`-------------`

			`.. autosummary::`
			`:toctree: _autosummary`

			`print_environment_info`