brainstate 0.1.10__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/nn/_activations.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -66,41 +66,49 @@ __all__ = [
 
 
 def tanh(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Hyperbolic tangent activation function.
+    r"""
+    Hyperbolic tangent activation function.
 
     Computes the element-wise function:
 
     .. math::
       \mathrm{tanh}(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}
 
-    Args:
-      x : input array
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    Returns:
-      An array.
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
     """
     return u.math.tanh(x)
 
 
 def softmin(x, axis=-1):
     r"""
-    Applies the Softmin function to an n-dimensional input Tensor
-    rescaling them so that the elements of the n-dimensional output Tensor
-    lie in the range `[0, 1]` and sum to 1.
+    Softmin activation function.
 
-    Softmin is defined as:
+    Applies the Softmin function to an n-dimensional input tensor, rescaling elements
+    so that they lie in the range [0, 1] and sum to 1 along the specified axis.
 
     .. math::
         \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)}
 
-    Shape:
-        - Input: :math:`(*)` where `*` means, any number of additional
-          dimensions
-        - Output: :math:`(*)`, same shape as the input
-
-    Args:
-        axis (int): A dimension along which Softmin will be computed (so every slice
-            along dim will sum to 1).
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array of any shape.
+    axis : int, optional
+        The axis along which Softmin will be computed. Every slice along this
+        dimension will sum to 1. Default is -1.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
     """
     unnormalized = u.math.exp(-x)
     return unnormalized / unnormalized.sum(axis, keepdims=True)
@@ -108,22 +116,36 @@ def softmin(x, axis=-1):
 
 def tanh_shrink(x):
     r"""
+    Tanh shrink activation function.
+
     Applies the element-wise function:
 
     .. math::
         \text{Tanhshrink}(x) = x - \tanh(x)
+
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
     """
     return x - u.math.tanh(x)
 
 
 def prelu(x, a=0.25):
     r"""
+    Parametric Rectified Linear Unit activation function.
+
     Applies the element-wise function:
 
     .. math::
         \text{PReLU}(x) = \max(0,x) + a * \min(0,x)
 
-    or
+    or equivalently:
 
     .. math::
         \text{PReLU}(x) =
@@ -132,16 +154,32 @@ def prelu(x, a=0.25):
         ax, & \text{ otherwise }
         \end{cases}
 
-    Here :math:`a` is a learnable parameter. When called without arguments, `nn.PReLU()` uses a single
-    parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,
-    a separate :math:`a` is used for each input channel.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    a : float or ArrayLike, optional
+        The negative slope coefficient. Can be a learnable parameter.
+        Default is 0.25.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
+
+    Notes
+    -----
+    When used in neural network layers, :math:`a` can be a learnable parameter
+    that is optimized during training.
     """
     return u.math.where(x >= 0., x, a * x)
 
 
 def soft_shrink(x, lambd=0.5):
     r"""
-    Applies the soft shrinkage function elementwise:
+    Soft shrinkage activation function.
+
+    Applies the soft shrinkage function element-wise:
 
     .. math::
         \text{SoftShrinkage}(x) =
@@ -151,43 +189,60 @@ def soft_shrink(x, lambd=0.5):
         0, & \text{ otherwise }
         \end{cases}
 
-    Args:
-        lambd: the :math:`\lambda` (must be no less than zero) value for the Softshrink formulation. Default: 0.5
-
-    Shape:
-        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
-        - Output: :math:`(*)`, same shape as the input.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array of any shape.
+    lambd : float, optional
+        The :math:`\lambda` value for the soft shrinkage formulation.
+        Must be non-negative. Default is 0.5.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
     """
-    return u.math.where(x > lambd,
-                        x - lambd,
-                        u.math.where(x < -lambd,
-                                     x + lambd,
-                                     u.Quantity(0., unit=u.get_unit(lambd))))
+    return u.math.where(
+        x > lambd,
+        x - lambd,
+        u.math.where(
+            x < -lambd,
+            x + lambd,
+            u.Quantity(0., unit=u.get_unit(lambd))
+        )
+    )
 
 
 def mish(x):
-    r"""Applies the Mish function, element-wise.
+    r"""
+    Mish activation function.
 
-    Mish: A Self Regularized Non-Monotonic Neural Activation Function.
+    Mish is a self-regularized non-monotonic activation function.
 
     .. math::
         \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
 
-    .. note::
-        See `Mish: A Self Regularized Non-Monotonic Neural Activation Function <https://arxiv.org/abs/1908.08681>`_
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array of any shape.
 
-    Shape:
-        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
-        - Output: :math:`(*)`, same shape as the input.
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
+
+    References
+    ----------
+    .. [1] Misra, D. (2019). "Mish: A Self Regularized Non-Monotonic Activation Function."
+           arXiv:1908.08681
     """
     return x * u.math.tanh(softplus(x))
 
 
 def rrelu(x, lower=0.125, upper=0.3333333333333333):
-    r"""Applies the randomized leaky rectified liner unit function, element-wise,
-    as described in the paper:
-
-    `Empirical Evaluation of Rectified Activations in Convolutional Network`_.
+    r"""
+    Randomized Leaky Rectified Linear Unit activation function.
 
     The function is defined as:
 
@@ -201,27 +256,36 @@ def rrelu(x, lower=0.125, upper=0.3333333333333333):
     where :math:`a` is randomly sampled from uniform distribution
     :math:`\mathcal{U}(\text{lower}, \text{upper})`.
 
-     See: https://arxiv.org/pdf/1505.00853.pdf
-
-    Args:
-        lower: lower bound of the uniform distribution. Default: :math:`\frac{1}{8}`
-        upper: upper bound of the uniform distribution. Default: :math:`\frac{1}{3}`
-
-    Shape:
-        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
-        - Output: :math:`(*)`, same shape as the input.
-
-    .. _`Empirical Evaluation of Rectified Activations in Convolutional Network`:
-        https://arxiv.org/abs/1505.00853
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array of any shape.
+    lower : float, optional
+        Lower bound of the uniform distribution for sampling the negative slope.
+        Default is 1/8.
+    upper : float, optional
+        Upper bound of the uniform distribution for sampling the negative slope.
+        Default is 1/3.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
+
+    References
+    ----------
+    .. [1] Xu, B., et al. (2015). "Empirical Evaluation of Rectified Activations
+           in Convolutional Network." arXiv:1505.00853
     """
     a = random.uniform(lower, upper, size=u.math.shape(x), dtype=x.dtype)
     return u.math.where(u.get_mantissa(x) >= 0., x, a * x)
 
 
 def hard_shrink(x, lambd=0.5):
-    r"""Applies the Hard Shrinkage (Hardshrink) function element-wise.
+    r"""
+    Hard shrinkage activation function.
 
-    Hardshrink is defined as:
+    Applies the hard shrinkage function element-wise:
 
     .. math::
         \text{HardShrink}(x) =
@@ -231,139 +295,202 @@ def hard_shrink(x, lambd=0.5):
         0, & \text{ otherwise }
         \end{cases}
 
-    Args:
-        lambd: the :math:`\lambda` value for the Hardshrink formulation. Default: 0.5
-
-    Shape:
-        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
-        - Output: :math:`(*)`, same shape as the input.
-
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array of any shape.
+    lambd : float, optional
+        The :math:`\lambda` threshold value for the hard shrinkage formulation.
+        Default is 0.5.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Output array with the same shape as the input.
     """
-    return u.math.where(x > lambd,
-                        x,
-                        u.math.where(x < -lambd,
-                                     x,
-                                     u.Quantity(0., unit=u.get_unit(x))))
+    return u.math.where(
+        x > lambd,
+        x,
+        u.math.where(
+            x < -lambd,
+            x,
+            u.Quantity(0., unit=u.get_unit(x))
+        )
+    )
 
 
 def relu(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Rectified linear unit activation function.
+    r"""
+    Rectified Linear Unit activation function.
 
     Computes the element-wise function:
 
     .. math::
       \mathrm{relu}(x) = \max(x, 0)
 
-    except under differentiation, we take:
+    Under differentiation, we take:
 
     .. math::
       \nabla \mathrm{relu}(0) = 0
 
-    For more information see
-    `Numerical influence of ReLU’(0) on backpropagation
-    <https://openreview.net/forum?id=urrcVI-_jRm>`_.
-
-    Args:
-      x : input array
-
-    Returns:
-      An array.
-
-    Example:
-      >>> jax.nn.relu(jax.numpy.array([-2., -1., -0.5, 0, 0.5, 1., 2.]))
-      Array([0. , 0. , 0. , 0. , 0.5, 1. , 2. ], dtype=float32)
-
-    See also:
-      :func:`relu6`
-
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> import brainstate 
+        >>> brainstate.nn.relu(jnp.array([-2., -1., -0.5, 0, 0.5, 1., 2.]))
+        Array([0. , 0. , 0. , 0. , 0.5, 1. , 2. ], dtype=float32)
+
+    See Also
+    --------
+    relu6 : ReLU6 activation function.
+    leaky_relu : Leaky ReLU activation function.
+
+    References
+    ----------
+    .. [1] For more information see "Numerical influence of ReLU'(0) on backpropagation"
+           https://openreview.net/forum?id=urrcVI-_jRm
     """
     return u.math.relu(x)
 
 
 def squareplus(x: ArrayLike, b: ArrayLike = 4) -> Union[jax.Array, u.Quantity]:
-    r"""Squareplus activation function.
+    r"""
+    Squareplus activation function.
 
-    Computes the element-wise function
+    Computes the element-wise function:
 
     .. math::
       \mathrm{squareplus}(x) = \frac{x + \sqrt{x^2 + b}}{2}
 
-    as described in https://arxiv.org/abs/2112.11687.
-
-    Args:
-      x : input array
-      b : smoothness parameter
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    b : ArrayLike, optional
+        Smoothness parameter. Default is 4.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    References
+    ----------
+    .. [1] So, D., et al. (2021). "Primer: Searching for Efficient Transformers
+           for Language Modeling." arXiv:2112.11687
     """
     return u.math.squareplus(x, b=b)
 
 
 def softplus(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Softplus activation function.
+    r"""
+    Softplus activation function.
 
-    Computes the element-wise function
+    Computes the element-wise function:
 
     .. math::
       \mathrm{softplus}(x) = \log(1 + e^x)
 
-    Args:
-      x : input array
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
     """
     return u.math.softplus(x)
 
 
 def soft_sign(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Soft-sign activation function.
+    r"""
+    Soft-sign activation function.
 
-    Computes the element-wise function
+    Computes the element-wise function:
 
     .. math::
       \mathrm{soft\_sign}(x) = \frac{x}{|x| + 1}
 
-    Args:
-      x : input array
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
     """
     return u.math.soft_sign(x)
 
 
 def sigmoid(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Sigmoid activation function.
+    r"""
+    Sigmoid activation function.
 
     Computes the element-wise function:
 
     .. math::
       \mathrm{sigmoid}(x) = \frac{1}{1 + e^{-x}}
 
-    Args:
-      x : input array
-
-    Returns:
-      An array.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    See also:
-      :func:`log_sigmoid`
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
 
+    See Also
+    --------
+    log_sigmoid : Logarithm of the sigmoid function.
     """
     return u.math.sigmoid(x)
 
 
 def silu(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""SiLU (a.k.a. swish) activation function.
+    r"""
+    SiLU (Sigmoid Linear Unit) activation function.
 
     Computes the element-wise function:
 
     .. math::
       \mathrm{silu}(x) = x \cdot \mathrm{sigmoid}(x) = \frac{x}{1 + e^{-x}}
 
-    :func:`swish` and :func:`silu` are both aliases for the same function.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    Args:
-      x : input array
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
 
-    Returns:
-      An array.
+    See Also
+    --------
+    sigmoid : The sigmoid function.
+    swish : Alias for silu.
 
-    See also:
-      :func:`sigmoid`
+    Notes
+    -----
+    `swish` and `silu` are both aliases for the same function.
     """
     return u.math.silu(x)
 
@@ -372,27 +499,34 @@ swish = silu
 
 
 def log_sigmoid(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Log-sigmoid activation function.
+    r"""
+    Log-sigmoid activation function.
 
     Computes the element-wise function:
 
     .. math::
       \mathrm{log\_sigmoid}(x) = \log(\mathrm{sigmoid}(x)) = -\log(1 + e^{-x})
 
-    Args:
-      x : input array
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    Returns:
-      An array.
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
 
-    See also:
-      :func:`sigmoid`
+    See Also
+    --------
+    sigmoid : The sigmoid function.
     """
     return u.math.log_sigmoid(x)
 
 
 def elu(x: ArrayLike, alpha: ArrayLike = 1.0) -> Union[jax.Array, u.Quantity]:
-    r"""Exponential linear unit activation function.
+    r"""
+    Exponential Linear Unit activation function.
 
     Computes the element-wise function:
 
@@ -402,21 +536,29 @@ def elu(x: ArrayLike, alpha: ArrayLike = 1.0) -> Union[jax.Array, u.Quantity]:
         \alpha \left(\exp(x) - 1\right), & x \le 0
       \end{cases}
 
-    Args:
-      x : input array
-      alpha : scalar or array of alpha values (default: 1.0)
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`selu`
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    alpha : ArrayLike, optional
+        Scalar or array of alpha values. Default is 1.0.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    selu : Scaled ELU activation function.
+    celu : Continuously-differentiable ELU activation function.
     """
     return u.math.elu(x, alpha=alpha)
 
 
 def leaky_relu(x: ArrayLike, negative_slope: ArrayLike = 1e-2) -> Union[jax.Array, u.Quantity]:
-    r"""Leaky rectified linear unit activation function.
+    r"""
+    Leaky Rectified Linear Unit activation function.
 
     Computes the element-wise function:
 
@@ -428,15 +570,22 @@ def leaky_relu(x: ArrayLike, negative_slope: ArrayLike = 1e-2) -> Union[jax.Arra
 
     where :math:`\alpha` = :code:`negative_slope`.
 
-    Args:
-      x : input array
-      negative_slope : array or scalar specifying the negative slope (default: 0.01)
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`relu`
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    negative_slope : ArrayLike, optional
+        Array or scalar specifying the negative slope. Default is 0.01.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    relu : Standard ReLU activation function.
+    prelu : Parametric ReLU with learnable slope.
     """
     return u.math.leaky_relu(x, negative_slope=negative_slope)
 
@@ -450,7 +599,8 @@ def hard_tanh(
     min_val: float = - 1.0,
     max_val: float = 1.0
 ) -> Union[jax.Array, u.Quantity]:
-    r"""Hard :math:`\mathrm{tanh}` activation function.
+    r"""
+    Hard hyperbolic tangent activation function.
 
     Computes the element-wise function:
 
@@ -461,13 +611,19 @@ def hard_tanh(
         1, & 1 < x
       \end{cases}
 
-    Args:
-      x : input array
-      min_val: float. minimum value of the linear region range. Default: -1
-      max_val: float. maximum value of the linear region range. Default: 1
-
-    Returns:
-      An array.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    min_val : float, optional
+        Minimum value of the linear region range. Default is -1.
+    max_val : float, optional
+        Maximum value of the linear region range. Default is 1.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
     """
     x = u.Quantity(x)
     min_val = u.Quantity(min_val).to(x.unit).mantissa
@@ -476,7 +632,8 @@ def hard_tanh(
 
 
 def celu(x: ArrayLike, alpha: ArrayLike = 1.0) -> Union[jax.Array, u.Quantity]:
-    r"""Continuously-differentiable exponential linear unit activation.
+    r"""
+    Continuously-differentiable Exponential Linear Unit activation.
 
     Computes the element-wise function:
 
@@ -486,22 +643,29 @@ def celu(x: ArrayLike, alpha: ArrayLike = 1.0) -> Union[jax.Array, u.Quantity]:
         \alpha \left(\exp(\frac{x}{\alpha}) - 1\right), & x \le 0
       \end{cases}
 
-    For more information, see
-    `Continuously Differentiable Exponential Linear Units
-    <https://arxiv.org/pdf/1704.07483.pdf>`_.
-
-    Args:
-      x : input array
-      alpha : array or scalar (default: 1.0)
-
-    Returns:
-      An array.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    alpha : ArrayLike, optional
+        Scalar or array value controlling the smoothness. Default is 1.0.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    References
+    ----------
+    .. [1] Barron, J. T. (2017). "Continuously Differentiable Exponential Linear Units."
+           arXiv:1704.07483
     """
     return u.math.celu(x, alpha=alpha)
 
 
 def selu(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Scaled exponential linear unit activation.
+    r"""
+    Scaled Exponential Linear Unit activation.
 
     Computes the element-wise function:
 
@@ -514,24 +678,31 @@ def selu(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
     where :math:`\lambda = 1.0507009873554804934193349852946` and
     :math:`\alpha = 1.6732632423543772848170429916717`.
 
-    For more information, see
-    `Self-Normalizing Neural Networks
-    <https://papers.nips.cc/paper/6698-self-normalizing-neural-networks.pdf>`_.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    Args:
-      x : input array
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
 
-    Returns:
-      An array.
+    See Also
+    --------
+    elu : Exponential Linear Unit activation function.
 
-    See also:
-      :func:`elu`
+    References
+    ----------
+    .. [1] Klambauer, G., et al. (2017). "Self-Normalizing Neural Networks."
+           NeurIPS 2017.
     """
     return u.math.selu(x)
 
 
 def gelu(x: ArrayLike, approximate: bool = True) -> Union[jax.Array, u.Quantity]:
-    r"""Gaussian error linear unit activation function.
+    r"""
+    Gaussian Error Linear Unit activation function.
 
     If ``approximate=False``, computes the element-wise function:
 
@@ -545,18 +716,30 @@ def gelu(x: ArrayLike, approximate: bool = True) -> Union[jax.Array, u.Quantity]
       \mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{tanh} \left(
         \sqrt{\frac{2}{\pi}} \left(x + 0.044715 x^3 \right) \right) \right)
 
-    For more information, see `Gaussian Error Linear Units (GELUs)
-    <https://arxiv.org/abs/1606.08415>`_, section 2.
-
-    Args:
-      x : input array
-      approximate: whether to use the approximate or exact formulation.
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    approximate : bool, optional
+        Whether to use the approximate (True) or exact (False) formulation.
+        Default is True.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    References
+    ----------
+    .. [1] Hendrycks, D., & Gimpel, K. (2016). "Gaussian Error Linear Units (GELUs)."
+           arXiv:1606.08415
     """
     return u.math.gelu(x, approximate=approximate)
 
 
 def glu(x: ArrayLike, axis: int = -1) -> Union[jax.Array, u.Quantity]:
-    r"""Gated linear unit activation function.
+    r"""
+    Gated Linear Unit activation function.
 
     Computes the function:
 
@@ -568,15 +751,22 @@ def glu(x: ArrayLike, axis: int = -1) -> Union[jax.Array, u.Quantity]:
     where the array is split into two along ``axis``. The size of the ``axis``
     dimension must be divisible by two.
 
-    Args:
-      x : input array
-      axis: the axis along which the split should be computed (default: -1)
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`sigmoid`
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array. The dimension specified by ``axis`` must be divisible by 2.
+    axis : int, optional
+        The axis along which the split should be computed. Default is -1.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as input except the ``axis`` dimension
+        is halved.
+
+    See Also
+    --------
+    sigmoid : The sigmoid activation function.
     """
     return u.math.glu(x, axis=axis)
 
@@ -584,26 +774,34 @@ def glu(x: ArrayLike, axis: int = -1) -> Union[jax.Array, u.Quantity]:
 def log_softmax(x: ArrayLike,
                 axis: int | tuple[int, ...] | None = -1,
                 where: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
-    r"""Log-Softmax function.
+    r"""
+    Log-Softmax function.
 
-    Computes the logarithm of the :code:`softmax` function, which rescales
+    Computes the logarithm of the softmax function, which rescales
     elements to the range :math:`[-\infty, 0)`.
 
     .. math ::
       \mathrm{log\_softmax}(x)_i = \log \left( \frac{\exp(x_i)}{\sum_j \exp(x_j)}
       \right)
 
-    Args:
-      x : input array
-      axis: the axis or axes along which the :code:`log_softmax` should be
-        computed. Either an integer or a tuple of integers.
-      where: Elements to include in the :code:`log_softmax`.
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`softmax`
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    axis : int or tuple of int, optional
+        The axis or axes along which the log-softmax should be computed.
+        Either an integer or a tuple of integers. Default is -1.
+    where : ArrayLike, optional
+        Elements to include in the log-softmax computation.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    softmax : The softmax function.
     """
     return jax.nn.log_softmax(x, axis=axis, where=where)
 
@@ -611,7 +809,8 @@ def log_softmax(x: ArrayLike,
 def softmax(x: ArrayLike,
             axis: int | tuple[int, ...] | None = -1,
             where: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
-    r"""Softmax function.
+    r"""
+    Softmax activation function.
 
     Computes the function which rescales elements to the range :math:`[0, 1]`
     such that the elements along :code:`axis` sum to :math:`1`.
@@ -619,20 +818,26 @@ def softmax(x: ArrayLike,
     .. math ::
       \mathrm{softmax}(x) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
 
-    Args:
-      x : input array
-      axis: the axis or axes along which the softmax should be computed. The
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    axis : int or tuple of int, optional
+        The axis or axes along which the softmax should be computed. The
         softmax output summed across these dimensions should sum to :math:`1`.
-        Either an integer or a tuple of integers.
-      where: Elements to include in the :code:`softmax`.
-      initial: The minimum value used to shift the input array. Must be present
-        when :code:`where` is not None.
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`log_softmax`
+        Either an integer or a tuple of integers. Default is -1.
+    where : ArrayLike, optional
+        Elements to include in the softmax computation.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    log_softmax : Logarithm of the softmax function.
+    softmin : Softmin activation function.
     """
     return jax.nn.softmax(x, axis=axis, where=where)
 
@@ -642,7 +847,32 @@ def standardize(x: ArrayLike,
                 variance: ArrayLike | None = None,
                 epsilon: ArrayLike = 1e-5,
                 where: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
-    r"""Normalizes an array by subtracting ``mean`` and dividing by :math:`\sqrt{\mathrm{variance}}`."""
+    r"""
+    Standardize (normalize) an array.
+
+    Normalizes an array by subtracting the mean and dividing by the standard
+    deviation :math:`\sqrt{\mathrm{variance}}`.
+
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+    axis : int or tuple of int, optional
+        The axis or axes along which to compute the mean and variance.
+        Default is -1.
+    variance : ArrayLike, optional
+        Pre-computed variance. If None, variance is computed from ``x``.
+    epsilon : ArrayLike, optional
+        A small constant added to the variance to avoid division by zero.
+        Default is 1e-5.
+    where : ArrayLike, optional
+        Elements to include in the computation.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        Standardized array with the same shape as the input.
+    """
     return jax.nn.standardize(x, axis=axis, where=where, variance=variance, epsilon=epsilon)
 
 
@@ -650,41 +880,60 @@ def one_hot(x: Any,
             num_classes: int, *,
             dtype: Any = jax.numpy.float_,
             axis: Union[int, Sequence[int]] = -1) -> Union[jax.Array, u.Quantity]:
-    """One-hot encodes the given indices.
+    """
+    One-hot encode the given indices.
 
     Each index in the input ``x`` is encoded as a vector of zeros of length
-    ``num_classes`` with the element at ``index`` set to one::
-
-      >>> one_hot(jnp.array([0, 1, 2]), 3)
-      Array([[1., 0., 0.],
-             [0., 1., 0.],
-             [0., 0., 1.]], dtype=float32)
-
-    Indices outside the range [0, num_classes) will be encoded as zeros::
-
-      >>> one_hot(jnp.array([-1, 3]), 3)
-      Array([[0., 0., 0.],
-             [0., 0., 0.]], dtype=float32)
-
-    Args:
-      x: A tensor of indices.
-      num_classes: Number of classes in the one-hot dimension.
-      dtype: optional, a float dtype for the returned values (default :obj:`jnp.float_`).
-      axis: the axis or axes along which the function should be
-        computed.
+    ``num_classes`` with the element at ``index`` set to one.
+
+    Indices outside the range [0, num_classes) will be encoded as zeros.
+
+    Parameters
+    ----------
+    x : ArrayLike
+        A tensor of indices.
+    num_classes : int
+        Number of classes in the one-hot dimension.
+    dtype : dtype, optional
+        The dtype for the returned values. Default is ``jnp.float_``.
+    axis : int or Sequence of int, optional
+        The axis or axes along which the function should be computed.
+        Default is -1.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        One-hot encoded array.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> import brainstate 
+        >>> brainstate.nn.one_hot(jnp.array([0, 1, 2]), 3)
+        Array([[1., 0., 0.],
+               [0., 1., 0.],
+               [0., 0., 1.]], dtype=float32)
+
+        >>> # Indices outside the range are encoded as zeros
+        >>> brainstate.nn.one_hot(jnp.array([-1, 3]), 3)
+        Array([[0., 0., 0.],
+               [0., 0., 0.]], dtype=float32)
     """
     return jax.nn.one_hot(x, axis=axis, num_classes=num_classes, dtype=dtype)
 
 
 def relu6(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Rectified Linear Unit 6 activation function.
+    r"""
+    Rectified Linear Unit 6 activation function.
 
-    Computes the element-wise function
+    Computes the element-wise function:
 
     .. math::
       \mathrm{relu6}(x) = \min(\max(x, 0), 6)
 
-    except under differentiation, we take:
+    Under differentiation, we take:
 
     .. math::
       \nabla \mathrm{relu}(0) = 0
@@ -694,57 +943,78 @@ def relu6(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
     .. math::
       \nabla \mathrm{relu}(6) = 0
 
-    Args:
-      x : input array
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    Returns:
-      An array.
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
 
-    See also:
-      :func:`relu`
+    See Also
+    --------
+    relu : Standard ReLU activation function.
     """
     return u.math.relu6(x)
 
 
 def hard_sigmoid(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Hard Sigmoid activation function.
+    r"""
+    Hard Sigmoid activation function.
 
-    Computes the element-wise function
+    Computes the element-wise function:
 
     .. math::
       \mathrm{hard\_sigmoid}(x) = \frac{\mathrm{relu6}(x + 3)}{6}
 
-    Args:
-      x : input array
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
 
-    Returns:
-      An array.
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
 
-    See also:
-      :func:`relu6`
+    See Also
+    --------
+    relu6 : ReLU6 activation function.
+    sigmoid : Standard sigmoid function.
     """
     return u.math.hard_sigmoid(x)
 
 
 def hard_silu(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Hard SiLU (swish) activation function
+    r"""
+    Hard SiLU (Swish) activation function.
 
-    Computes the element-wise function
+    Computes the element-wise function:
 
     .. math::
       \mathrm{hard\_silu}(x) = x \cdot \mathrm{hard\_sigmoid}(x)
 
-    Both :func:`hard_silu` and :func:`hard_swish` are aliases for the same
-    function.
-
-    Args:
-      x : input array
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`hard_sigmoid`
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    hard_sigmoid : Hard sigmoid activation function.
+    silu : Standard SiLU activation function.
+    hard_swish : Alias for hard_silu.
+
+    Notes
+    -----
+    Both `hard_silu` and `hard_swish` are aliases for the same function.
     """
     return u.math.hard_silu(x)
 
@@ -753,7 +1023,8 @@ hard_swish = hard_silu
 
 
 def sparse_plus(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Sparse plus function.
+    r"""
+    Sparse plus activation function.
 
     Computes the function:
 
@@ -765,19 +1036,31 @@ def sparse_plus(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
         x, & 1 \leq x
       \end{cases}
 
-    This is the twin function of the softplus activation ensuring a zero output
+    This is the twin function of the softplus activation, ensuring a zero output
     for inputs less than -1 and a linear output for inputs greater than 1,
-    while remaining smooth, convex, monotonic by an adequate definition between
-    -1 and 1.
-
-    Args:
-      x: input (float)
+    while remaining smooth, convex, and monotonic between -1 and 1.
+
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    sparse_sigmoid : Derivative of sparse_plus.
+    softplus : Standard softplus activation function.
     """
     return u.math.sparse_plus(x)
 
 
 def sparse_sigmoid(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
-    r"""Sparse sigmoid activation function.
+    r"""
+    Sparse sigmoid activation function.
 
     Computes the function:
 
@@ -789,20 +1072,29 @@ def sparse_sigmoid(x: ArrayLike) -> Union[jax.Array, u.Quantity]:
         1, & 1 \leq x
       \end{cases}
 
-    This is the twin function of the ``sigmoid`` activation ensuring a zero output
-    for inputs less than -1, a 1 output for inputs greater than 1, and a linear
-    output for inputs between -1 and 1. It is the derivative of ``sparse_plus``.
-
-    For more information, see `Learning with Fenchel-Young Losses (section 6.2)
-    <https://arxiv.org/abs/1901.02324>`_.
-
-    Args:
-      x : input array
-
-    Returns:
-      An array.
-
-    See also:
-      :func:`sigmoid`
+    This is the twin function of the standard sigmoid activation, ensuring a zero
+    output for inputs less than -1, a 1 output for inputs greater than 1, and a
+    linear output for inputs between -1 and 1. It is the derivative of `sparse_plus`.
+
+    Parameters
+    ----------
+    x : ArrayLike
+        Input array.
+
+    Returns
+    -------
+    jax.Array or Quantity
+        An array with the same shape as the input.
+
+    See Also
+    --------
+    sigmoid : Standard sigmoid activation function.
+    sparse_plus : Sparse plus activation function.
+
+    References
+    ----------
+    .. [1] Martins, A. F. T., & Astudillo, R. F. (2016). "From Softmax to Sparsemax:
+           A Sparse Model of Attention and Multi-Label Classification."
+           In ICML. See also "Learning with Fenchel-Young Losses", arXiv:1901.02324
     """
     return u.math.sparse_sigmoid(x)