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

brainstate/{augment → transform}/_autograd.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.
@@ -35,7 +35,7 @@ import jax
 
 from brainstate._state import State
 from brainstate._utils import set_module_as
-from brainstate.compile._make_jaxpr import StatefulFunction
+from brainstate.transform._make_jaxpr import StatefulFunction
 from brainstate.typing import PyTree, Missing
 from brainstate.util import PrettyType, PrettyAttr, PrettyRepr
 
@@ -153,16 +153,90 @@ class GradientTransform(PrettyRepr):
     It allows for flexible configuration of gradient computation with respect to specified states
     and function arguments.
 
-    Attributes:
-        target (Callable): The function to be transformed.
-        stateful_target (StatefulFunction): A wrapper around the target function for state management.
-        raw_argnums (Optional[Union[int, Sequence[int]]]): The original argnums specified by the user.
-        true_argnums (Union[int, Tuple[int, ...]]): The adjusted argnums used internally.
-        return_value (bool): Whether to return the function's value along with gradients.
-        has_aux (bool): Whether the function returns auxiliary data.
+    Parameters
+    ----------
+    target : callable
+        The function to be transformed.
+    transform : callable
+        The transformation function to apply.
+    grad_states : State, sequence of State, or dict of State, optional
+        States to compute gradients for.
+    argnums : int or sequence of int, optional
+        Indices of arguments to differentiate with respect to.
+    return_value : bool, default False
+        Whether to return the function's value along with gradients.
+    has_aux : bool, default False
+        Whether the function returns auxiliary data.
+    transform_params : dict, optional
+        Additional parameters for the transformation function.
+    check_states : bool, default True
+        Whether to check that all grad_states are found in the function.
+
+    Attributes
+    ----------
+    target : callable
+        The function to be transformed.
+    stateful_target : StatefulFunction
+        A wrapper around the target function for state management.
+    raw_argnums : int, sequence of int, or None
+        The original argnums specified by the user.
+    true_argnums : int or tuple of int
+        The adjusted argnums used internally.
+    return_value : bool
+        Whether to return the function's value along with gradients.
+    has_aux : bool
+        Whether the function returns auxiliary data.
+
+    Examples
+    --------
+    Basic gradient computation with states:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # Create states
+        >>> weight = brainstate.State(jnp.array([[1.0, 2.0], [3.0, 4.0]]))
+        >>> bias = brainstate.State(jnp.array([0.5, -0.5]))
+        >>>
+        >>> def loss_fn(x):
+        ...     y = x @ weight.value + bias.value
+        ...     return jnp.sum(y ** 2)
+        >>>
+        >>> # Create gradient transform
+        >>> grad_transform = brainstate.transform.GradientTransform(
+        ...     target=loss_fn,
+        ...     transform=jax.grad,
+        ...     grad_states=[weight, bias]
+        ... )
+        >>>
+        >>> # Compute gradients
+        >>> x = jnp.array([1.0, 2.0])
+        >>> grads = grad_transform(x)
+
+    With function arguments and auxiliary data:
+
+    .. code-block:: python
+
+        >>> def loss_fn_with_aux(x, scale):
+        ...     y = x @ weight.value + bias.value
+        ...     loss = jnp.sum((y * scale) ** 2)
+        ...     return loss, {"predictions": y, "scale": scale}
+        >>>
+        >>> grad_transform = brainstate.transform.GradientTransform(
+        ...     target=loss_fn_with_aux,
+        ...     transform=jax.grad,
+        ...     grad_states=[weight, bias],
+        ...     argnums=[0, 1],  # gradient w.r.t x and scale
+        ...     has_aux=True,
+        ...     return_value=True
+        ... )
+        >>>
+        >>> grads, loss_value, aux_data = grad_transform(x, 2.0)
     """
 
-    __module__ = "brainstate.augment"
+    __module__ = "brainstate.transform"
 
     def __init__(
         self,
@@ -178,17 +252,29 @@ class GradientTransform(PrettyRepr):
         """
         Initialize a ``GradientTransform`` instance.
 
