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

Merge pull request #1217 from shoyer/doc-index-page

Browse Source 
Fix index of API docs
This commit is contained in:
Matthew Johnson 2019-08-21 16:46:29 -07:00 committed by GitHub
commit 1e375dafe0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 24 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
docs/jax.rst
View File

@ -1,3 +1,5 @@
.. currentmodule:: jax
jax package
===========
@ -18,33 +20,38 @@ Subpackages
Just-in-time compilation (:code:`jit`)
--------------------------------------
.. automodule:: jax
:members: jit, disable_jit, xla_computation, make_jaxpr, eval_shape
:undoc-members:
:show-inheritance:
.. autofunction:: jit
.. autofunction:: disable_jit
.. autofunction:: xla_computation
.. autofunction:: make_jaxpr
.. autofunction:: eval_shape
Automatic differentiation
-------------------------
.. automodule:: jax
:members: grad, value_and_grad, jacfwd, jacrev, hessian, jvp, linearize, vjp, custom_transforms, defjvp, defjvp_all, defvjp, defvjp_all, custom_gradient
:undoc-members:
:show-inheritance:
.. 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`)
----------------------------
.. automodule:: jax
:members: vmap
:undoc-members:
:show-inheritance:
.. autofunction:: vmap
Parallelization (:code:`pmap`)
----------------------------
------------------------------
.. automodule:: jax
:members: pmap
:undoc-members:
:show-inheritance:
.. autofunction:: pmap