mirror of
https://github.com/ROCm/jax.git
synced 2025-04-19 05:16:06 +00:00

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.
24 lines
372 B
ReStructuredText
24 lines
372 B
ReStructuredText
jax package
|
|
===========
|
|
|
|
Subpackages
|
|
-----------
|
|
|
|
.. toctree::
|
|
:maxdepth: 1
|
|
|
|
jax.numpy
|
|
jax.scipy
|
|
jax.experimental
|
|
jax.lax
|
|
jax.ops
|
|
jax.random
|
|
|
|
Module contents
|
|
---------------
|
|
|
|
.. automodule:: jax
|
|
:members: jit, disable_jit, grad, value_and_grad, vmap, jacfwd, jacrev, hessian, jvp, vjp, make_jaxpr
|
|
:undoc-members:
|
|
:show-inheritance:
|