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

brainstate/_state.py

@@ -33,6 +33,7 @@ from typing import (
     Generator,
 )
 
+import brainunit as u
 import jax
 import numpy as np
 from jax.api_util import shaped_abstractify
@@ -47,6 +48,8 @@ __all__ = [
     'ShortTermState',
     'LongTermState',
     'HiddenState',
+    'HiddenGroupState',
+    'HiddenTreeState',
     'ParamState',
     'BatchState',
     'TreefyState',
@@ -213,7 +216,7 @@ class State(Generic[A], PrettyObject):
         tag (Optional[str]): An optional tag for categorizing or grouping states.
 
     Args:
-        value (Union[PyTree[ArrayLike], StateMetadata[PyTree[ArrayLike]]]): 
+        value (Union[PyTree[ArrayLike], StateMetadata[PyTree[ArrayLike]]]):
             The initial value for the state. Can be a PyTree of array-like objects
             or a StateMetadata object.
         name (Optional[str]): An optional name for the state.
@@ -233,13 +236,13 @@ class State(Generic[A], PrettyObject):
          [0. 0. 0.]]
 
     Note:
-        - Subclasses of :class:`State` (e.g., ShortTermState, LongTermState, ParamState, 
+        - Subclasses of :class:`State` (e.g., ShortTermState, LongTermState, ParamState,
           RandomState) are typically used for specific purposes in a program.
-        - The class integrates with BrainState's tracing system to track state 
+        - The class integrates with BrainState's tracing system to track state
           creation and modifications.
 
         The typical examples of :py:class:`~.State` subclass are:
-    
+
         - :py:class:`ShortTermState`: The short-term state, which is used to store the short-term data in the program.
         - :py:class:`LongTermState`: The long-term state, which is used to store the long-term data in the program.
         - :py:class:`ParamState`: The parameter state, which is used to store the parameters in the program.
@@ -271,7 +274,7 @@ class State(Generic[A], PrettyObject):
         handling various input types and metadata.
 
         Args:
-            value (Union[PyTree[ArrayLike], StateMetadata[PyTree[ArrayLike]]]): 
+            value (Union[PyTree[ArrayLike], StateMetadata[PyTree[ArrayLike]]]):
                 The initial value for the hidden state. Can be a PyTree of array-like objects
                 or a StateMetadata object containing both value and metadata.
             name (Optional[str], optional): A name for the hidden state. Defaults to None.
@@ -738,7 +741,7 @@ class BatchState(LongTermState):
 
 class HiddenState(ShortTermState):
     """
-     Represents hidden state variables in neurons or synapses.
+    Represents hidden state variables in neurons or synapses.
 
     This class extends :class:`ShortTermState` and is specifically designed to represent
     and manage hidden states within dynamic models, such as recurrent neural networks.
@@ -751,6 +754,12 @@ class HiddenState(ShortTermState):
         or networks. The latter is used to store the trainable parameters in the model,
         such as synaptic weights.
 
+    Note:
+        From version 0.2.2, :class:`HiddenState` only supports value of numpy.ndarray,
+        jax.Array or brainunit.Quantity. Moreover, it is equivalent to :class:`brainscale.ETraceState`.
+        Dynamics models defined with :class:`HiddenState` can be seamlessly integrated with
+        BrainScale online learning.
+
     Example:
         >>> lstm_hidden = HiddenState(np.zeros(128), name="lstm_hidden_state")
         >>> gru_hidden = HiddenState(np.zeros(64), name="gru_hidden_state")
@@ -758,6 +767,502 @@ class HiddenState(ShortTermState):
 
     __module__ = 'brainstate'
 
+    value: ArrayLike
+
+    def __init__(self, value: ArrayLike, name: Optional[str] = None):
+        self._check_value(value)
+        super().__init__(value, name=name)
+
+    @property
+    def varshape(self) -> Tuple[int, ...]:
+        """
+        Get the shape of the hidden state variable.
+
+        This property returns the shape of the hidden state variable stored in the instance.
+        It provides the dimensions of the array representing the hidden state.
+
+        Returns:
+            Tuple[int, ...]: A tuple representing the shape of the hidden state variable.
+        """
+        return self.value.shape
+
+    @property
+    def num_state(self) -> int:
+        """
+        Get the number of hidden states.
+
+        This property returns the number of hidden states represented by the instance.
+        For the `ETraceState` class, this is always 1, as it represents a single hidden state.
+
+        Returns:
+            int: The number of hidden states, which is 1 for this class.
+        """
+        return 1
+
+    def _check_value(self, value: ArrayLike):
+        if not isinstance(value, (np.ndarray, jax.Array, u.Quantity)):
+            raise TypeError(
+                f'Currently, {HiddenState.__name__} only supports '
+                f'numpy.ndarray, jax.Array or brainunit.Quantity. '
+                f'But we got {type(value)}.'
+            )
+
+
+class HiddenGroupState(HiddenState):
+    """
+    A group of multiple hidden states for eligibility trace-based learning.
+
+    This class is used to define multiple hidden states within a single instance
+    of :py:class:`ETraceState`. Normally, you should define multiple instances
+    of :py:class:`ETraceState` to represent multiple hidden states. But
+    :py:class:`HiddenGroupState` let your define multiple hidden states within
+    a single instance.
+
+    The following is the way to initialize the hidden states.
+
+    .. code-block:: python
+
+        import brainunit as u
+        value = np.random.randn(10, 10, 5) * u.mV
+        state = HiddenGroupState(value)
+
+    Then, you can retrieve the hidden state value with the following method.
+
+    .. code-block:: python
+
+        state.get_value(0)  # get the first hidden state
+        # or
+        state.get_value('0')  # get the hidden state with the name '0'
+
+    You can write the hidden state value with the following method.
+
+    .. code-block:: python
+
+        state.set_value({0: np.random.randn(10, 10) * u.mV})  # set the first hidden state
+        # or
+        state.set_value({'0': np.random.randn(10, 10) * u.mV})  # set the hidden state with the name '0'
+        # or
+        state.value = np.random.randn(10, 10, 5) * u.mV  # set all hidden state value
+
+    Args:
+        value: The values of the hidden states. It can be a sequence of hidden states,
+            or a single hidden state with the last dimension as the number of hidden states,
+            or a dictionary of hidden states.
+    """
+
+    __module__ = 'brainstate'
+    value: ArrayLike
+    name2index: Dict[str, int]
+
+    def __init__(self, value: ArrayLike):
+        value, name2index = self._check_value(value)
+        self.name2index = name2index
+        ShortTermState.__init__(self, value)
+
+    @property
+    def varshape(self) -> Tuple[int, ...]:
+        """
+        Get the shape of each hidden state variable.
+
+        This property returns the shape of the hidden state variables, excluding
+        the last dimension which represents the number of hidden states.
+
+        Returns:
+            Tuple[int, ...]: A tuple representing the shape of each hidden state variable.
+        """
+        return self.value.shape[:-1]
+
+    @property
+    def num_state(self) -> int:
+        """
+        Get the number of hidden states.
+
+        This property returns the number of hidden states represented by the last dimension
+        of the value array.
+
+        Returns:
+            int: The number of hidden states.
+        """
+        return self.value.shape[-1]
+
+    def _check_value(self, value) -> Tuple[ArrayLike, Dict[str, int]]:
+        """
+        Validates the input value for hidden states and returns a tuple containing
+        the processed value and a dictionary mapping state names to indices.
+
+        This function ensures that the input value is of a supported type and has
+        the required dimensionality for hidden states. It also constructs a mapping
+        from string representations of indices to their integer counterparts.
+
+        Parameters
+        ----------
+        value (ArrayLike): The input value representing hidden states.
+            It must be an instance of numpy.ndarray, jax.Array, or brainunit.Quantity
+            with at least two dimensions.
+
+        Returns
+        -------
+        Tuple[ArrayLike, Dict[str, int]]: A tuple containing:
+            - The validated and possibly modified input value.
+            - A dictionary mapping string representations of indices to integer indices.
+
+        Raises
+        ------
+        TypeError: If the input value is not of a supported type.
+        ValueError: If the input value does not have the required number of dimensions.
+        """
+        if not isinstance(value, (np.ndarray, jax.Array, u.Quantity)):
+            raise TypeError(
+                f'Currently, {self.__class__.__name__} only supports '
+                f'numpy.ndarray, jax.Array or brainunit.Quantity. '
+                f'But we got {type(value)}.'
+            )
+        if value.ndim < 2:
+            raise ValueError(
+                f'Currently, {self.__class__.__name__} only supports '
+                f'hidden states with more than 2 dimensions, where the last '
+                f'dimension is the number of state size and the other dimensions '
+                f'are the hidden shape. '
+                f'But we got {value.ndim} dimensions.'
+            )
+        name2index = {str(i): i for i in range(value.shape[-1])}
+        return value, name2index
+
+    def get_value(self, item: int | str) -> ArrayLike:
+        """
+        Get the value of the hidden state with the item.
+
+        Args:
+            item: int or str. The index of the hidden state.
+                - If int, the index of the hidden state.
+                - If str, the name of the hidden state.
+        Returns:
+            The value of the hidden state.
+        """
+        if isinstance(item, int):
+            assert item < self.value.shape[-1], (f'Index {item} out of range. '
+                                                 f'The maximum index is {self.value.shape[-1] - 1}.')
+            return self.value[..., item]
+        elif isinstance(item, str):
+            assert item in self.name2index, (f'Hidden state name {item} not found. '
+                                             f'Please check the hidden state names.')
+            index = self.name2index[item]
+            return self.value[..., index]
+        else:
+            raise TypeError(
+                f'Currently, {self.__class__.__name__} only supports '
+                f'int or str for getting the hidden state. '
+                f'But we got {type(item)}.'
+            )
+
+    def set_value(
+        self,
+        val: Dict[int | str, ArrayLike] | Sequence[ArrayLike]
+    ) -> None:
+        """
+        Set the value of the hidden state with the specified item.
+
+        This method updates the hidden state values based on the provided dictionary or sequence.
+        The values are set according to the indices or names specified in the input.
+
+        Parameters
+        ----------
+        val (Dict[int | str, ArrayLike] | Sequence[ArrayLike]):
+            A dictionary or sequence containing the new values for the hidden states.
+            - If a dictionary, keys can be integers (indices) or strings (names) of the hidden states.
+            - If a sequence, it is converted to a dictionary with indices as keys.
+
+        Returns
+        -------
+        None: This method does not return any value. It updates the hidden state values in place.
+        """
+        if isinstance(val, (tuple, list)):
+            val = {i: v for i, v in enumerate(val)}
+        assert isinstance(val, dict), (
+            f'Currently, {self.__class__.__name__}.set_value() only supports '
+            f'dictionary of hidden states. But we got {type(val)}.'
+        )
+        indices = []
+        values = []
+        for k, v in val.items():
+            if isinstance(k, str):
+                k = self.name2index[k]
+            assert isinstance(k, int), (
+                f'Key {k} should be int or str. '
+                f'But we got {type(k)}.'
+            )
+            assert v.shape == self.varshape, (
+                f'The shape of the hidden state should be {self.varshape}. '
+                f'But we got {v.shape}.'
+            )
+            indices.append(k)
+            values.append(v)
+        values = u.math.stack(values, axis=-1)
+        self.value = self.value.at[..., indices].set(values)
+
+
+class HiddenTreeState(HiddenGroupState):
+    """
+    A pytree of multiple hidden states for eligibility trace-based learning.
+
+    .. note::
+
+        The value in this state class behaves likes a dictionary/sequence of hidden states.
+        However, the state is actually stored as a single dimensionless array.
+
+    There are two ways to define the hidden states.
+
+    1. The first is to define a sequence of hidden states.
+
+    .. code-block:: python
+
+        import brainunit as u
+        value = [np.random.randn(10, 10) * u.mV,
+                 np.random.randn(10, 10) * u.mA,
+                 np.random.randn(10, 10) * u.mS]
+        state = HiddenTreeState(value)
+
+    Then, you can retrieve the hidden state value with the following method.
+
+    .. code-block:: python
+
+        state.get_value(0)  # get the first hidden state
+        # or
+        state.get_value('0')  # get the hidden state with the name '0'
+
+    You can write the hidden state value with the following method.
+
+    .. code-block:: python
+
+        state.set_value({0: np.random.randn(10, 10) * u.mV})  # set the first hidden state
+        # or
+        state.set_value({'1': np.random.randn(10, 10) * u.mA})  # set the hidden state with the name '1'
+        # or
+        state.set_value([np.random.randn(10, 10) * u.mV,
+                         np.random.randn(10, 10) * u.mA,
+                         np.random.randn(10, 10) * u.mS])  # set all hidden state value
+        # or
+        state.set_value({
+            0: np.random.randn(10, 10) * u.mV,
+            1: np.random.randn(10, 10) * u.mA,
+            2: np.random.randn(10, 10) * u.mS
+        })  # set all hidden state value
+
+    2. The second is to define a dictionary of hidden states.
+
+    .. code-block:: python
+
+        import brainunit as u
+        value = {'v': np.random.randn(10, 10) * u.mV,
+                 'i': np.random.randn(10, 10) * u.mA,
+                 'g': np.random.randn(10, 10) * u.mS}
+        state = HiddenTreeState(value)
+
+    Then, you can retrieve the hidden state value with the following method.
+
+    .. code-block:: python
+
+        state.get_value('v')  # get the hidden state with the name 'v'
+        # or
+        state.get_value('i')  # get the hidden state with the name 'i'
+
+    You can write the hidden state value with the following method.
+
+    .. code-block:: python
+
+        state.set_value({'v': np.random.randn(10, 10) * u.mV})  # set the hidden state with the name 'v'
+        # or
+        state.set_value({'i': np.random.randn(10, 10) * u.mA})  # set the hidden state with the name 'i'
+        # or
+        state.set_value([np.random.randn(10, 10) * u.mV,
+                         np.random.randn(10, 10) * u.mA,
+                         np.random.randn(10, 10) * u.mS])  # set all hidden state value
+        # or
+        state.set_value({
+            'v': np.random.randn(10, 10) * u.mV,
+            'g': np.random.randn(10, 10) * u.mA,
+            'i': np.random.randn(10, 10) * u.mS
+        })  # set all hidden state value
+
+    .. note::
+
+        Avoid using ``HiddenTreeState.value`` to get the state value, or
+        ``HiddenTreeState.value =`` to assign the state value.
+
+        Instead, use ``HiddenTreeState.get_value()`` and ``HiddenTreeState.set_value()``.
+        This is because ``.value`` loss hidden state units and other information,
+        and it is only dimensionless data.
+
+        This design aims to ensure that any etrace hidden state has only one array.
+
+
+    Args:
+        value: The values of the hidden states.
+    """
+
+    __module__ = 'brainstate'
+    value: ArrayLike
+
+    def __init__(
+        self,
+        value: Dict[str, ArrayLike] | Sequence[ArrayLike],
+    ):
+        value, name2unit, name2index = self._check_value(value)
+        self.name2unit: Dict[str, u.Unit] = name2unit
+        self.name2index: Dict[str, int] = name2index
+        self.index2unit: Dict[int, u.Unit] = {i: v for i, v in enumerate(name2unit.values())}
+        self.index2name: Dict[int, str] = {v: k for k, v in name2index.items()}
+        ShortTermState.__init__(self, value)
+
+    @property
+    def varshape(self) -> Tuple[int, ...]:
+        """
+        The shape of each hidden state variable.
+        """
+        return self.value.shape[:-1]
+
+    @property
+    def num_state(self) -> int:
+        """
+        The number of hidden states.
+        """
+        assert self.value.shape[-1] == len(self.name2index), (
+            f'The number of hidden states '
+            f'is not equal to the number of hidden state names.'
+        )
+        return self.value.shape[-1]
+
+    def _check_value(
+        self,
+        value: dict | Sequence
+    ) -> Tuple[ArrayLike, Dict[str, u.Unit], Dict[str, int]]:
+        """
+        Validates and processes the input value to ensure it conforms to the expected format
+        and structure for hidden states.
+
+        This function checks if the input value is a dictionary or sequence of hidden states,
+        verifies that all hidden states have the same shape, and extracts units and indices
+        for each hidden state.
+
+        Args:
+            value (dict | Sequence): A dictionary or sequence representing hidden states.
+                - If a sequence, it is converted to a dictionary with string indices as keys.
+                - Each hidden state should be a numpy.ndarray, jax.Array, or brainunit.Quantity.
+
+        Returns:
+            Tuple[ArrayLike, Dict[str, u.Unit], Dict[str, int]]:
+                - A stacked array of hidden state magnitudes.
+                - A dictionary mapping hidden state names to their units.
+                - A dictionary mapping hidden state names to their indices.
+
+        Raises:
+            TypeError: If any hidden state is not a numpy.ndarray, jax.Array, or brainunit.Quantity.
+            ValueError: If hidden states do not have the same shape.
+        """
+        if isinstance(value, (tuple, list)):
+            value = {str(i): v for i, v in enumerate(value)}
+        assert isinstance(value, dict), (
+            f'Currently, {self.__class__.__name__} only supports '
+            f'dictionary/sequence of hidden states. But we got {type(value)}.'
+        )
+        shapes = []
+        for k, v in value.items():
+            if not isinstance(v, (np.ndarray, jax.Array, u.Quantity)):
+                raise TypeError(
+                    f'Currently, {self.__class__.__name__} only supports '
+                    f'numpy.ndarray, jax.Array or brainunit.Quantity. '
+                    f'But we got {type(v)} for key {k}.'
+                )
+            shapes.append(v.shape)
+        if len(set(shapes)) > 1:
+            info = {k: v.shape for k, v in value.items()}
+            raise ValueError(
+                f'Currently, {self.__class__.__name__} only supports '
+                f'hidden states with the same shape. '
+                f'But we got {info}.'
+            )
+        name2unit = {k: u.get_unit(v) for k, v in value.items()}
+        name2index = {k: i for i, k in enumerate(value.keys())}
+        value = u.math.stack([u.get_magnitude(v) for v in value.values()], axis=-1)
+        return value, name2unit, name2index
+
+    def get_value(self, item: str | int) -> ArrayLike:
+        """
+        Get the value of the hidden state with the key.
+
+        Args:
+            item: The key of the hidden state.
+                - If int, the index of the hidden state.
+                - If str, the name of the hidden state.
+        """
+        if isinstance(item, int):
+            assert item < self.value.shape[-1], (f'Index {item} out of range. '
+                                                 f'The maximum index is {self.value.shape[-1] - 1}.')
+            val = self.value[..., item]
+        elif isinstance(item, str):
+            assert item in self.name2index, (f'Hidden state name {item} not found. '
+                                             f'Please check the hidden state names.')
+            item = self.name2index[item]
+            val = self.value[..., item]
+        else:
+            raise TypeError(
+                f'Currently, {self.__class__.__name__} only supports '
+                f'int or str for getting the hidden state. '
+                f'But we got {type(item)}.'
+            )
+        if self.index2unit[item].dim.is_dimensionless:
+            return val
+        else:
+            return val * self.index2unit[item]
+
+    def set_value(
+        self,
+        val: Dict[int | str, ArrayLike] | Sequence[ArrayLike]
+    ) -> None:
+        """
+        Set the value of the hidden state with the specified item.
+
+        This method updates the hidden state values based on the provided dictionary or sequence.
+        The values are set according to the indices or names specified in the input.
+
+        Parameters
+        ----------
+        val (Dict[int | str, ArrayLike] | Sequence[ArrayLike]):
+            A dictionary or sequence containing the new values for the hidden states.
+            - If a dictionary, keys can be integers (indices) or strings (names) of the hidden states.
+            - If a sequence, it is converted to a dictionary with indices as keys.
+
+        Returns
+        -------
+        None: This method does not return any value. It updates the hidden state values in place.
+        """
+        if isinstance(val, (tuple, list)):
+            val = {i: v for i, v in enumerate(val)}
+        assert isinstance(val, dict), (f'Currently, {self.__class__.__name__}.set_value() only supports '
+                                       f'dictionary of hidden states. But we got {type(val)}.')
+        indices = []
+        values = []
+        for index, v in val.items():
+            if isinstance(index, str):
+                index = self.name2index[index]
+            assert isinstance(index, int), (f'Key {index} should be int or str. '
+                                            f'But we got {type(index)}.')
+            assert v.shape == self.varshape, (f'The shape of the hidden state should be {self.varshape}. '
+                                              f'But we got {v.shape}.')
+            indices.append(index)
+            values.append(u.Quantity(v).to(self.index2unit[index]).mantissa)
+        if len(indices) == 0:
+            raise ValueError(
+                f'No hidden state is set. Please check the hidden state names or indices.'
+            )
+        if len(indices) == 1:
+            indices = indices[0]
+            values = values[0]
+        else:
+            indices = np.asarray(indices)
+            values = u.math.stack(values, axis=-1)
+        self.value = self.value.at[..., indices].set(values)
+
 
 class ParamState(LongTermState):
     """
@@ -967,17 +1472,17 @@ class StateTraceStack(Generic[A]):
     def new_arg(self, state: State) -> None:
         """
         Apply a transformation to the value of a given state using a predefined function.
-    
+
         This method is used internally to transform the value of a state during tracing.
         If a transformation function (``_jax_trace_new_arg``) is defined, it applies this
         function to each element of the state's value using JAX's tree mapping.
-    
+
         Args:
             state (State): The State object whose value needs to be transformed.
-    
+
         Returns:
             None: This function modifies the state in-place and doesn't return anything.
-    
+
         Note:
             This method is intended for internal use and relies on the presence of
             a ``_jax_trace_new_arg`` function, which should be set separately.
@@ -997,18 +1502,18 @@ class StateTraceStack(Generic[A]):
     def read_its_value(self, state: State) -> None:
         """
         Record that a state's value has been read during tracing.
-    
+
         This method marks the given state as having been read in the current
         tracing context. If the state hasn't been encountered before, it adds
         it to the internal tracking structures and applies any necessary
         transformations via the new_arg method.
-    
+
         Args:
             state (State): The State object whose value is being read.
-    
+
         Returns:
             None
-    
+
         Note:
             This method updates the internal tracking of state accesses.
             It doesn't actually read or return the state's value.
@@ -1052,11 +1557,11 @@ class StateTraceStack(Generic[A]):
     ) -> Sequence[PyTree] | Tuple[Sequence[PyTree], Sequence[PyTree]]:
         """
         Retrieve the values of all states in the StateTraceStack.
-    
+
         This method returns the values of all states, optionally separating them
         into written and read states, and optionally replacing values with None
         for states that weren't accessed in a particular way.
-    
+
         Args:
             separate (bool, optional): If True, separate the values into written
                 and read states. If False, return all values in a single sequence.
@@ -1064,7 +1569,7 @@ class StateTraceStack(Generic[A]):
             replace (bool, optional): If True and separate is True, replace values
                 with None for states that weren't written/read. If False, only
                 include values for states that were written/read. Defaults to False.
-    
+
         Returns:
             Sequence[PyTree] | Tuple[Sequence[PyTree], Sequence[PyTree]]:
                 If separate is False:
@@ -1075,7 +1580,7 @@ class StateTraceStack(Generic[A]):
                     - The second sequence contains values of read states.
                     If replace is True, these sequences will have None for
                     states that weren't written/read respectively.
-    
+
         """
         if separate:
             if replace:
