Elementwise and Activation Function Modules

PReLU	Element-wise Parametric Rectified Linear Unit (PReLU) activation function.
ReLU	Elementwise activation function for jax.nn.relu.
ReLU6	Elementwise activation function for jax.nn.relu6.
Sigmoid	Elementwise activation function for jax.nn.sigmoid.
Softplus	Elementwise activation function for jax.nn.softplus.
SparsePlus	Elementwise activation function for jax.nn.sparse_plus.
SparseSigmoid	Elementwise activation function for jax.nn.sparse_sigmoid.
SoftSign	Elementwise activation function for jax.nn.soft_sign.
SiLU	Elementwise activation function for jax.nn.silu.
Swish	Elementwise activation function for jax.nn.swish.
LogSigmoid	Elementwise activation function for jax.nn.log_sigmoid.
LeakyReLU	Elementwise activation function for jax.nn.leaky_relu.
HardSigmoid	Elementwise activation function for jax.nn.hard_sigmoid.
HardSiLU	Elementwise activation function for jax.nn.hard_silu.
HardSwish	Elementwise activation function for jax.nn.hard_swish.
HardTanh	Elementwise activation function for jax.nn.hard_tanh.
ELU	Elementwise activation function for jax.nn.elu.
CELU	Elementwise activation function for jax.nn.celu.
SELU	Elementwise activation function for jax.nn.selu.
GELU	Elementwise activation function for jax.nn.gelu.
GLU	Elementwise activation function for jax.nn.glu.
SquarePlus	Elementwise activation function for jax.nn.squareplus.
Mish	Elementwise activation function for jax.nn.mish.
Identity	Elementwise activation function for jax.nn.identity.
Softmax	Elementwise activation function for jax.nn.softmax.
LogSoftmax	Elementwise activation function for jax.nn.log_softmax.
LogSumExp	Elementwise activation function for jax.nn.logsumexp.
Standardize	Elementwise activation function for jax.nn.standardize.
OneHot	Elementwise activation function for jax.nn.one_hot.
Abs	Elementwise activation function for jax.numpy.abs.
Absolute	Elementwise activation function for jax.numpy.absolute.
ACos	Elementwise activation function for jax.numpy.acos.
ACosh	Elementwise activation function for jax.numpy.acosh.
AMax	Elementwise activation function for jax.numpy.amax.
AMin	Elementwise activation function for jax.numpy.amin.
Angle	Elementwise activation function for jax.numpy.angle.
ArcCos	Elementwise activation function for jax.numpy.arccos.
ArcCosh	Elementwise activation function for jax.numpy.arccosh.
ArcSin	Elementwise activation function for jax.numpy.arcsin.
ArcSinh	Elementwise activation function for jax.numpy.arcsinh.
ArcTan	Elementwise activation function for jax.numpy.arctan.
ArcTan2	Elementwise activation function for jax.numpy.arctan2.
ArcTanh	Elementwise activation function for jax.numpy.arctanh.
ASin	Elementwise activation function for jax.numpy.asin.
ASinh	Elementwise activation function for jax.numpy.asinh.
ATan	Elementwise activation function for jax.numpy.atan.
ATanh	Elementwise activation function for jax.numpy.atanh.
Cbrt	Elementwise activation function for jax.numpy.cbrt.
Ceil	Elementwise activation function for jax.numpy.ceil.
Clip	Elementwise activation function for jax.numpy.clip.
Conj	Elementwise activation function for jax.numpy.conj.
Conjugate	Elementwise activation function for jax.numpy.conjugate.
Cos	Elementwise activation function for jax.numpy.cos.
Cosh	Elementwise activation function for jax.numpy.cosh.
Deg2Rad	Elementwise activation function for jax.numpy.deg2rad.
Degrees	Elementwise activation function for jax.numpy.degrees.
Exp	Elementwise activation function for jax.numpy.exp.
Exp2	Elementwise activation function for jax.numpy.exp2.
Expm1	Elementwise activation function for jax.numpy.expm1.
FAbs	Elementwise activation function for jax.numpy.fabs.
Fix	Elementwise activation function for jax.numpy.fix.
FloatPower	Elementwise activation function for jax.numpy.float_power.
Floor	Elementwise activation function for jax.numpy.floor.
FloorDivide	Elementwise activation function for jax.numpy.floor_divide.
FrExp	Elementwise activation function for jax.numpy.frexp.
I0	Elementwise activation function for jax.numpy.i0.
Imag	Elementwise activation function for jax.numpy.imag.
Invert	Elementwise activation function for jax.numpy.invert.
LDExp	Elementwise activation function for jax.numpy.ldexp.
Log	Elementwise activation function for jax.numpy.log.
Log10	Elementwise activation function for jax.numpy.log10.
Log1p	Elementwise activation function for jax.numpy.log1p.
Log2	Elementwise activation function for jax.numpy.log2.
NaNToNum	Elementwise activation function for jax.numpy.nan_to_num.
NanToNum	Elementwise activation function for jax.numpy.nan_to_num.
NextAfter	Elementwise activation function for jax.numpy.nextafter.
Packbits	Elementwise activation function for jax.numpy.packbits.
Piecewise	Elementwise activation function for jax.numpy.piecewise.
Positive	Elementwise activation function for jax.numpy.positive.
Pow	Elementwise activation function for jax.numpy.pow.
Power	Elementwise activation function for jax.numpy.power.
Rad2Deg	Elementwise activation function for jax.numpy.rad2deg.
Radians	Elementwise activation function for jax.numpy.radians.
Real	Elementwise activation function for jax.numpy.real.
Reciprocal	Elementwise activation function for jax.numpy.reciprocal.
RInt	Elementwise activation function for jax.numpy.rint.
Round	Elementwise activation function for jax.numpy.round.
Sign	Elementwise activation function for jax.numpy.sign.
Signbit	Elementwise activation function for jax.numpy.signbit.
Sin	Elementwise activation function for jax.numpy.sin.
Sinc	Elementwise activation function for jax.numpy.sinc.
Sinh	Elementwise activation function for jax.numpy.sinh.
Sqrt	Elementwise activation function for jax.numpy.sqrt.
Square	Elementwise activation function for jax.numpy.square.
Tan	Elementwise activation function for jax.numpy.tan.
Tanh	Elementwise activation function for jax.numpy.tanh.
Trunc	Elementwise activation function for jax.numpy.trunc.
Unpackbits	Elementwise activation function for jax.numpy.unpackbits.