-        Args:
-            target (Callable): The function to be transformed.
-            transform (TransformFn): The transformation function to apply.
-            grad_states (Optional[Union[State, Sequence[State], Dict[str, State]]]): States to compute gradients for.
-            argnums (Optional[Union[int, Sequence[int]]]): Indices of arguments to differentiate with respect to.
-            return_value (bool): Whether to return the function's value along with gradients.
-            has_aux (bool): Whether the function returns auxiliary data.
-            transform_params (Optional[Dict[str, Any]]): Additional parameters for the transformation function.
-
-        Raises:
-            TypeError: If any grad_states are not State instances.
+        Parameters
+        ----------
+        target : callable
+            The function to be transformed.
+        transform : callable
+            The transformation function to apply.
+        grad_states : State, sequence of State, or dict of State, optional
+            States to compute gradients for.
+        argnums : int or sequence of int, optional
+            Indices of arguments to differentiate with respect to.
+        return_value : bool, default False
+            Whether to return the function's value along with gradients.
+        has_aux : bool, default False
+            Whether the function returns auxiliary data.
+        transform_params : dict, optional
+            Additional parameters for the transformation function.
+        check_states : bool, default True
+            Whether to check that all grad_states are found in the function.
+
+        Raises
+        ------
+        TypeError
+            If any grad_states are not State instances.
         """
         # gradient variables
         if isinstance(grad_states, dict):
@@ -221,7 +307,7 @@ class GradientTransform(PrettyRepr):
         # target
         assert callable(target), "The target should be a callable object."
         self.target = target
-        self.stateful_target = StatefulFunction(target, name='gradient')
+        self.stateful_target = StatefulFunction(target, name='gradient', return_only_write=False)
 
         # transform
         grad_setting = dict() if transform_params is None else transform_params
@@ -307,8 +393,7 @@ class GradientTransform(PrettyRepr):
         Returns:
             Tuple: A tuple containing updated state values and the function output.
         """
-        cache = self.stateful_target.get_arg_cache_key(*args, **kwargs)
-        state_trace = self.stateful_target.get_state_trace(cache)
+        state_trace = self.stateful_target.get_state_trace(*args, **kwargs, compile_if_miss=True)
         state_vals = self._merge_state_vals(grad_vals, other_vals, state_trace)
         state_vals, out = self.stateful_target.jaxpr_call(state_vals, *args, **kwargs)
         return state_vals, out
@@ -403,12 +488,18 @@ class GradientTransform(PrettyRepr):
         """
         Compute gradients by calling the transformed function.
 
-        Args:
-            *args: Positional arguments to pass to the target function.
-            **kwargs: Keyword arguments to pass to the target function.
-
-        Returns:
-            Union[Gradient, Tuple]: The computed gradients, potentially including function value and/or auxiliary data.
+        Parameters
+        ----------
+        *args
+            Positional arguments to pass to the target function.
+        **kwargs
+            Keyword arguments to pass to the target function.
+
+        Returns
+        -------
+        Gradient or tuple
+            The computed gradients, potentially including function value and/or auxiliary data.
+            The exact return structure depends on the settings of return_value and has_aux.
         """
 
         # TODO: support jax.disable_jit()
@@ -418,79 +509,135 @@ class GradientTransform(PrettyRepr):
         cache = self.stateful_target.get_arg_cache_key(*args, **kwargs)
 
         # apply the gradient transformation
-        state_trace = self.stateful_target.get_state_trace(cache)
+        state_trace = self.stateful_target.get_state_trace_by_cache(cache)
         rets = self._transform(*self._split_state_vals(state_trace), *args, **kwargs)
 
         # analyze and return the results
         return self._return(rets, state_trace)
 
 