@@ -1121,19 +1626,19 @@ class StateTraceStack(Generic[A]):
     def merge(self, *traces) -> 'StateTraceStack':
         """
         Merge other state traces into the current ``StateTraceStack``.
-    
+
         This method combines the states, their write status, and original values from
         other ``StateTraceStack`` instances into the current one. If a state from another
         trace is not present in the current trace, it is added. If a state is already
         present, its write status is updated if necessary.
-    
+
         Args:
             *traces: Variable number of ``StateTraceStack`` instances to be merged into
                      the current instance.
-    
+
         Returns:
             StateTraceStack: The current ``StateTraceStack`` instance with merged traces.
-    
+
         Note:
             This method modifies the current ``StateTraceStack`` in-place and also returns it.
         """
@@ -1152,16 +1657,16 @@ class StateTraceStack(Generic[A]):
     def get_read_states(self, replace_writen: bool = False) -> Tuple[State, ...]:
         """
         Retrieve the states that were read during the function execution.
-    
+
         This method returns the states that were accessed (read from) during
         the traced function's execution. It can optionally replace written
         states with None.
-    
+
         Args:
             replace_writen (bool, optional): If True, replace written states with None
                 in the returned tuple. If False, exclude written states entirely from
                 the result. Defaults to False.
-    
+
         Returns:
             Tuple[State, ...]: A tuple containing the read states.
                 If replace_writen is True, the tuple will have the same length as the
@@ -1177,15 +1682,15 @@ class StateTraceStack(Generic[A]):
     def get_read_state_values(self, replace_writen: bool = False) -> Tuple[PyTree, ...]:
         """
         Retrieve the values of states that were read during the function execution.
