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

brainstate/graph/_operation.py

@@ -1262,11 +1262,7 @@ def treefy_states(
     if len(filters) == 0:
         return state_mapping
     else:
-        state_mappings = state_mapping.filter(*filters)
-        if len(filters) == 1:
-            return state_mappings[0]
-        else:
-            return state_mappings
+        return state_mapping.filter(*filters)
 
 
 def _graph_update_dynamic(node: Any, state: Mapping) -> None:

brainstate/mixin.py

@@ -380,6 +380,20 @@ class ParamDescriber(metaclass=NoSubclassMeta):
         merged_kwargs = {**self.kwargs, **kwargs}
         return self.cls(*self.args, *args, **merged_kwargs)
 
+    def __repr__(self):
+        """
+        Return a string representation of the ParamDescriber.
+
+        Returns
+        -------
+        str
+            A string showing the class and stored parameters.
+        """
+        args_str = ', '.join(repr(a) for a in self.args)
+        kwargs_str = ', '.join(f'{k}={v!r}' for k, v in self.kwargs.items())
+        all_params = ', '.join(filter(None, [args_str, kwargs_str]))
+        return f'ParamDescriber({self.cls.__name__}({all_params}))'
+
     def init(self, *args, **kwargs):
         """
         Alias for __call__, explicitly named for clarity.

brainstate/nn/__init__.py

@@ -87,43 +87,52 @@ del (
 
 # Deprecated names that redirect to brainpy
 _DEPRECATED_NAMES = {
-    'SpikeTime': 'brainpy.SpikeTime',
-    'PoissonSpike': 'brainpy.PoissonSpike',
-    'PoissonEncoder': 'brainpy.PoissonEncoder',
-    'PoissonInput': 'brainpy.PoissonInput',
-    'poisson_input': 'brainpy.poisson_input',
-    'Neuron': 'brainpy.Neuron',
-    'IF': 'brainpy.IF',
-    'LIF': 'brainpy.LIF',
-    'LIFRef': 'brainpy.LIFRef',
-    'ALIF': 'brainpy.ALIF',
-    'LeakyRateReadout': 'brainpy.LeakyRateReadout',
-    'LeakySpikeReadout': 'brainpy.LeakySpikeReadout',
-    'STP': 'brainpy.STP',
-    'STD': 'brainpy.STD',
-    'Synapse': 'brainpy.Synapse',
-    'Expon': 'brainpy.Expon',
-    'DualExpon': 'brainpy.DualExpon',
-    'Alpha': 'brainpy.Alpha',
-    'AMPA': 'brainpy.AMPA',
-    'GABAa': 'brainpy.GABAa',
-    'COBA': 'brainpy.COBA',
-    'CUBA': 'brainpy.CUBA',
-    'MgBlock': 'brainpy.MgBlock',
-    'SynOut': 'brainpy.SynOut',
-    'AlignPostProj': 'brainpy.AlignPostProj',
-    'DeltaProj': 'brainpy.DeltaProj',
-    'CurrentProj': 'brainpy.CurrentProj',
-    'align_pre_projection': 'brainpy.align_pre_projection',
-    'Projection': 'brainpy.Projection',
-    'SymmetryGapJunction': 'brainpy.SymmetryGapJunction',
-    'AsymmetryGapJunction': 'brainpy.AsymmetryGapJunction',
+    'SpikeTime': 'brainpy.state.SpikeTime',
+    'PoissonSpike': 'brainpy.state.PoissonSpike',
+    'PoissonEncoder': 'brainpy.state.PoissonEncoder',
+    'PoissonInput': 'brainpy.state.PoissonInput',
+    'poisson_input': 'brainpy.state.poisson_input',
+    'Neuron': 'brainpy.state.Neuron',
+    'IF': 'brainpy.state.IF',
+    'LIF': 'brainpy.state.LIF',
+    'LIFRef': 'brainpy.state.LIFRef',
+    'ALIF': 'brainpy.state.ALIF',
+    'LeakyRateReadout': 'brainpy.state.LeakyRateReadout',
+    'LeakySpikeReadout': 'brainpy.state.LeakySpikeReadout',
+    'STP': 'brainpy.state.STP',
+    'STD': 'brainpy.state.STD',
+    'Synapse': 'brainpy.state.Synapse',
+    'Expon': 'brainpy.state.Expon',
+    'DualExpon': 'brainpy.state.DualExpon',
+    'Alpha': 'brainpy.state.Alpha',
+    'AMPA': 'brainpy.state.AMPA',
+    'GABAa': 'brainpy.state.GABAa',
+    'COBA': 'brainpy.state.COBA',
+    'CUBA': 'brainpy.state.CUBA',
+    'MgBlock': 'brainpy.state.MgBlock',
+    'SynOut': 'brainpy.state.SynOut',
+    'AlignPostProj': 'brainpy.state.AlignPostProj',
+    'DeltaProj': 'brainpy.state.DeltaProj',
+    'CurrentProj': 'brainpy.state.CurrentProj',
+    'align_pre_projection': 'brainpy.state.align_pre_projection',
+    'Projection': 'brainpy.state.Projection',
+    'SymmetryGapJunction': 'brainpy.state.SymmetryGapJunction',
+    'AsymmetryGapJunction': 'brainpy.state.AsymmetryGapJunction',
 }
 
 
 def __getattr__(name: str):
+    import warnings
+    if name == 'DynamicsGroup':
+        warnings.warn(
+            f"'brainstate.nn.{name}' is deprecated and will be removed in a future version. "
+            f"Please use 'brainstate.nn.Module' instead.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+        return Module
+
     if name in _DEPRECATED_NAMES:
-        import warnings
         new_name = _DEPRECATED_NAMES[name]
         warnings.warn(
             f"'brainstate.nn.{name}' is deprecated and will be removed in a future version. "
@@ -133,5 +142,5 @@ def __getattr__(name: str):
         )
         # Import and return the actual brainpy object
         import brainpy
-        return getattr(brainpy, name)
+        return getattr(brainpy.state, name)
     raise AttributeError(f"module 'brainstate.nn' has no attribute '{name}'")

brainstate/nn/_collective_ops.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
+
 import warnings
 from collections.abc import Sequence, Mapping
 from typing import Callable, TypeVar, Any

brainstate/nn/_common_test.py

@@ -132,23 +132,3 @@ class TestFilterStates(unittest.TestCase):
 
         self.mock_module.states.assert_called_once_with(filter_obj)
         self.assertEqual(result, ['test1', 'test2'])
-
-    def test_filter_states_dict_filters(self):
-        """Test _filter_states with dictionary of filters.
-
-        Note: Current implementation expects dict to be iterable as tuples,
-        which suggests it's meant to be passed as a dict that yields tuples when iterated.
-        This is likely a bug - should use filters.items().
-        """
-        # Skip this test as the current implementation has a bug
-        self.skipTest("Current implementation has a bug in dict iteration")
-
-    def test_filter_states_dict_invalid_axis(self):
-        """Test _filter_states with non-integer axis in dictionary."""
-        # Skip this test as the current implementation has a bug in dict iteration
-        self.skipTest("Current implementation has a bug in dict iteration")
-
-    def test_filter_states_dict_multiple_filters_same_axis(self):
-        """Test _filter_states with multiple filters for the same axis."""
-        # Skip this test as the current implementation has a bug in dict iteration
-        self.skipTest("Current implementation has a bug in dict iteration")

brainstate/nn/_delay.py

@@ -285,7 +285,7 @@ class Delay(Module):
         Returns:
             DelayAccess: An object that provides access to the delay data for the specified entry and time.
         """
-        return DelayAccess(self, delay_time, entry=entry)
+        return DelayAccess(self, *delay_time, entry=entry)
 
     def at(self, entry: str) -> ArrayLike:
         """

brainstate/nn/_dropout_test.py

@@ -146,16 +146,19 @@ class TestDropout1d(parameterized.TestCase):
 
 
 class TestDropout2d(parameterized.TestCase):
+    def setUp(self):
+        brainstate.random.seed(0)
 
     def test_dropout2d_basic(self):
         """Test basic Dropout2d functionality."""
-        dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
-        input_data = brainstate.random.randn(2, 3, 4, 5)  # (N, C, H, W)
+        with brainstate.random.seed_context(42):
+            dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
+            input_data = brainstate.random.randn(2, 3, 4, 5)  # (N, C, H, W)
 
-        with brainstate.environ.context(fit=True):
-            output_data = dropout_layer(input_data)
-            self.assertEqual(input_data.shape, output_data.shape)
-            self.assertTrue(np.any(output_data == 0))
+            with brainstate.environ.context(fit=True):
+                output_data = dropout_layer(input_data)
+                self.assertEqual(input_data.shape, output_data.shape)
+                self.assertTrue(np.any(output_data == 0))
 
     def test_dropout2d_channel_wise(self):
         """Test that Dropout2d applies dropout."""