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

brainstate-0.2.0.dist-info/RECORD

@@ -0,0 +1,111 @@
+brainstate/__init__.py,sha256=bmCIZG6xMiVlZ5qWk-d5qpkqdnrWxFhLePtpu3ff5WI,6024
+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/_error.py,sha256=6A5ILy17ZMMZIjS8LkajTZBDRnwv_Qait5x__h2Levo,1522
+brainstate/_state.py,sha256=YCYKrX2xCTojIY15vtQl8WWjtJ0HNMg8y5LCUk1d_3Q,58096
+brainstate/_state_test.py,sha256=wCoWTbvARVkhNokMLLBV6sEoHSE9OJRyuUVrNDcRiG8,1589
+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_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_test.py,sha256=IVyrJh4io3sDgtrTEIAItGzNs2XEf7rO1rvI1r_KiII,39119
+brainstate/nn/__init__.py,sha256=oEQ81xpWppPMUMCRalMhAhUU5faYs0aKcyydwjoxGVo,4759
+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_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/_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_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/_elementwise.py,sha256=4kKzrbKn5luwnpY8n7IeaMOtBVVie8oPHEc64hSn8-w,34858
+brainstate/nn/_elementwise_test.py,sha256=sbWlUyTB8oiu3PRHObTvmUaob99EjIGYE6k6bEpg6K8,27296
+brainstate/nn/_embedding.py,sha256=eMwbz9udm6WGVEGeHwhvCr7sJA6kMFyaZsli7lL30eg,14942
+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_test.py,sha256=21qomGOo96YLmlEQok2hCByAcRpmqREhSr3kxmhKOm8,13014
+brainstate/nn/_linear.py,sha256=olbo9AmC35UQBxECyAQW06XeucsJfBAAMoSfwOW6c7s,24047
+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_test.py,sha256=znjB7FU5evJENQ1Pqw7ZlOGC5faQe4-4VpjW60H8UWI,1414
+brainstate/nn/_normalizations.py,sha256=exdIn627ph5pVHdYx_4NIaK_f1xcLaBj-YQ0Ui22CsA,50185
+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
+brainstate/nn/_poolings.py,sha256=aiLDTgtbDEse2OZlJONdU1CCsBjZtsWqsnL8ffTSGpk,86045
+brainstate/nn/_poolings_test.py,sha256=cZ4lyutUY-Iti6rkGWjqWZ72fXF5QvwX9mki4CQ1xDs,34164
+brainstate/nn/_rnns.py,sha256=FmuVsUwQCl7QbePmYokd-ZxKwVgDqa8FdCto_dNFDjk,32947
+brainstate/nn/_rnns_test.py,sha256=gqPGz64i_QUEXi2fUBrIaBEPZjTs8hm8S6HyPbs3I6M,22112
+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/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_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/_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/_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
+brainstate/util/_others.py,sha256=HSP3ynNv2ocxPz5omQl2rxMTOjGkQNNprsTOCEANCjA,30837
+brainstate/util/_others_test.py,sha256=gEiUybMxtv11tsD0jge9Y0K8idkWkHVBlTZz1gKXNAU,30535
+brainstate/util/_pretty_pytree.py,sha256=0fVwW8qHtKrpJU4q2tbn8usSKOlG8lBALn6E6crpf64,46815
+brainstate/util/_pretty_pytree_test.py,sha256=6DiYX_Rwyvv3XcuQHsU-PuWXmvOqSmiVVN4diVvp76o,22362
+brainstate/util/_pretty_repr.py,sha256=fafG6SIFoPjaWmQcTHwmnEqbVvcTCoN_lKiBKBk1QTQ,13958
+brainstate/util/_pretty_repr_test.py,sha256=yfaANMfozlM5E3IZv-LBaNpyvSVB8ftRp0r7_oAr5yA,23122
+brainstate/util/filter.py,sha256=wY_XUF3OhrXSV1bZkTcVhlEPba4HP1l9N5aRW2zgxqQ,27410
+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/augment/__init__.py

@@ -1,30 +0,0 @@
-# Copyright 2024 BDP 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.
-# ==============================================================================
-
-"""
-This module includes transformations for augmenting the functionalities of JAX code.
-"""
-
-from ._autograd import GradientTransform, grad, vector_grad, hessian, jacobian, jacrev, jacfwd
-from ._eval_shape import abstract_init
-from ._mapping import vmap, pmap, map, vmap_new_states
-from ._random import restore_rngs
-
-__all__ = [
-    'GradientTransform', 'grad', 'vector_grad', 'hessian', 'jacobian', 'jacrev', 'jacfwd',
-    'abstract_init',
-    'vmap', 'pmap', 'map', 'vmap_new_states',
-    'restore_rngs',
-]

brainstate/augment/_eval_shape.py

@@ -1,99 +0,0 @@
-# Copyright 2024 BDP 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.graph import Node, flatten, unflatten
-from ._random import restore_rngs
-
-__all__ = [
-    'abstract_init',
-]
-
-A = TypeVar('A')
-
-
-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.
-
-    Here's an example::
-
-        >>> import brainstate
-        >>> 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)
-
-        >>> r = brainstate.augment.abstract_init(lambda: MLP(1, 2, 3))
-        >>> r
-        MLP(
-          dense1=Linear(
-            in_size=(1,),
-            out_size=(2,),
-            w_mask=None,
-            weight=ParamState(
-              value={'bias': ShapeDtypeStruct(shape=(2,), dtype=float32), 'weight': ShapeDtypeStruct(shape=(1, 2), dtype=float32)}
-            )
-          ),
-          dense2=Linear(
-            in_size=(2,),
-            out_size=(3,),
-            w_mask=None,
-            weight=ParamState(
-              value={'bias': ShapeDtypeStruct(shape=(3,), dtype=float32), 'weight': ShapeDtypeStruct(shape=(2, 3), dtype=float32)}
-            )
-          )
-        )
-
-    Args:
-        fn: The function whose output shape should be evaluated.
-        *args: a 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).
-        **kwargs: a 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.
-        rngs: 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.
-
-    Returns:
-        out: a nested PyTree containing :class:`jax.ShapeDtypeStruct` objects as leaves.
-
-    """
-
-    @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_)