-_doc_of_return = '''
+@set_module_as("brainstate.transform")
+def grad(
+    fun: Callable = Missing(),
+    grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
+    argnums: Optional[Union[int, Sequence[int]]] = None,
+    holomorphic: Optional[bool] = False,
+    allow_int: Optional[bool] = False,
+    has_aux: Optional[bool] = None,
+    return_value: Optional[bool] = False,
+    unit_aware: bool = False,
+    check_states: bool = True,
+) -> GradientTransform | Callable[[Callable], GradientTransform]:
+    """
+    Compute the gradient of a scalar-valued function with respect to its arguments.
+
 
     1. When ``grad_states`` is None
+
         - ``has_aux=False`` + ``return_value=False`` => ``arg_grads``.
         - ``has_aux=True`` + ``return_value=False`` => ``(arg_grads, aux_data)``.
         - ``has_aux=False`` + ``return_value=True`` => ``(arg_grads, loss_value)``.
         - ``has_aux=True`` + ``return_value=True`` => ``(arg_grads, loss_value, aux_data)``.
     2. When ``grad_states`` is not None and ``argnums`` is None
+
         - ``has_aux=False`` + ``return_value=False`` => ``var_grads``.
         - ``has_aux=True`` + ``return_value=False`` => ``(var_grads, aux_data)``.
         - ``has_aux=False`` + ``return_value=True`` => ``(var_grads, loss_value)``.
         - ``has_aux=True`` + ``return_value=True`` => ``(var_grads, loss_value, aux_data)``.
     3. When ``grad_states`` is not None and ``argnums`` is not None
+
         - ``has_aux=False`` + ``return_value=False`` => ``(var_grads, arg_grads)``.
         - ``has_aux=True`` + ``return_value=False`` => ``((var_grads, arg_grads), aux_data)``.
         - ``has_aux=False`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value)``.
         - ``has_aux=True`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value, aux_data)``.
 
-'''
-
-
-@set_module_as("brainstate.augment")
-def grad(
-    fun: Callable = Missing(),
-    grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
-    argnums: Optional[Union[int, Sequence[int]]] = None,
-    holomorphic: Optional[bool] = False,
-    allow_int: Optional[bool] = False,
-    has_aux: Optional[bool] = None,
-    return_value: Optional[bool] = False,
-    unit_aware: bool = False,
-    check_states: bool = True,
-) -> GradientTransform | Callable[[Callable], GradientTransform]:
-    """
-    Compute the gradient of a scalar-valued function with respect to its arguments.
 
-    %s
-
-    Args:
-        fun: callable. the scalar-valued function to be differentiated.
-        allow_int: (bool) optional. Whether to allow differentiating with respect to
-            integer valued inputs. The gradient of an integer input will have a trivial
-            vector-space dtype (float0). Default False.
-        holomorphic: (bool) optional. Whether fun is promised to be holomorphic.
-            Default False.
-        grad_states: (State, Sequence[State], Dict[str, State]) optional. The variables
-            in fun to take their gradients.
-        fun: the scalar-valued function to be differentiated.
-        argnums: (int or tuple of ints) optional. Specifies which positional
-            argument(s) to differentiate with respect to.
-        has_aux: (bool) optional. Indicates whether fun returns a pair where the
-            first element is considered the output of the mathematical function to be
-            differentiated and the second element is auxiliary data. Default False.
-        return_value: (bool) optional. Indicates whether to return the value of the
-            function along with the gradient. Default False.
-        unit_aware: (bool) optional. Whether to return the gradient in the unit-aware
-            mode. Default False.
-
-    Returns:
-      A function which computes the gradient of fun. The function takes the same
-      arguments as `fun`, but returns the gradient instead. If `has_aux` is True,
-      the function returns a pair where the first element is the gradient and the
-      second element is the auxiliary data. If `return_value` is True, the function
-      returns a pair where the first element is the gradient and the second element
-      is the value of the function.
+    Parameters
+    ----------
+    fun : callable, optional
+        The scalar-valued function to be differentiated.
+    grad_states : State, sequence of State, or dict of State, optional
+        The variables in fun to take their gradients.
+    argnums : int or sequence of int, optional
+        Specifies which positional argument(s) to differentiate with respect to.
+    holomorphic : bool, default False
+        Whether fun is promised to be holomorphic.
+    allow_int : bool, default False
+        Whether to allow differentiating with respect to
+        integer valued inputs. The gradient of an integer input will have a trivial
+        vector-space dtype (float0).
+    has_aux : bool, optional
+        Indicates whether fun returns a pair where the
+        first element is considered the output of the mathematical function to be
+        differentiated and the second element is auxiliary data.
+    return_value : bool, default False
+        Indicates whether to return the value of the
+        function along with the gradient.
+    unit_aware : bool, default False
+        Whether to return the gradient in the unit-aware mode.
+    check_states : bool, default True
+        Whether to check that all grad_states are found in the function.
 
+    Returns
+    -------
+    GradientTransform or callable
+        A function which computes the gradient of fun. The function takes the same
+        arguments as `fun`, but returns the gradient instead. If `has_aux` is True,
+        the function returns a pair where the first element is the gradient and the
+        second element is the auxiliary data. If `return_value` is True, the function
+        returns a pair where the first element is the gradient and the second element
+        is the value of the function.
+
+    Examples
+    --------
+    Basic gradient computation:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # Simple function gradient
+        >>> def f(x):
+        ...     return jnp.sum(x ** 2)
+        >>>
+        >>> grad_f = brainstate.transform.grad(f)
+        >>> x = jnp.array([1.0, 2.0, 3.0])
+        >>> gradient = grad_f(x)
+
+    Gradient with respect to states:
+
+    .. code-block:: python
+
+        >>> # Create states
+        >>> weight = brainstate.State(jnp.array([1.0, 2.0]))
+        >>> bias = brainstate.State(jnp.array([0.5]))
+        >>>
+        >>> def loss_fn(x):
+        ...     prediction = jnp.dot(x, weight.value) + bias.value
+        ...     return prediction ** 2
+        >>>
+        >>> # Compute gradients with respect to states
+        >>> grad_fn = brainstate.transform.grad(loss_fn, grad_states=[weight, bias])
+        >>> x = jnp.array([1.0, 2.0])
+        >>> state_grads = grad_fn(x)
+
+    With auxiliary data and return value:
+
+    .. code-block:: python
+
+        >>> def loss_with_aux(x):
+        ...     prediction = jnp.dot(x, weight.value) + bias.value
+        ...     loss = prediction ** 2
+        ...     return loss, {"prediction": prediction}
+        >>>
+        >>> grad_fn = brainstate.transform.grad(
+        ...     loss_with_aux,
+        ...     grad_states=[weight, bias],
+        ...     has_aux=True,
+        ...     return_value=True
+        ... )
+        >>> grads, loss_value, aux_data = grad_fn(x)
     """
     if isinstance(fun, Missing):
         def transform(fun) -> GradientTransform:
