rocm_jax/docs/jax.rst

.. currentmodule:: jax

jax package
===========

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

.. toctree::
    :maxdepth: 1

    jax.numpy
    jax.scipy
    jax.experimental
    jax.lax
    jax.nn
    jax.ops
    jax.random
    jax.tree_util
    jax.dlpack

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

.. autofunction:: jit
.. autofunction:: disable_jit
.. autofunction:: xla_computation
.. autofunction:: make_jaxpr
.. autofunction:: eval_shape

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

.. autofunction:: grad
.. autofunction:: value_and_grad
.. autofunction:: jacfwd
.. autofunction:: jacrev
.. autofunction:: hessian
.. autofunction:: jvp
.. autofunction:: linearize
.. autofunction:: vjp
.. autofunction:: custom_transforms
.. autofunction:: defjvp
.. autofunction:: defjvp_all
.. autofunction:: defvjp
.. autofunction:: defvjp_all
.. autofunction:: custom_gradient


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

.. autofunction:: vmap


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

.. autofunction:: pmap
.. autofunction:: devices
.. autofunction:: local_devices
.. autofunction:: host_id
.. autofunction:: host_ids
.. autofunction:: device_count
.. autofunction:: local_device_count
.. autofunction:: host_count
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`

Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00			`jax package`
			`===========`

			`Subpackages`
			`-----------`

			`.. toctree::`
			`:maxdepth: 1`

			`jax.numpy`
			`jax.scipy`
			`jax.experimental`
			`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.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`
Update documentation and changelog to mention DLPack and array interface support. (#2134) 2020-01-31 11:15:04 -05:00			`jax.dlpack`
Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00
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
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			`.. autofunction:: jit`
			`.. autofunction:: disable_jit`
			`.. autofunction:: xla_computation`
			`.. autofunction:: make_jaxpr`
			`.. autofunction:: eval_shape`
Add documentation to JAX modules. 2019-07-20 14:40:31 +01:00
			`Automatic differentiation`
			`-------------------------`

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			`.. autofunction:: grad`
			`.. autofunction:: value_and_grad`
			`.. autofunction:: jacfwd`
			`.. autofunction:: jacrev`
			`.. autofunction:: hessian`
			`.. autofunction:: jvp`
			`.. autofunction:: linearize`
			`.. autofunction:: vjp`
			`.. autofunction:: custom_transforms`
			`.. autofunction:: defjvp`
			`.. autofunction:: defjvp_all`
			`.. autofunction:: defvjp`
			`.. autofunction:: defvjp_all`
			`.. autofunction:: custom_gradient`
Add documentation to JAX modules. 2019-07-20 14:40:31 +01:00

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

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			`.. autofunction:: vmap`
Add documentation to JAX modules. 2019-07-20 14:40:31 +01:00

			Parallelization (:code:`pmap`)
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			`------------------------------`
Add documentation to JAX modules. 2019-07-20 14:40:31 +01:00
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			`.. autofunction:: pmap`
Add jax.devices(), etc. to the docs. 2019-11-22 11:03:26 -08:00			`.. autofunction:: devices`
			`.. autofunction:: local_devices`
			`.. autofunction:: host_id`
			`.. autofunction:: host_ids`
			`.. autofunction:: device_count`
			`.. autofunction:: local_device_count`
			`.. autofunction:: host_count`