Revision 66b182b735503db3a7bbc956b2fc1d30d897505b authored by Matthew Johnson on 06 December 2019, 15:44:12 UTC, committed by Matthew Johnson on 06 December 2019, 19:12:59 UTC
1 parent 1da950d
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
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
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...