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

brainstate/transform/_mapping_test.py

@@ -1,194 +1,104 @@
-# 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.
-# 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.
-# ==============================================================================
-
-
 import unittest
 
 import jax
 import jax.numpy as jnp
-from jax import vmap
-from jax.lax import psum, pmean, pmax
-
-import brainstate
-import brainstate.transform
-from brainstate._error import BatchAxisError
 
+import brainstate as bst
+from brainstate.transform import StatefulMapping, vmap, vmap_new_states, pmap, map as bst_map
+from brainstate.util import filter as state_filter
 
 
 class TestMap(unittest.TestCase):
-    def test_map(self):
-        for dim in [(10,), (10, 10), (10, 10, 10)]:
-            x = brainstate.random.rand(*dim)
-            r1 = brainstate.transform.map(lambda a: a + 1, x, batch_size=None)
-            r2 = brainstate.transform.map(lambda a: a + 1, x, batch_size=2)
-            r3 = brainstate.transform.map(lambda a: a + 1, x, batch_size=4)
-            r4 = brainstate.transform.map(lambda a: a + 1, x, batch_size=5)
-            true_r = x + 1
-
-            self.assertTrue(jnp.allclose(r1, true_r))
-            self.assertTrue(jnp.allclose(r2, true_r))
-            self.assertTrue(jnp.allclose(r3, true_r))
-            self.assertTrue(jnp.allclose(r4, true_r))
-
-
-class TestAxisName:
-    def test1(self):
-        def compute_stats_with_axis_name(x):
-            """Compute statistics using named axis operations"""
-            # Sum across the named axis 'batch'
-            total_sum = psum(x, axis_name='batch')
-
-            # Mean across the named axis 'batch'
-            mean_val = pmean(x, axis_name='batch')
-
-            # Max across the named axis 'batch'
-            max_val = pmax(x, axis_name='batch')
-
-            return {
-                'sum': total_sum,
-                'mean': mean_val,
-                'max': max_val,
-                'original': x
-            }
-
-        batch_data = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
-        print("Input batch data:", batch_data)
-
-        # vmap with axis name 'batch'
-        vectorized_stats_jax = jax.jit(vmap(compute_stats_with_axis_name, axis_name='batch'))
-        result_jax = vectorized_stats_jax(batch_data)
-
-        # vmap with axis name 'batch'
-        vectorized_stats = brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
-        result = vectorized_stats(batch_data)
-
-        # vmap with axis name 'batch'
-        vectorized_stats_v2 = brainstate.transform.jit(
-            brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+    def test_map_matches_vectorized(self):
+        xs = jnp.arange(6.0).reshape(6, 1)
+
+        def fn(x):
+            return x + 1.0
+
+        expected = jax.vmap(fn)(xs)
+        result = bst_map(fn, xs)
+        self.assertTrue(jnp.allclose(result, expected))
+
+    def test_map_multiple_inputs_and_batch_size(self):
+        xs = jnp.arange(5.0)
+        ys = jnp.ones_like(xs) * 2.0
+
+        def fn(a, b):
+            return a * a + b
+
+        expected = jax.vmap(fn)(xs, ys)
+        result = bst_map(fn, xs, ys, batch_size=2)
+        self.assertTrue(jnp.allclose(result, expected))
+
+
+class TestVmapIntegration(unittest.TestCase):
+    def test_decorator_batched_stateful_function(self):
+        counter = bst.ShortTermState(jnp.zeros(3))
+
+        @vmap(
+            in_axes=0,
+            out_axes=0,
+            state_in_axes={0: state_filter.OfType(bst.ShortTermState)},
+            state_out_axes={0: state_filter.OfType(bst.ShortTermState)},
         )
-        result_v2 = vectorized_stats_v2(batch_data)
-
-        for key in result_jax.keys():
-            print(f"  {key}: {result_jax[key]}")
-            assert jnp.allclose(result_jax[key], result[key]), f"Mismatch in {key}"
-            assert jnp.allclose(result_jax[key], result_v2[key]), f"Mismatch in {key}"
-
-    def test_nested_vmap(self):
-        def nested_computation(x):
-            """Computation with multiple named axes"""
-            # Sum over 'inner' axis, then mean over 'outer' axis
-            inner_sum = psum(x, axis_name='inner')
-            outer_mean = pmean(inner_sum, axis_name='outer')
-            return outer_mean
-
-        # Create 2D batch data
-        data_2d = jnp.arange(12.0).reshape(3, 4)  # Shape: [outer_batch=3, inner_batch=4]
-        print("Input 2D data shape:", data_2d.shape)
-        print("Input 2D data:\n", data_2d)
-
-        # Nested vmap: first over inner dimension, then outer dimension
-        inner_vmap = vmap(nested_computation, axis_name='inner')
-        nested_vmap = vmap(inner_vmap, axis_name='outer')
-
-        result_2d = nested_vmap(data_2d)
-        print("Result after nested vmap:", result_2d)
-
-        inner_vmap_bst = brainstate.transform.vmap(nested_computation, axis_name='inner')
-        nested_vmap_bst = brainstate.transform.vmap(inner_vmap_bst, axis_name='outer')
-        result_2d_bst = nested_vmap_bst(data_2d)
-        print("Result after nested vmap:", result_2d_bst)
-
-        assert jnp.allclose(result_2d, result_2d_bst)
-
-    def _gradient_averaging_simulation_bst(self):
-        def loss_function(params, x, y):
-            """Simple quadratic loss"""
-            pred = params * x
-            return (pred - y) ** 2
-
-        def compute_gradients_with_averaging(params, batch_x, batch_y):
-            """Compute gradients and average them across the batch"""
-            # Compute per-sample gradients
-            grad_fn = jax.grad(loss_function, argnums=0)
-            per_sample_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
-
-            # Average gradients across batch using named axis
-            def average_grads(grads):
-                return pmean(grads, axis_name='batch')
-
-            # Apply averaging with named axis
-            averaged_grads = vmap(average_grads, axis_name='batch')(per_sample_grads)
-            return averaged_grads
-
-        # Example data
-        params = 2.0
-        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
-        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
-
-        print("Parameters:", params)
-        print("Batch X:", batch_x)
-        print("Batch Y:", batch_y)
-
-        # Compute individual gradients first
-        grad_fn = jax.grad(loss_function, argnums=0)
-        individual_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
-        print("Individual gradients:", individual_grads)
-
-        # Now compute averaged gradients using axis names
-        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
-        print("Averaged gradients:", averaged_grads)
-
-        return individual_grads, averaged_grads
-
-    def _gradient_averaging_simulation_jax(self):
-        def loss_function(params, x, y):
-            """Simple quadratic loss"""
-            pred = params * x
-            return (pred - y) ** 2
-
-        def compute_gradients_with_averaging(params, batch_x, batch_y):
-            """Compute gradients and average them across the batch"""
-            # Compute per-sample gradients
-            grad_fn = jax.grad(loss_function, argnums=0)
-            per_sample_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
-
-            # Average gradients across batch using named axis
-            def average_grads(grads):
-                return pmean(grads, axis_name='batch')
-
-            # Apply averaging with named axis
-            averaged_grads = brainstate.transform.vmap(average_grads, axis_name='batch')(per_sample_grads)
-            return averaged_grads
-
-        # Example data
-        params = 2.0
-        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
-        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
-
-        print("Parameters:", params)
-        print("Batch X:", batch_x)
-        print("Batch Y:", batch_y)
-
-        # Compute individual gradients first
-        grad_fn = jax.grad(loss_function, argnums=0)
-        individual_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
-        print("Individual gradients:", individual_grads)
-
-        # Now compute averaged gradients using axis names
-        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
-        print("Averaged gradients:", averaged_grads)
-
-        return individual_grads, averaged_grads
+        def accumulate(x):
+            counter.value = counter.value + x
+            return counter.value
+
+        xs = jnp.asarray([1.0, 2.0, 3.0])
+        result = accumulate(xs)
+        self.assertTrue(jnp.allclose(result, xs))
+        self.assertTrue(jnp.allclose(counter.value, xs))
+
+    def test_vmap_partial_returns_stateful_mapping(self):
+        builder = vmap(in_axes=0, out_axes=0)
+
+        def fn(x):
+            return x * 2.0
+
+        mapped = builder(fn)
+        self.assertIsInstance(mapped, StatefulMapping)
+        xs = jnp.arange(3.0)
+        self.assertTrue(jnp.allclose(mapped(xs), xs * 2.0))
+
+
+class TestVmapNewStates(unittest.TestCase):
+    def test_new_states_are_vectorized(self):
+        @vmap_new_states(in_axes=0, out_axes=0)
+        def build(x):
+            scratch = bst.ShortTermState(jnp.array(0.0), tag='scratch')
+            scratch.value = scratch.value + x
+            return scratch.value
+
+        xs = jnp.arange(4.0)
+        result_first = build(xs)
+        result_second = build(xs)
+        self.assertTrue(jnp.allclose(result_first, xs))
+        self.assertTrue(jnp.allclose(result_second, xs))
+
+
+class TestPmapIntegration(unittest.TestCase):
+    @unittest.skipIf(jax.local_device_count() < 2, "Requires at least 2 devices")
+    def test_pmap_stateful_execution(self):
+        param = bst.ParamState(jnp.ones((4,)))
+
+        @pmap(
+            in_axes=0,
+            out_axes=0,
+            axis_name='devices',
+            state_in_axes={0: state_filter.OfType(bst.ParamState)},
+            state_out_axes={0: state_filter.OfType(bst.ParamState)},
+        )
+        def update(delta):
+            param.value = param.value + delta
+            return param.value
+
+        device_count = jax.local_device_count()
+        deltas = jnp.arange(device_count * 4.0, dtype=param.value.dtype).reshape(device_count, 4)
+        updated = update(deltas)
+        self.assertEqual(updated.shape, (device_count, 4))
+        self.assertTrue(jnp.all(updated >= 1.0))
+
 
+if __name__ == "__main__":
+    unittest.main()

{brainstate-0.2.0.dist-info → brainstate-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: brainstate
-Version: 0.2.0
+Version: 0.2.2
 Summary: A State-based Transformation System for Brain Modeling.
 Author-email: BrainState Developers <chao.brain@qq.com>
 License: Apache-2.0 license

{brainstate-0.2.0.dist-info → brainstate-0.2.2.dist-info}/RECORD

@@ -1,55 +1,55 @@
-brainstate/__init__.py,sha256=bmCIZG6xMiVlZ5qWk-d5qpkqdnrWxFhLePtpu3ff5WI,6024
+brainstate/__init__.py,sha256=XZZctdLIf_6-IJIbeTyoEHR6r3pyKF2n6MfikCffASA,5915
 brainstate/_compatible_import.py,sha256=7thV_2F0FD5AF2DETjBfmtNb_2ZQzki8NxFgC62frg0,11037
 brainstate/_compatible_import_test.py,sha256=6ka26Sa_Kk6F-Ar1HR6UaKJTHquXcUCWglgXBUOovcg,22762
 brainstate/_deprecation.py,sha256=gSh36_TWLBgQAo0gNfOzscV9ssa26k3te9y25BG6O2w,8381
-brainstate/_deprecation_test.py,sha256=5_tJ9JDhG79zLwe5MuBLsqcgl7gi-oXPgDTCujrPmz0,88325
+brainstate/_deprecation_test.py,sha256=vnmcZ7l_FmQRgsMYUA9wMiWL1ltPc01plUTK-pWKCzk,87550
 brainstate/_error.py,sha256=6A5ILy17ZMMZIjS8LkajTZBDRnwv_Qait5x__h2Levo,1522
-brainstate/_state.py,sha256=YCYKrX2xCTojIY15vtQl8WWjtJ0HNMg8y5LCUk1d_3Q,58096
-brainstate/_state_test.py,sha256=wCoWTbvARVkhNokMLLBV6sEoHSE9OJRyuUVrNDcRiG8,1589
+brainstate/_state.py,sha256=G5SCjxi42Lb1EoqaIZcyy0fmjNonwrQLLb4ujKt_99E,77867
+brainstate/_state_test.py,sha256=ptul4-BHagOIhnAvnWmb2BFyXnIfWpK0rQClLxhzLBY,39895
 brainstate/_utils.py,sha256=cmUyO9ds1etrrpV4ucp1G8mDqE15g4ZtbivblH_cD9o,1613
 brainstate/environ.py,sha256=BmQsvo1aZaMpckHXlJ45dZh9DUdnHHP9Q9JdLNCa9wA,42169
 brainstate/environ_test.py,sha256=RdVmeP7irbk3_qNjwWoy-DSdpkTRpxqAABtNGZcbB2w,42418
-brainstate/mixin.py,sha256=O58rGznYowz-hBQy7iKOhbrg3Dze7FAR8AU4CX8hUEI,44635
+brainstate/mixin.py,sha256=QoZPhBs4hGXTzTfXaF0M8XhsbcbOtT4Ylm54D7pCOSA,45124
 brainstate/mixin_test.py,sha256=6WmqJf34kT3Z5WaiCNDo3OV3ci0DIHCQp474zECDxEU,34718
 brainstate/typing.py,sha256=pYiNI-9oHpH7HfjRKYxugK03KGiamCwweagMyO0rsi4,26301
 brainstate/typing_test.py,sha256=2mmMW0uAzIo3_VXpT5Boq79BohxYkzBlHexBysFUGII,26240
 brainstate/graph/__init__.py,sha256=kGVtHAnkiWR5MqDYQU0G3AobWnioGeDqjILA--RyDz8,846
 brainstate/graph/_node.py,sha256=_XH8xx6_glsCK4KCsQnarACK8meyhCdfh3nWfUDko0k,6407
 brainstate/graph/_node_test.py,sha256=sD2DS0AhDKOU5ZQm0cYz0llnJ6D60ftNfDpztk4i8cM,18687
-brainstate/graph/_operation.py,sha256=n2HqfwPzG3f6QHish_7d8lNZRhHUyS4_YHTdPMGgdUk,54096
+brainstate/graph/_operation.py,sha256=5EN_qL08M_BIV-L3YkVP3OwlsDuNKmVV5NaCLvrl7kw,53971
 brainstate/graph/_operation_test.py,sha256=IVyrJh4io3sDgtrTEIAItGzNs2XEf7rO1rvI1r_KiII,39119
-brainstate/nn/__init__.py,sha256=oEQ81xpWppPMUMCRalMhAhUU5faYs0aKcyydwjoxGVo,4759
+brainstate/nn/__init__.py,sha256=JcaVVt01LNbAwxHymNmBUFAZwchtucyHq_sT70rvYc4,5245
 brainstate/nn/_activations.py,sha256=6jHR67obYR1lpo-imVXmfd3m_NDyU0XZb8t-pVYDvUU,26917
 brainstate/nn/_activations_test.py,sha256=Ikr8RYBaIpApVKUhY-XAWr6llEG7vYWS8YuqDHyTtBY,13438
-brainstate/nn/_collective_ops.py,sha256=UditkuCy0f9ggkFwGUXyJDoFtDdVUIQU3xmD5YOxRKg,21280
+brainstate/nn/_collective_ops.py,sha256=AHda68XjoNUyLOwxRhwQLnXr4zLBqnB8UDLAkfx2eAg,21282
 brainstate/nn/_collective_ops_test.py,sha256=8mKQkfTjfwuO7DA1i_Yr4QD1yg2ZJgVpX7zhHFo9CuQ,25600
 brainstate/nn/_common.py,sha256=UyJMJoVF9KfrToOX5Dbv-2s3CD49SsroLRfL17DLl4Q,7184
-brainstate/nn/_common_test.py,sha256=NpKQpKPPU4I2nVpz1X_bOGC5R-RyCP2kYlyYqXWoKAE,5966
+brainstate/nn/_common_test.py,sha256=xztvEPYCurvY9LiaZOtHjWcFrwJaYknA8_RLy_XHI1Y,4895
 brainstate/nn/_conv.py,sha256=3cGToc5UoGN5jp4BUHlrK6md8O_0IcWMCBXFHBQg7nE,82106
 brainstate/nn/_conv_test.py,sha256=65FlrteUxLQb8ckUUJFhaPZDsQZTlGVAp0HUGUJtt1M,30173
-brainstate/nn/_delay.py,sha256=0brSPlKOC7_nNWr1zcdjXNHJFw_I9KUX7PGyVWXENTk,22331
+brainstate/nn/_delay.py,sha256=qssPPKgWeS-vcxKJG6ewn2bQ1KvdHHBrbI6bIJytSfI,22332
 brainstate/nn/_delay_test.py,sha256=FzBb8vXfse8HEcEid83hpa6aag6oj90mtcHYsDs0DOE,10376
 brainstate/nn/_dropout.py,sha256=UotjW0PQO4gypfhtSqzkR4UVGkY0kNBncEB_poGL7Sc,22555
-brainstate/nn/_dropout_test.py,sha256=5LEN_G6t2TZAbmvJdZ75h1nyKZTHHiNT_i1W-im5wo8,21461
-brainstate/nn/_dynamics.py,sha256=V7GA1enppLRRDJCw146RI_CdwpVyRgHUYH9fu2-PJOA,44181
-brainstate/nn/_dynamics_test.py,sha256=G1VkrWgKK4CXtcKJmevxiQQ9DVeyf1zcdMsAWyCN0Jw,2356
+brainstate/nn/_dropout_test.py,sha256=87ETtMZ7Aeaa0les7Z5ICO9KVKMewm_WWd-ITCFDo_w,21589
+brainstate/nn/_dynamics.py,sha256=H83WUFWv6ZLBEn01U30qInjgfh8WAgZ5jwgXWA9TEkg,29161
+brainstate/nn/_dynamics_test.py,sha256=XwTBCXXKN8hAaLFtQjmWEOf-ry0MJZmfgx9dAju-rXA,1744
 brainstate/nn/_elementwise.py,sha256=4kKzrbKn5luwnpY8n7IeaMOtBVVie8oPHEc64hSn8-w,34858
 brainstate/nn/_elementwise_test.py,sha256=sbWlUyTB8oiu3PRHObTvmUaob99EjIGYE6k6bEpg6K8,27296
-brainstate/nn/_embedding.py,sha256=eMwbz9udm6WGVEGeHwhvCr7sJA6kMFyaZsli7lL30eg,14942
+brainstate/nn/_embedding.py,sha256=MXkGja6uCB2OsbwqArEQ8ez_GgwZ2s3DINaUnFyeZWk,14991
 brainstate/nn/_embedding_test.py,sha256=Gc0y6gHMEagaDrBJoAYQZMDTd47TYQNVrenWUwLWK_w,6242
 brainstate/nn/_event_fixedprob.py,sha256=ZEnIyjDksxtUWWG5GXLcF-RHR1S33DZTw-rN2lHKs0g,9395
 brainstate/nn/_event_fixedprob_test.py,sha256=rvTKxEzKwvctQc8-AxXjJ4p4D-if1va2O8KqQYW5nxY,3836
 brainstate/nn/_event_linear.py,sha256=d0J54Sf9zBl926BzXoy4Oc1p96h9veU3f8YhWrZLRPk,2554
 brainstate/nn/_event_linear_test.py,sha256=qzcGplDIwxTnZOs4JzD5GX_oNBYtcYFNaf3GpWo8pZY,3765
-brainstate/nn/_exp_euler.py,sha256=W5ofUPe6UK8NXO917PY9_nNXOI9chSHgBmnDHMVXHoo,8635
+brainstate/nn/_exp_euler.py,sha256=deKqWu5RvN_Jvj_QHg5j1xkKhlysBFhi3Qvo_Z-QNDo,8698
 brainstate/nn/_exp_euler_test.py,sha256=21qomGOo96YLmlEQok2hCByAcRpmqREhSr3kxmhKOm8,13014
-brainstate/nn/_linear.py,sha256=olbo9AmC35UQBxECyAQW06XeucsJfBAAMoSfwOW6c7s,24047
+brainstate/nn/_linear.py,sha256=LLWBE6eBXjpXUvx5mjdwOjFnolNRGBxNb_ZbhsYOu4M,24201
 brainstate/nn/_linear_test.py,sha256=5fHx4v4_54dH3Bsyapl2cobvVtMEu3DFR-jDKLbkJFw,17876
 brainstate/nn/_metrics.py,sha256=TgALwv6i9La4Dm1WAkWDWxvxr9rkd7CJLGFy2sOGQbQ,36481
 brainstate/nn/_metrics_test.py,sha256=XZiRndchRgEH0X8zsHzg0fsHNMxj43mnd883QChfSik,24104
-brainstate/nn/_module.py,sha256=to280ubWAP-HiCV7LknMBqOhet0UkH9Oh-PkLkual_g,12775
+brainstate/nn/_module.py,sha256=W9iXXmLxaX_QyNxukTxajHJfah7s9GqJXZRHQhZqCfU,13076
 brainstate/nn/_module_test.py,sha256=znjB7FU5evJENQ1Pqw7ZlOGC5faQe4-4VpjW60H8UWI,1414
-brainstate/nn/_normalizations.py,sha256=exdIn627ph5pVHdYx_4NIaK_f1xcLaBj-YQ0Ui22CsA,50185
+brainstate/nn/_normalizations.py,sha256=3ivJdQUhhWeMhDpzWhXhRQnu8dRR4zDzM1cOlPvPo2s,50374
 brainstate/nn/_normalizations_test.py,sha256=y5n7aTaUHRkyAAjq4Oj8dfButJC3ehm4KdUB7226Bow,23350
 brainstate/nn/_paddings.py,sha256=3u3dbRFtPSlIsLMBYZmHcYDQ6HFl0u_d-yP7ZoYcCrA,32415
 brainstate/nn/_paddings_test.py,sha256=uY9CRexf9sM5V6AzTfzm8214y4IoGqEbahy92yaKsWM,27409
@@ -61,36 +61,36 @@ brainstate/nn/_utils.py,sha256=VK-Se53e1q-Ip4AtMOZ3SUzYw8u2UllLJRLRtEFRCRE,7403
 brainstate/nn/_utils_test.py,sha256=uim2SkfNHrBZzNDvN0WOK8qeZC1kaeOd-UQDvrn_M24,14266
 brainstate/nn/init.py,sha256=7iLHrL-ZHpU-g5d0PlusaqmtkWO7X_KNr2eLx90oHrc,25656
 brainstate/nn/init_test.py,sha256=bfby6kovvhbc7CCEaohtQawUQj3w3GvaMSDAdTiP2ps,6200
-brainstate/random/__init__.py,sha256=2k5aI3GI_ftOtuWO9rosRUhyr5OjnMty82b-CTBkOUU,8383
-brainstate/random/_rand_funs.py,sha256=hBkZYTlBNHBLdHTLBecT7sja7mVItMF7TV3ao26vfj0,135339
-brainstate/random/_rand_funs_test.py,sha256=mZaSRtlXCwJP3YfF5GsFelR-YACg-uH1gHWG-coZT_k,23040
-brainstate/random/_rand_seed.py,sha256=mzSe2lsyhN_0eXhkD8gpRKmSbReQfZF0pZA89cjiiHU,24927
-brainstate/random/_rand_seed_test.py,sha256=Y2VCAkUzciDaCfYZWPe_Ewmi3MylK-WzfPA7TzorV8Q,1491
-brainstate/random/_rand_state.py,sha256=_62wIrdZKKvmtjBErB0lRCuf1vYEskFyPd4Nhp_FKDU,53591
-brainstate/random/_rand_state_test.py,sha256=0Y1kx7rvkACQiiiAHK0plePtLxYeYTzvWUGt3MimSGI,19227
-brainstate/transform/__init__.py,sha256=CfzUYGQFt9hAp57ZvffUBBhLjudgzJ9n_aTaor57iOk,2241
+brainstate/random/__init__.py,sha256=yeWQ3RUcFXtXcDkFhaOYa_nwQ_M7hlH_W69YEcGF3Oo,8351
+brainstate/random/_fun.py,sha256=fW2bc2i15sQNodmKT66ZCfXJSIZ8Ygxmn_Xf80Zjygw,135334
+brainstate/random/_fun_test.py,sha256=EcHOcCOsmZPbm1n5TUyiTDxS13Pn3MJbq5IG-twbnPI,22866
+brainstate/random/_impl.py,sha256=A79IK8YNZuN8RGLftASKHcZI0Cdsep7mGYNtCfKl2fQ,21830
+brainstate/random/_seed.py,sha256=mLHqOu-lJQjsXo4nODACY78SSKgYayHx3n5IBZ9L6J0,24922
+brainstate/random/_seed_test.py,sha256=Y2VCAkUzciDaCfYZWPe_Ewmi3MylK-WzfPA7TzorV8Q,1491
+brainstate/random/_state.py,sha256=JPK6-jqFwrzsrb3lrxZ4GahP0TtXycUfWrJSdjHIEg8,42946
+brainstate/random/_state_test.py,sha256=OfW0WxTpJZm_kT_7bZjJJ8ZtLM_X8NBF90h1sDSZxmw,19221
+brainstate/transform/__init__.py,sha256=P7MAmt4pJYcpxLO30gFoT4BO_AHHM3VNOoG2j3OkNRk,2126
 brainstate/transform/_ad_checkpoint.py,sha256=4dcNCEQVV_CPMSkE32URERDMpQHbyfdGeLT_Nvhyd4o,6912
 brainstate/transform/_ad_checkpoint_test.py,sha256=fPXBjDxsLHbL2mhIU3x_F5BpitkXLpgIsRCRgm2Us6w,1697
 brainstate/transform/_autograd.py,sha256=4zGSYa9TMn6bqzPJNLfU9UZGZRyYxmwMXTKZWO4w3QQ,39991
 brainstate/transform/_autograd_test.py,sha256=saWG1_k3cRXpsyQDzQkOLGvsF7IIxG9aGnjrf5B3HNk,44112
-brainstate/transform/_conditions.py,sha256=IIu_V2f7R74saOoURGUxwzI1RpsnxkzpfB2fes9plOs,11399
+brainstate/transform/_conditions.py,sha256=nLc_m0bLlybEPTi42feSlz2zTDIHB-BknSqPTIL_I0w,11376
 brainstate/transform/_conditions_test.py,sha256=MEuqRq6IFmyORRDi0qWvNo4pWKFyc8aNrW1v9Saqxj0,8493
 brainstate/transform/_error_if.py,sha256=e9tp3wT5p4bEyjn_Za_SrPNOG3OIoPBMIrvG2CsZzvw,2680
 brainstate/transform/_error_if_test.py,sha256=yn-qcZ6lZUWciIif4fJOpdpKzJFAAdfgzZm6FfPeq7U,1848
-brainstate/transform/_eval_shape.py,sha256=BNbjiFHUsk-qfENiZf1K8yX-x7eIIuAyWR0CjBIBr5w,5355
-brainstate/transform/_eval_shape_test.py,sha256=4A2NdHcpksiGPf_UmPlMPgHJnes9ciiQO20ZJHVzA9g,1355
+brainstate/transform/_find_state.py,sha256=nUrVp_DUP78E2H8UHyhH8RL03kwcKjepnkDTqbISFmI,7405
+brainstate/transform/_find_state_test.py,sha256=KahI4HSj6MTlC4ccCPyLiyERy3MryZd_iczGwFXb2bM,2913
 brainstate/transform/_jit.py,sha256=qYsL3Z9nAAW0UyQe_AyvBEuJvqAan6iw3lN49o1oC0A,15421
 brainstate/transform/_jit_test.py,sha256=ecw54dGQYdJq2J94itPrXBQrSDNvCY_htUD7z7y4HUM,4013
 brainstate/transform/_loop_collect_return.py,sha256=HhjC2gq6qzliw4ofP16VxdtR5hW-NmDZdeHxuiLdYGk,25899
 brainstate/transform/_loop_collect_return_test.py,sha256=BVK-b3CuDtTXciRaA_8t4751N4taQOnIPNzAelSts-k,1753
 brainstate/transform/_loop_no_collection.py,sha256=ArPpNemMh4jJsq_vUWPxuagCnxTlONN--3P_-44qYq8,10156
 brainstate/transform/_loop_no_collection_test.py,sha256=3bRo9_Oaypbw3asEevrgTK0WksxDAIZKJgeaWpt7nl8,1371
-brainstate/transform/_make_jaxpr.py,sha256=A9ivekPuzLUCoUcgSLm5hUuWPj5_1_ksqn2KvztcWhY,73327
-brainstate/transform/_make_jaxpr_test.py,sha256=dFAnEQavYlEddex6h42ulRkXbV5TSGCDHFXWAvPsy7w,49128
-brainstate/transform/_mapping.py,sha256=L9Q_1M_GXJA0ZCJtty9x01wViyefgeHyOLga-Jo2AE8,21570
-brainstate/transform/_mapping_test.py,sha256=yGG-QSo3epqtBvGX8VxrNDKiC-d6TMN_jv_uKPad6eM,7647
+brainstate/transform/_make_jaxpr.py,sha256=_pSw1oumQ504Emvcgy7eXBTxOw1c1pbnvbcdveKMG8s,80444
+brainstate/transform/_make_jaxpr_test.py,sha256=EfY-ZKR-u-aLfjOTPY6QiwEzZU-QpIUsB5oCCZi5Z-I,53014
+brainstate/transform/_mapping.py,sha256=wBTvmXSg8TW-ZjIhKU5nQLLb6BIkLTCaJWwESZ72Hks,22021
+brainstate/transform/_mapping_test.py,sha256=BkoL9peJGiJvkp7aLErrB6GnDasSviMbBkBqh1yD8LM,3255
 brainstate/transform/_progress_bar.py,sha256=kZ-mI5hbUQXhqKFVyo0qeKG_LvrR9ZIar7WkXyOeET0,8961
-brainstate/transform/_random.py,sha256=ZTH5Smx5SvFy6El7qk-ihoYE9WII6TJXcsC9Vm7VgjA,5259
 brainstate/transform/_unvmap.py,sha256=cW6fjs5Iy1YBB6Nx2mxlM4IzV8U99bEX5QjT8rBRDho,6319
 brainstate/transform/_util.py,sha256=IYqJj7oyAYzm_m3d9WEsUQRKDdVLQaAWpwM5O8PD4YQ,11304
 brainstate/util/__init__.py,sha256=anHdG5BIsMqcBQy7gt8lErKInZv1wf2NOLJGUqltyAQ,1154
@@ -104,8 +104,8 @@ brainstate/util/filter.py,sha256=wY_XUF3OhrXSV1bZkTcVhlEPba4HP1l9N5aRW2zgxqQ,274
 brainstate/util/filter_test.py,sha256=ZfrEeOc1yMHYzrcSR3p4jbZGj7c_tXC8VcPq7H13q8E,31653
 brainstate/util/struct.py,sha256=LYPLGDGfPuw14hhx5k4rb8msSH3yZPdAu_0CvjxPWwE,24505
 brainstate/util/struct_test.py,sha256=q_fWsUH1ON35DKjUUAMq6VtYglqTDyBvd6WMVGD89EI,16526
-brainstate-0.2.0.dist-info/licenses/LICENSE,sha256=RJ40fox7u2in2H8wvIS5DsPGlNHaA7JI024thFUlaZE,11348
-brainstate-0.2.0.dist-info/METADATA,sha256=GrI-RT31978SYlGNnRP9xodMOYl1g6mPpIiP1XdUaOY,4421
-brainstate-0.2.0.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-brainstate-0.2.0.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.2.0.dist-info/RECORD,,
+brainstate-0.2.2.dist-info/licenses/LICENSE,sha256=RJ40fox7u2in2H8wvIS5DsPGlNHaA7JI024thFUlaZE,11348
+brainstate-0.2.2.dist-info/METADATA,sha256=tbhbVxomU2orW-T3WaEjVtYAyZSn4T4uJxAXzPmU9JY,4421
+brainstate-0.2.2.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+brainstate-0.2.2.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.2.2.dist-info/RECORD,,

brainstate/transform/_eval_shape.py

@@ -1,145 +0,0 @@
-# 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.
-# 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.
-# ==============================================================================
-
-import functools
-from typing import Any, TypeVar, Callable, Sequence, Union
-
-import jax
-
-from brainstate import random
-from brainstate._utils import set_module_as
-from brainstate.graph import Node, flatten, unflatten
-from ._random import restore_rngs
-
-__all__ = [
-    'abstract_init',
-]
-
-A = TypeVar('A')
-
-
-@set_module_as('brainstate.transform')
-def abstract_init(
-    fn: Callable[..., A],
-    *args: Any,
-    rngs: Union[random.RandomState, Sequence[random.RandomState]] = random.DEFAULT,
-    **kwargs: Any,
-) -> A:
-    """
-    Compute the shape/dtype of ``fn`` without any FLOPs.
-
-    This function evaluates the shape and dtype of the output of a function without
-    actually executing the computational operations. It's particularly useful for
-    initializing neural network models to understand their structure and parameter
-    shapes without performing expensive computations.
-
-    Parameters
-    ----------
-    fn : callable
-        The function whose output shape should be evaluated.
-    *args
-        Positional argument tuple of arrays, scalars, or (nested) standard
-        Python containers (tuples, lists, dicts, namedtuples, i.e. pytrees) of
-        those types. Since only the ``shape`` and ``dtype`` attributes are
-        accessed, one can use :class:`jax.ShapeDtypeStruct` or another container
-        that duck-types as ndarrays (note however that duck-typed objects cannot
-        be namedtuples because those are treated as standard Python containers).
-    rngs : RandomState or sequence of RandomState, default random.DEFAULT
-        A :class:`RandomState` or a sequence of :class:`RandomState` objects
-        representing the random number generators to use. If not provided, the
-        default random number generator will be used.
-    **kwargs
-        Keyword argument dict of arrays, scalars, or (nested) standard
-        Python containers (pytrees) of those types. As in ``args``, array values
-        need only be duck-typed to have ``shape`` and ``dtype`` attributes.
-
-    Returns
-    -------
-    A
-        A nested PyTree containing :class:`jax.ShapeDtypeStruct` objects as leaves,
-        representing the structure and shape/dtype information of the function output.
-
-    Examples
-    --------
-    Basic usage with neural network initialization:
-
-    .. code-block:: python
-
-        >>> import brainstate
-        >>> import jax.numpy as jnp
-        >>>
-        >>> class MLP:
-        ...     def __init__(self, n_in, n_mid, n_out):
-        ...         self.dense1 = brainstate.nn.Linear(n_in, n_mid)
-        ...         self.dense2 = brainstate.nn.Linear(n_mid, n_out)
-        >>>
-        >>> # Get shape information without actual computation
-        >>> model_shape = brainstate.transform.abstract_init(lambda: MLP(1, 2, 3))
-
-    With function arguments:
-
-    .. code-block:: python
-
-        >>> def create_model(input_size, hidden_size, output_size):
-        ...     return brainstate.nn.Sequential([
-        ...         brainstate.nn.Linear(input_size, hidden_size),
-        ...         brainstate.nn.ReLU(),
-        ...         brainstate.nn.Linear(hidden_size, output_size)
-        ...     ])
-        >>>
-        >>> # Abstract initialization with arguments
-        >>> model_shape = brainstate.transform.abstract_init(
-        ...     create_model, 784, 256, 10
-        ... )
-
-    Using custom random number generators:
-
-    .. code-block:: python
-
-        >>> import brainstate.random as random
-        >>>
-        >>> # Create custom RNG
-        >>> rng = random.RandomState(42)
-        >>>
-        >>> def init_with_custom_weights():
-        ...     return brainstate.nn.Linear(10, 5)
-        >>>
-        >>> model_shape = brainstate.transform.abstract_init(
-        ...     init_with_custom_weights, rngs=rng
-        ... )
-
-    Evaluating function with array inputs:
-
-    .. code-block:: python
-
-        >>> def model_forward(x):
-        ...     layer = brainstate.nn.Linear(x.shape[-1], 128)
-        ...     return layer(x)
-        >>>
-        >>> # Use ShapeDtypeStruct to represent input without actual data
-        >>> input_shape = jax.ShapeDtypeStruct((32, 784), jnp.float32)
-        >>> output_shape = brainstate.transform.abstract_init(model_forward, input_shape)
-    """
-
-    @functools.wraps(fn)
-    @restore_rngs(rngs=rngs)
-    def _eval_shape_fn(*args_, **kwargs_):
-        out = fn(*args_, **kwargs_)
-        assert isinstance(out, Node), 'The output of the function must be Node'
-        graph_def, treefy_states = flatten(out)
-        return graph_def, treefy_states
-
-    graph_def_, treefy_states_ = jax.eval_shape(_eval_shape_fn, *args, **kwargs)
-    return unflatten(graph_def_, treefy_states_)

brainstate/transform/_eval_shape_test.py

@@ -1,38 +0,0 @@
-# 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.
-# 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.
-# ==============================================================================
-
-
-import unittest
-
-import brainstate
-
-
-class TestEvalShape(unittest.TestCase):
-    def test1(self):
-        class MLP(brainstate.nn.Module):
-            def __init__(self, n_in, n_mid, n_out):
-                super().__init__()
-                self.dense1 = brainstate.nn.Linear(n_in, n_mid)
-                self.dense2 = brainstate.nn.Linear(n_mid, n_out)
-
-            def __call__(self, x):
-                x = self.dense1(x)
-                x = brainstate.functional.relu(x)
-                x = self.dense2(x)
-                return x
-
-        r = brainstate.augment.abstract_init(lambda: MLP(1, 2, 3))
-        print(r)
-        print(brainstate.random.DEFAULT)