mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 20:06:05 +00:00
43 lines
490 B
ReStructuredText
43 lines
490 B
ReStructuredText
|
|
||
|
|
jax.nn package
|
||
|
|
=================
|
||
|
|
|
||
|
|
.. currentmodule:: jax.nn
|
||
|
|
|
||
|
|
.. toctree::
|
||
|
|
:maxdepth: 1
|
||
|
|
|
||
|
|
jax.nn.initializers
|
||
|
|
|
||
|
|
.. automodule:: jax.nn
|
||
|
|
|
||
|
|
|
||
|
|
Activation functions
|
||
|
|
------------------------
|
||
|
|
|
||
|
|
.. autosummary::
|
||
|
|
:toctree: _autosummary
|
||
|
|
|
||
|
|
relu
|
||
|
|
sigmoid
|
||
|
|
softplus
|
||
|
|
soft_sign
|
||
|
|
swish
|
||
|
|
log_sigmoid
|
||
|
|
leaky_relu
|
||
|
|
hard_tanh
|
||
|
|
elu
|
||
|
|
celu
|
||
|
|
selu
|
||
|
|
gelu
|
||
|
|
|
||
|
|
Other functions
|
||
|
|
---------------
|
||
|
|
|
||
|
|
.. autosummary::
|
||
|
|
:toctree: _autosummary
|
||
|
|
|
||
|
|
softmax
|
||
|
|
log_softmax
|
||
|
|
normalize
|