Content - 31ad050df31e5c3e419380626986cb0bd51cbfb4 - 7e0026a/docs/jax.nn.rst

visit type:

Tip revision: fe0f12e00e53fd9557723c6398797cc4514992f1 authored by Adam Paszke on 18 September 2023, 17:53:10 UTC
[Pallas] Wire up cost estimates in Mosaic params

Tip revision: fe0f12e

jax.nn.rst


``jax.nn`` module
=================

.. 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

Browse the archive

https://github.com/google/jax