mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 11:56:07 +00:00
00140f07e2
* Add jax.numpy.vectorize This is basically a non-experimental version of the machinery in `jax.experimental.vectorize`, except: - It adds the `excluded` argument from NumPy, which works just like `static_argnums` in `jax.jit`. - It doesn't include the `axis` argument yet (which NumPy doesn't have). Eventually we might want want to consolidate the specification of signatures with signatures used by shape-checking machinery, but it's nice to emulate NumPy's existing interface, and this is already useful (e.g., for writing vectorized linear algebra routines). * Add deprecation warning to jax.experimental.vectorize * improve implementation
67 lines
1.3 KiB
ReStructuredText
67 lines
1.3 KiB
ReStructuredText
.. 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
|
|
.. autofunction:: numpy.vectorize
|
|
|
|
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
|