jax.nn package jax.nn jax.nn.initializers jax.nn Activation functions 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 softmax log_softmax normalize one_hot