-    
+
         This method returns the values of states that were accessed (read from) during
         the traced function's execution. It can optionally replace written states with None.
-    
+
         Args:
             replace_writen (bool, optional): If True, replace the values of written
                 states with None in the returned tuple. If False, exclude written
                 states entirely from the result. Defaults to False.
-    
+
         Returns:
             Tuple[PyTree, ...]: A tuple containing the values of read states.
                 If replace_writen is True, the tuple will have the same length as the
@@ -1204,16 +1709,16 @@ class StateTraceStack(Generic[A]):
     def get_write_states(self, replace_read: bool = False) -> Tuple[State, ...]:
         """
         Retrieve the states that were written during the function execution.
-    
+
         This method returns the states that were modified (written to) during
         the traced function's execution. It can optionally replace unwritten (read-only)
         states with None.
-    
+
         Args:
             replace_read (bool, optional): If True, replace read-only states with None
                 in the returned tuple. If False, exclude read-only states entirely from
                 the result. Defaults to False.
-    
+
         Returns:
             Tuple[State, ...]: A tuple containing the written states.
                 If replace_read is True, the tuple will have the same length as the
@@ -1229,23 +1734,23 @@ class StateTraceStack(Generic[A]):
     def get_write_state_values(self, replace_read: bool = False) -> Tuple[PyTree, ...]:
         """
         Retrieve the values of states that were written during the function execution.
-    
+
         This method returns the values of states that were modified (written to) during
         the traced function's execution. It can optionally replace unwritten (read-only)
         states with None.
-    
+
         Args:
             replace_read (bool, optional): If True, replace the values of read-only
                 states with None in the returned tuple. If False, exclude read-only
                 states entirely from the result. Defaults to False.
-    
+
         Returns:
             Tuple[PyTree, ...]: A tuple containing the values of written states.
                 If replace_read is True, the tuple will have the same length as the
                 total number of states, with None for read-only states.
                 If replace_read is False, the tuple will only contain values of
                 written states.
-    
+
         """
         if replace_read:
             return tuple([st.value if been_writen else None for st, been_writen in zip(self.states, self.been_writen)])