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

132 lines
2.2 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
jax.experimental
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.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 a jax.profiler module that exposes the new TensorFlow profiler in… (#2354) * Add a jax.profiler module that exposes the new TensorFlow profiler integration. * Fix documentation.
2020-03-05 12:07:57 -05:00
jax.profiler
Add Sphinx-generated reference documentation for JAX.
2019-01-15 20:14:19 -05:00
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::
jit
disable_jit
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 `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
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::
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::
vmap
jax.numpy.vectorize
Parallelization (:code:`pmap`)
------------------------------
.. autosummary::
pmap
devices
local_devices
Internal change PiperOrigin-RevId: 369345279
2021-04-19 18:22:38 -07:00
host_id
host_ids
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
Internal change PiperOrigin-RevId: 369345279
2021-04-19 18:22:38 -07:00
host_count
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
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
Document jax.device_put (#2366)
2020-03-05 14:45:01 -08:00
.. autofunction:: device_put
DOC: minor documentation & formatting fixes
2021-02-23 10:31:44 -08:00
.. autofunction:: device_put_replicated
.. autofunction:: device_put_sharded
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
.. autofunction:: default_backend
Add named_call autofunction description to docs
2020-11-30 14:23:52 +01:00
.. autofunction:: named_call
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:: grad
.. autofunction:: value_and_grad
.. autofunction:: jacfwd
.. autofunction:: jacrev
.. autofunction:: hessian
.. autofunction:: jvp
.. autofunction:: linearize
Fixes to jax.api docs
2020-11-08 22:22:36 +00:00
.. autofunction:: linear_transpose
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:: vjp
Improve docs for custom_jvp and custom_vjp Correct the custom_jvp docstring to include the defjvps instance method. Add the defjvp/defvjp instance methods to the sphinx doc.
2020-11-06 11:37:15 +00:00
.. autoclass:: custom_jvp
.. automethod:: defjvp
.. automethod:: defjvps
.. autoclass:: custom_vjp
.. automethod:: defvjp
include closure_convert in generated docs
2021-01-25 17:42:46 -08:00
.. autofunction:: closure_convert
Fix jax.checkpoint in API docs (#3980) On the index API doc page, it turns out functions need to be listed *twice* to appear.
2020-08-06 09:39:33 -07:00
.. autofunction:: checkpoint
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:: vmap
Minor fixes to docs related to jax.numpy.vectorize (#2278) - Show `numpy.jax.vectorize` explicitly in the JAX docs, rather than the original `numpy.vectorize. - Updated regex for identifying function signatures in NumPy. This now correctly parses `np.vectorize` and `np.einsum`. - Removed docs for `jax.experimental.vectorize`. There's still some good narrative content in the docstring but it should go somewhere else.
2020-02-23 10:10:39 -08:00
.. autofunction:: jax.numpy.vectorize
DOC: minor documentation & formatting fixes
2021-02-23 10:31:44 -08:00
:noindex:
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
Internal change PiperOrigin-RevId: 369345279
2021-04-19 18:22:38 -07:00
.. autofunction:: host_id
.. autofunction:: host_ids
Add jax.devices(), etc. to the docs.
2019-11-22 11:03:26 -08:00
.. autofunction:: device_count
.. autofunction:: local_device_count
Internal change PiperOrigin-RevId: 369345279
2021-04-19 18:22:38 -07:00
.. autofunction:: host_count
Reference in New Issue Copy Permalink