@@ -519,10 +666,7 @@ def grad(
     )
 
 
-grad.__doc__ = grad.__doc__ % _doc_of_return
-
-
-@set_module_as("brainstate.augment")
+@set_module_as("brainstate.transform")
 def vector_grad(
     func: Callable = Missing(),
     grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
@@ -532,34 +676,91 @@ def vector_grad(
     unit_aware: bool = False,
     check_states: bool = True,
 ) -> GradientTransform | Callable[[Callable], GradientTransform]:
-    """Take vector-valued gradients for function ``func``.
+    """
+    Take vector-valued gradients for function ``func``.
 
-    Same as :py:func:`grad`, :py:func:`jacrev`, and :py:func:`jacfwd`, 
+    Same as :py:func:`grad`, :py:func:`jacrev`, and :py:func:`jacfwd`,
     the returns in this function are different for different argument settings.
 
-    %s
+
+    1. When ``grad_states`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``arg_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(arg_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(arg_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(arg_grads, loss_value, aux_data)``.
+    2. When ``grad_states`` is not None and ``argnums`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``var_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(var_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(var_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(var_grads, loss_value, aux_data)``.
+    3. When ``grad_states`` is not None and ``argnums`` is not None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``(var_grads, arg_grads)``.
+        - ``has_aux=True`` + ``return_value=False`` => ``((var_grads, arg_grads), aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value, aux_data)``.
+
 
     Parameters
     ----------
-    func: Callable
+    func : callable, optional
         Function whose gradient is to be computed.
-    grad_states : optional, ArrayType, sequence of ArrayType, dict
+    grad_states : State, sequence of State, or dict of State, optional
         The variables in ``func`` to take their gradients.
-    has_aux: optional, bool
+    argnums : int or sequence of int, optional
+        Specifies which positional argument(s) to differentiate with respect to.
+    return_value : bool, default False
+        Whether to return the loss value.
+    has_aux : bool, optional
         Indicates whether ``fun`` returns a pair where the
         first element is considered the output of the mathematical function to be
-        differentiated and the second element is auxiliary data. Default False.
-    return_value : bool
-        Whether return the loss value.
-    argnums: Optional, integer or sequence of integers. Specifies which
-        positional argument(s) to differentiate with respect to (default ``0``).
-    unit_aware: (bool) optional. Whether to return the gradient in the unit-aware
-        mode. Default False.
+        differentiated and the second element is auxiliary data.
+    unit_aware : bool, default False
+        Whether to return the gradient in the unit-aware mode.
+    check_states : bool, default True
+        Whether to check that all grad_states are found in the function.
 
     Returns
     -------
-    func : GradientTransform
+    GradientTransform or callable
         The vector gradient function.
+
+    Examples
+    --------
+    Basic vector gradient computation:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # Vector-valued function
+        >>> def f(x):
+        ...     return jnp.array([x[0]**2, x[1]**3, x[0]*x[1]])
+        >>>
+        >>> vector_grad_f = brainstate.transform.vector_grad(f)
+        >>> x = jnp.array([2.0, 3.0])
+        >>> gradients = vector_grad_f(x)  # Shape: (3, 2)
+
+    With states:
+
+    .. code-block:: python
+
+        >>> params = brainstate.State(jnp.array([1.0, 2.0]))
+        >>>
+        >>> def model(x):
+        ...     return jnp.array([
+        ...         x * params.value[0],
+        ...         x**2 * params.value[1]
+        ...     ])
+        >>>
+        >>> vector_grad_fn = brainstate.transform.vector_grad(
+        ...     model, grad_states=[params]
+        ... )
+        >>> x = 3.0
+        >>> param_grads = vector_grad_fn(x)
     """
 
     if isinstance(func, Missing):
@@ -588,10 +789,7 @@ def vector_grad(
         )
 
 
-vector_grad.__doc__ = vector_grad.__doc__ % _doc_of_return
-
-
-@set_module_as("brainstate.augment")
+@set_module_as("brainstate.transform")
 def jacrev(
     fun: Callable,
     grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
@@ -610,7 +808,26 @@ def jacrev(
     computation on functions and class functions. Moreover, it supports returning
     value ("return_value") and returning auxiliary data ("has_aux").
 
-    %s
+
+    1. When ``grad_states`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``arg_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(arg_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(arg_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(arg_grads, loss_value, aux_data)``.
+    2. When ``grad_states`` is not None and ``argnums`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``var_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(var_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(var_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(var_grads, loss_value, aux_data)``.
+    3. When ``grad_states`` is not None and ``argnums`` is not None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``(var_grads, arg_grads)``.
+        - ``has_aux=True`` + ``return_value=False`` => ``((var_grads, arg_grads), aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value, aux_data)``.
+
 
 
     Parameters
@@ -657,12 +874,10 @@ def jacrev(
     )
 
 
-jacrev.__doc__ = jacrev.__doc__ % _doc_of_return
-
 jacobian = jacrev
 
 
-@set_module_as("brainstate.augment")
+@set_module_as("brainstate.transform")
 def jacfwd(
     func: Callable,
     grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
@@ -679,7 +894,26 @@ def jacfwd(
     computation on functions and class functions. Moreover, it supports returning
     value ("return_value") and returning auxiliary data ("has_aux").
 
-    %s
+
+    1. When ``grad_states`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``arg_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(arg_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(arg_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(arg_grads, loss_value, aux_data)``.
+    2. When ``grad_states`` is not None and ``argnums`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``var_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(var_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(var_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(var_grads, loss_value, aux_data)``.
+    3. When ``grad_states`` is not None and ``argnums`` is not None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``(var_grads, arg_grads)``.
+        - ``has_aux=True`` + ``return_value=False`` => ``((var_grads, arg_grads), aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value, aux_data)``.
+
 
     Parameters
     ----------
@@ -717,10 +951,7 @@ def jacfwd(
     )
 
 
-jacfwd.__doc__ = jacfwd.__doc__ % _doc_of_return
-
-
-@set_module_as("brainstate.augment")
+@set_module_as("brainstate.transform")
 def hessian(
     func: Callable,
     grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
@@ -734,7 +965,26 @@ def hessian(
     """
     Hessian of ``func`` as a dense array.
 
-    %s
+
+    1. When ``grad_states`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``arg_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(arg_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(arg_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(arg_grads, loss_value, aux_data)``.
+    2. When ``grad_states`` is not None and ``argnums`` is None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``var_grads``.
+        - ``has_aux=True`` + ``return_value=False`` => ``(var_grads, aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``(var_grads, loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``(var_grads, loss_value, aux_data)``.
+    3. When ``grad_states`` is not None and ``argnums`` is not None
+
+        - ``has_aux=False`` + ``return_value=False`` => ``(var_grads, arg_grads)``.
+        - ``has_aux=True`` + ``return_value=False`` => ``((var_grads, arg_grads), aux_data)``.
+        - ``has_aux=False`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value)``.
+        - ``has_aux=True`` + ``return_value=True`` => ``((var_grads, arg_grads), loss_value, aux_data)``.
+
 
     Parameters
     ----------
@@ -773,6 +1023,3 @@ def hessian(
         transform_params=dict(holomorphic=holomorphic),
         check_states=check_states
     )
-
-
-hessian.__doc__ = hessian.__doc__ % _doc_of_return

brainstate/{augment → transform}/_autograd_test.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.
@@ -24,7 +24,7 @@ import jax.numpy as jnp
 import pytest
 
 import brainstate
-from brainstate.augment._autograd import _jacfwd
+from brainstate.transform._autograd import _jacfwd
 
 
 class TestPureFuncGrad(unittest.TestCase):