``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 soft_sign silu swish log_sigmoid leaky_relu hard_sigmoid hard_silu hard_swish hard_tanh elu celu selu gelu glu squareplus Other functions --------------- .. autosummary:: :toctree: _autosummary softmax log_softmax logsumexp standardize one_hot