2019-08-29 17:51:15 -07:00
|
|
|
|
2023-01-18 10:59:42 -08:00
|
|
|
``jax.nn`` module
|
2019-08-29 17:51:15 -07:00
|
|
|
=================
|
|
|
|
|
|
|
|
|
|
.. currentmodule:: jax.nn
|
|
|
|
|
|
|
|
|
|
.. toctree::
|
|
|
|
|
:maxdepth: 1
|
|
|
|
|
|
|
|
|
|
jax.nn.initializers
|
|
|
|
|
|
|
|
|
|
.. automodule:: jax.nn
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
Activation functions
|
2023-01-18 10:59:42 -08:00
|
|
|
--------------------
|
2019-08-29 17:51:15 -07:00
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
relu
|
2020-04-29 14:59:23 -04:00
|
|
|
relu6
|
2019-08-29 17:51:15 -07:00
|
|
|
sigmoid
|
|
|
|
|
softplus
|
2024-03-15 04:39:48 -07:00
|
|
|
sparse_plus
|
2024-05-08 12:23:01 +05:30
|
|
|
sparse_sigmoid
|
2019-08-29 17:51:15 -07:00
|
|
|
soft_sign
|
2020-07-06 18:08:16 -04:00
|
|
|
silu
|
2019-08-29 17:51:15 -07:00
|
|
|
swish
|
|
|
|
|
log_sigmoid
|
|
|
|
|
leaky_relu
|
2020-04-29 14:59:23 -04:00
|
|
|
hard_sigmoid
|
2020-07-06 18:08:16 -04:00
|
|
|
hard_silu
|
2020-04-29 14:59:23 -04:00
|
|
|
hard_swish
|
2019-08-29 17:51:15 -07:00
|
|
|
hard_tanh
|
|
|
|
|
elu
|
|
|
|
|
celu
|
|
|
|
|
selu
|
|
|
|
|
gelu
|
2019-12-02 14:21:10 -05:00
|
|
|
glu
|
2023-11-14 23:52:41 -05:00
|
|
|
squareplus
|
2024-04-03 16:37:07 -04:00
|
|
|
mish
|
2019-08-29 17:51:15 -07:00
|
|
|
|
|
|
|
|
Other functions
|
|
|
|
|
---------------
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
softmax
|
|
|
|
|
log_softmax
|
2020-10-15 11:10:56 -04:00
|
|
|
logsumexp
|
2022-10-20 21:25:37 -07:00
|
|
|
standardize
|
2020-04-06 12:39:28 -04:00
|
|
|
one_hot
|
