https://github.com/google/jax
Raw File
Tip revision: 0118f8d568d4d94665e01f284f63f0a5d9865f5e authored by Yash Katariya on 08 December 2022, 00:01:43 UTC
Prepare for jax and jaxlib 0.4.0 release
Tip revision: 0118f8d
jax.nn.rst

jax.nn package
=================

.. currentmodule:: jax.nn

.. toctree::
    :maxdepth: 1

    jax.nn.initializers

.. automodule:: jax.nn


Activation functions
------------------------

.. autosummary::
  :toctree: _autosummary

    relu
    relu6
    sigmoid
    softplus
    soft_sign
    silu
    swish
    log_sigmoid
    leaky_relu
    hard_sigmoid
    hard_silu
    hard_swish
    hard_tanh
    elu
    celu
    selu
    gelu
    glu

Other functions
---------------

.. autosummary::
  :toctree: _autosummary

    softmax
    log_softmax
    logsumexp
    standardize
    one_hot
back to top