Revision 01194bd2fb1f7b6a46db88be14324f7acdb1723e authored by Sergei Lebedev on 16 May 2024, 18:27:56 UTC, committed by jax authors on 16 May 2024, 18:29:09 UTC
The callback APIs were migrated to use jax.Arrays for both inputs and outputs in JAX 0.4.27. PiperOrigin-RevId: 634473890
1 parent 380503b
jax.nn.rst
``jax.nn`` module
=================
.. currentmodule:: jax.nn
.. toctree::
:maxdepth: 1
jax.nn.initializers
.. automodule:: jax.nn
Activation functions
--------------------
.. autosummary::
:toctree: _autosummary
relu
relu6
sigmoid
softplus
sparse_plus
soft_sign
silu
swish
log_sigmoid
leaky_relu
hard_sigmoid
hard_silu
hard_swish
hard_tanh
elu
celu
selu
gelu
glu
squareplus
mish
Other functions
---------------
.. autosummary::
:toctree: _autosummary
softmax
log_softmax
logsumexp
standardize
one_hot
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...