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

brainstate/random/{_rand_state_test.py → _state_test.py}

@@ -19,7 +19,7 @@ import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 
-from brainstate.random._rand_state import RandomState, DEFAULT, _formalize_key, _size2shape, _check_py_seq
+from brainstate.random._state import RandomState, DEFAULT, formalize_key, _size2shape, _check_py_seq
 
 
 class TestRandomStateInitialization(unittest.TestCase):
@@ -437,32 +437,32 @@ class TestUtilityFunctions(unittest.TestCase):
 
     def test_formalize_key_with_int(self):
         """Test _formalize_key with integer."""
-        key = _formalize_key(42)
+        key = formalize_key(42)
         expected = jr.PRNGKey(42)
         np.testing.assert_array_equal(key, expected)
 
     def test_formalize_key_with_array(self):
         """Test _formalize_key with array."""
         input_key = jr.PRNGKey(123)
-        key = _formalize_key(input_key)
+        key = formalize_key(input_key, True)
         np.testing.assert_array_equal(key, input_key)
 
     def test_formalize_key_with_uint32_array(self):
         """Test _formalize_key with uint32 array."""
         input_array = np.array([123, 456], dtype=np.uint32)
-        key = _formalize_key(input_array)
+        key = formalize_key(input_array)
         np.testing.assert_array_equal(key, input_array)
 
     def test_formalize_key_invalid_input(self):
         """Test _formalize_key with invalid input."""
         with self.assertRaises(TypeError):
-            _formalize_key("invalid")
+            formalize_key("invalid")
 
         with self.assertRaises(TypeError):
-            _formalize_key(np.array([1, 2, 3], dtype=np.uint32))  # Wrong size
+            formalize_key(np.array([1, 2, 3], dtype=np.uint32))  # Wrong size
 
         with self.assertRaises(TypeError):
-            _formalize_key(np.array([1, 2], dtype=np.int32))  # Wrong dtype
+            formalize_key(np.array([1, 2], dtype=np.int32))  # Wrong dtype
 
     def test_size2shape(self):
         """Test _size2shape function."""

brainstate/transform/__init__.py

@@ -22,8 +22,8 @@ from ._conditions import *
 from ._conditions import __all__ as _conditions_all
 from ._error_if import *
 from ._error_if import __all__ as _error_if_all
-from ._eval_shape import *
-from ._eval_shape import __all__ as _eval_shape_all
+from ._find_state import *
+from ._find_state import __all__ as _find_all
 from ._jit import *
 from ._jit import __all__ as _jit_all
 from ._loop_collect_return import *
@@ -36,24 +36,21 @@ from ._mapping import *
 from ._mapping import __all__ as _mapping_all
 from ._progress_bar import *
 from ._progress_bar import __all__ as _progress_bar_all
-from ._random import *
-from ._random import __all__ as _random_all
 from ._unvmap import *
 from ._unvmap import __all__ as _unvmap_all
 
-__all__ = _ad_checkpoint_all + _autograd_all + _conditions_all + _error_if_all
-__all__ += _eval_shape_all + _jit_all + _loop_collect_return_all + _loop_no_collection_all
-__all__ += _make_jaxpr_all + _mapping_all + _progress_bar_all + _random_all + _unvmap_all
+__all__ = _ad_checkpoint_all + _autograd_all + _conditions_all + _error_if_all + _find_all
+__all__ += _jit_all + _loop_collect_return_all + _loop_no_collection_all
+__all__ += _make_jaxpr_all + _mapping_all + _progress_bar_all + _unvmap_all
+del _find_all
 del _ad_checkpoint_all
 del _autograd_all
 del _conditions_all
 del _error_if_all
-del _eval_shape_all
 del _jit_all
 del _loop_collect_return_all
 del _loop_no_collection_all
 del _make_jaxpr_all
 del _mapping_all
 del _progress_bar_all
-del _random_all
 del _unvmap_all

brainstate/transform/_conditions.py

@@ -218,7 +218,7 @@ def switch(index, branches: Sequence[Callable], *operands):
     index = jax.lax.clamp(lo, index, hi)
 
     # not jit
-    if jax.config.jax_disable_jit and isinstance(jax.core.core.get_aval(index), jax.core.ConcreteArray):
+    if jax.config.jax_disable_jit and not isinstance(to_concrete_aval(index), Tracer):
         return branches[int(index)](*operands)
 
     # evaluate jaxpr
@@ -311,6 +311,6 @@ def ifelse(conditions, branches, *operands, check_cond: bool = True):
     # format index
     conditions = jnp.asarray(conditions, np.int32)
     if check_cond:
-        jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {}.", err_arg=conditions)
+        jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {c}.", c=conditions)
     index = jnp.where(conditions, size=1, fill_value=len(conditions) - 1)[0][0]
     return switch(index, branches, *operands)

