rocm_jax/docs/jax.nn.rst
2020-04-06 12:39:28 -04:00

45 lines
510 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
glu
Other functions
---------------
.. autosummary::
:toctree: _autosummary
softmax
log_softmax
normalize
one_hot