ContextQMD
Back to jax

``jax.nn`` module

docs/jax.nn.rst

0.3.25798 B
Original Source

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
sparse_sigmoid
soft_sign
silu
swish
log_sigmoid
leaky_relu
hard_sigmoid
hard_silu
hard_swish
hard_tanh
tanh
elu
celu
selu
gelu
glu
squareplus
mish
identity

Other functions

.. autosummary:: :toctree: _autosummary

softmax
log_softmax
logmeanexp
logsumexp
standardize
one_hot
dot_product_attention
scaled_matmul
get_scaled_dot_general_config
scaled_dot_general
log1mexp