brainstate/transform/_find_state.py

@@ -0,0 +1,200 @@
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Callable, Dict, Hashable, Literal, Sequence, Any
+
+from brainstate._state import State
+from brainstate.util.filter import Filter, to_predicate
+from ._make_jaxpr import StatefulFunction
+
+__all__ = [
+    'StateFinder'
+]
+
+
+class StateFinder:
+    """
+    Discover :class:`~brainstate.State` instances touched by a callable.
+
+    ``StateFinder`` wraps a function in :class:`~brainstate.transform.StatefulFunction`
+    and exposes the collection of states the function reads or writes. The finder
+    can filter states by predicates, request only read or write states, and return
+    the result in several convenient formats.
+
+    Parameters
+    ----------
+    fn : callable
+        Function whose state usage should be inspected.
+    filter : Filter, optional
+        Predicate (see :mod:`brainstate.util.filter`) used to select states.
+    usage : {'all', 'read', 'write', 'both'}, default 'all'
+        Portion of the state trace to return. ``'both'`` returns a mapping with
+        separate read and write entries.
+    return_type : {'dict', 'list', 'tuple'}, default 'dict'
+        Controls the container type returned for the selected states. When
+        ``usage='both'``, the same container type is used for the ``'read'`` and
+        ``'write'`` entries.
+    key_fn : callable, optional
+        Callable ``key_fn(index, state)`` that produces dictionary keys when
+        ``return_type='dict'``. Defaults to using the positional index so existing
+        code continues to receive integer keys.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> param = brainstate.ParamState(jnp.ones(()), name='weights')
+        >>> bias = brainstate.ParamState(jnp.zeros(()), name='bias')
+        >>>
+        >>> def forward(x):
+        ...     _ = bias.value  # read-only
+        ...     param.value = param.value * x  # read + write
+        ...     return param.value + bias.value
+        >>>
+        >>> finder = brainstate.transform.StateFinder(
+        ...     forward,
+        ...     filter=brainstate.ParamState,
+        ...     usage='both',
+        ...     key_fn=lambda i, st: st.name or i,
+        ... )
+        >>> finder(2.0)['write']  # doctest: +ELLIPSIS
+        {'weights': ParamState(...}
+
+    Notes
+    -----
+    The underlying :class:`StatefulFunction` is cached, so subsequent calls with
+    compatible arguments will reuse the compiled trace.
+    """
+
+    _VALID_USAGE: tuple[str, ...] = ('all', 'read', 'write', 'both')
+    _VALID_RETURN_TYPE: tuple[str, ...] = ('dict', 'list', 'tuple')
+
+    def __init__(
+        self,
+        fn: Callable,
+        filter: Filter = None,
+        *,
+        usage: Literal['all', 'read', 'write', 'both'] = 'all',
+        return_type: Literal['dict', 'list', 'tuple'] = 'dict',
+        key_fn: Callable[[int, State], Hashable] | None = None,
+    ) -> None:
+        if usage not in self._VALID_USAGE:
+            raise ValueError(f"Invalid usage '{usage}'. Expected one of {self._VALID_USAGE}.")
+        if return_type not in self._VALID_RETURN_TYPE:
+            raise ValueError(
+                f"Invalid return_type '{return_type}'. Expected one of {self._VALID_RETURN_TYPE}."
+            )
+
+        self.fn = fn
+        self._usage = usage
+        self._return_type = return_type
+        self._key_fn = key_fn if key_fn is not None else self._default_key_fn
+        self._filter = to_predicate(filter) if filter is not None else None
+        self.stateful_fn = StatefulFunction(self.fn)
+
+    def __call__(self, *args, **kwargs):
+        """
+        Invoke :meth:`find` to retrieve states touched by ``fn``.
+        """
+        return self.find(*args, **kwargs)
+
+    def find(self, *args, **kwargs):
+        """
+        Execute the wrapped function symbolically and return the selected states.
+
+        Parameters
+        ----------
+        *args, **kwargs
+            Arguments forwarded to ``fn`` to determine the state trace.
+
+        Returns
+        -------
+        Any
+            Container holding the requested states as configured by ``usage`` and
+            ``return_type``.
+        """
+        if self._usage == 'both':
+            read_states = self._collect_states('read', *args, **kwargs)
+            write_states = self._collect_states('write', *args, **kwargs)
+            return {
+                'read': self._format_states(read_states),
+                'write': self._format_states(write_states),
+            }
+
+        states = self._collect_states(self._usage, *args, **kwargs)
+        return self._format_states(states)
+
+    def _collect_states(self, usage: str, *args, **kwargs) -> Sequence[State]:
+        usage_map = {
+            'all': self.stateful_fn.get_states,
+            'read': self.stateful_fn.get_read_states,
+            'write': self.stateful_fn.get_write_states,
+        }
+        collector = usage_map.get(usage)
+        if collector is None:
+            raise ValueError(f"Unsupported usage '{usage}'.")
+        states = list(collector(*args, **kwargs))
+        if self._filter is not None:
+            states = [st for st in states if self._filter(tuple(), st)]
+        return states
+
+    def _format_states(self, states: Sequence[State]):
+        if self._return_type == 'list':
+            return list(states)
+        if self._return_type == 'tuple':
+            return tuple(states)
+        return self._states_to_dict(states)
+
+    def _states_to_dict(self, states: Sequence[State]) -> Dict[Hashable, State]:
+        result: Dict[Hashable, State] = {}
+        used_keys: set[Hashable] = set()
+        for idx, state in enumerate(states):
+            key = self._key_fn(idx, state)
+            key = self._ensure_hashable(key)
+            key = self._ensure_unique_key(key, idx, state, used_keys)
+            result[key] = state
+            used_keys.add(key)
+        return result
+
+    @staticmethod
+    def _default_key_fn(idx: int, state: State) -> Hashable:
+        return idx
+
+    @staticmethod
+    def _ensure_hashable(key: Any) -> Hashable:
+        if key is None:
+            return None
+        try:
+            hash(key)
+        except TypeError:
+            return str(key)
+        return key
+
+    @staticmethod
+    def _ensure_unique_key(key: Hashable, idx: int, state: State, used: set[Hashable]) -> Hashable:
+        if key is None or key in used:
+            base_name = getattr(state, 'name', None)
+            base = base_name if base_name not in (None, '') else f"state_{idx}"
+            candidate = base
+            suffix = 1
+            while candidate in used:
+                candidate = f"{base}_{suffix}"
+                suffix += 1
+            return candidate
+        return key
+

brainstate/transform/_find_state_test.py

@@ -0,0 +1,84 @@
+import unittest
+
+import jax.numpy as jnp
+
+import brainstate as bst
+from brainstate.transform import StateFinder
+
+
+class TestStateFinder(unittest.TestCase):
+    def test_default_dictionary_output(self):
+        read_state = bst.State(jnp.array(0.0), name='read_state')
+        param_state = bst.ParamState(jnp.array(1.0), name='param_state')
+
+        def fn(scale):
+            _ = read_state.value
+            param_state.value = param_state.value * scale
+            return param_state.value + _
+
+        finder = StateFinder(fn)
+        result = finder(2.0)
+        self.assertEqual(len(result), 2)
+        self.assertEqual(set(result.values()), {read_state, param_state})
+
+    def test_filter_and_usage_read(self):
+        buffer_state = bst.State(jnp.array(1.0), name='buffer')
+        param_state = bst.ParamState(jnp.array(3.0), name='param')
+
+        def fn(offset):
+            _ = buffer_state.value
+            param_state.value = param_state.value + offset
+            return param_state.value
+
+        read_finder = StateFinder(fn, usage='read', return_type='list')
+        reads = read_finder(1.0)
+        self.assertEqual(reads, [buffer_state])
+
+        param_finder = StateFinder(fn, filter=bst.ParamState, usage='all')
+        param_states = param_finder(1.0)
+        self.assertEqual(list(param_states.values()), [param_state])
+
+    def test_usage_write_with_custom_key(self):
+        param_state = bst.ParamState(jnp.array(5.0), name='param')
+
+        def fn(scale):
+            param_state.value = param_state.value * scale
+            return param_state.value
+
+        finder = StateFinder(fn, usage='write', return_type='dict', key_fn=lambda idx, st: f"w_{idx}")
+        write_states = finder(2.0)
+        self.assertIn('w_0', write_states)
+        self.assertIs(write_states['w_0'], param_state)
+
+    def test_usage_both_returns_separated_collections(self):
+        read_state = bst.State(jnp.array(2.0), name='read')
+        write_state = bst.ParamState(jnp.array(4.0), name='write')
+
+        def fn(delta):
+            _ = read_state.value
+            write_state.value = write_state.value + delta
+            return write_state.value
+
+        finder = StateFinder(fn, usage='both', return_type='tuple')
+        result = finder(1.5)
+        self.assertEqual(set(result.keys()), {'read', 'write'})
+        self.assertEqual(result['read'], (read_state,))
+        self.assertEqual(result['write'], (write_state,))
+
+    def test_duplicate_names_are_disambiguated(self):
+        first = bst.State(jnp.array(0.0), name='dup')
+        second = bst.State(jnp.array(1.0), name='dup')
+
+        def fn():
+            _ = first.value
+            _ = second.value
+            return None
+
+        finder = StateFinder(fn)
+        states = finder()
+        self.assertEqual(len(states), 2)
+        self.assertEqual(set(states.values()), {first, second})
+
+
+if __name__ == "__main__":
+    unittest.main()