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

brainstate/transform/_random.py

@@ -1,171 +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 Callable, Sequence, Union
-
-from brainstate._utils import set_module_as
-from brainstate.random import DEFAULT, RandomState
-from brainstate.typing import Missing
-from brainstate.util import PrettyObject
-
-__all__ = [
-    'restore_rngs'
-]
-
-
-class RngRestore(PrettyObject):
-    """
-    Manage backing up and restoring multiple random states.
-
-    Parameters
-    ----------
-    rngs : Sequence[RandomState]
-        Sequence of :class:`~brainstate.random.RandomState` instances whose
-        states should be captured and restored.
-
-    Attributes
-    ----------
-    rngs : Sequence[RandomState]
-        Managed random-state instances.
-    rng_keys : list
-        Cached keys captured by :meth:`backup` until :meth:`restore` runs.
-
-    Examples
-    --------
-    .. code-block:: python
-
-        >>> import brainstate
-        >>>
-        >>> rng = brainstate.random.RandomState(0)
-        >>> restorer = brainstate.transform.RngRestore([rng])
-        >>> restorer.backup()
-        >>> _ = rng.random()
-        >>> restorer.restore()
-    """
-    __module__ = 'brainstate.transform'
-
-    def __init__(self, rngs: Sequence[RandomState]):
-        """
-        Initialize a restorer for the provided random states.
-
-        Parameters
-        ----------
-        rngs : Sequence[RandomState]
-            Random states that will be backed up and restored.
-        """
-        self.rngs: Sequence[RandomState] = rngs
-        self.rng_keys = []
-
-    def backup(self):
-        """
-        Cache the current key for each managed random state.
-
-        Notes
-        -----
-        The cached keys persist until :meth:`restore` is called, after which the
-        internal cache is cleared.
-        """
-        self.rng_keys = [rng.value for rng in self.rngs]
-
-    def restore(self):
-        """
-        Restore each random state to the cached key.
-
-        Raises
-        ------
-        ValueError
-            Raised when the number of stored keys does not match ``rngs``.
-        """
-        if len(self.rng_keys) != len(self.rngs):
-            raise ValueError('The number of random keys does not match the number of random states.')
-        for rng, key in zip(self.rngs, self.rng_keys):
-            rng.restore_value(key)
-        self.rng_keys.clear()
-
-
-def _rng_backup(
-    fn: Callable,
-    rngs: Union[RandomState, Sequence[RandomState]]
-) -> Callable:
-    rng_restorer = RngRestore(rngs)
-
-    @functools.wraps(fn)
-    def wrapper(*args, **kwargs):
-        # backup the random state
-        rng_restorer.backup()
-        # call the function
-        out = fn(*args, **kwargs)
-        # restore the random state
-        rng_restorer.restore()
-        return out
-
-    return wrapper
-
-
-@set_module_as('brainstate.transform')
-def restore_rngs(
-    fn: Callable = Missing(),
-    rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
-) -> Callable:
-    """
-    Decorate a function so specified random states are restored after execution.
-
-    Parameters
-    ----------
-    fn : Callable, optional
-        Function to wrap. When omitted, :func:`restore_rngs` returns a decorator
-        preconfigured with ``rngs``.
-    rngs : Union[RandomState, Sequence[RandomState]], optional
-        Random states whose keys should be backed up before running ``fn`` and
-        restored afterwards. Defaults to :data:`brainstate.random.DEFAULT`.
-
-    Returns
-    -------
-    Callable
-        Wrapped callable that restores the random state or a partially applied
-        decorator depending on how :func:`restore_rngs` is used.
-
-    Raises
-    ------
-    AssertionError
-        If ``rngs`` is neither a :class:`~brainstate.random.RandomState` instance nor
-        a sequence of such instances.
-
-    Examples
-    --------
-    .. code-block:: python
-
-        >>> import brainstate
-        >>>
-        >>> rng = brainstate.random.RandomState(0)
-        >>>
-        >>> @brainstate.transform.restore_rngs(rngs=rng)
-        ... def sample_pair():
-        ...     first = rng.random()
-        ...     second = rng.random()
-        ...     return first, second
-        >>>
-        >>> assert sample_pair()[0] == sample_pair()[0]
-    """
-    if isinstance(fn, Missing):
-        return functools.partial(restore_rngs, rngs=rngs)
-
-    if isinstance(rngs, RandomState):
-        rngs = [rngs]
-    assert isinstance(rngs, Sequence), 'rngs must be a RandomState or a sequence of RandomState instances.'
-    for rng in rngs:
-        assert isinstance(rng, RandomState), 'rngs must be a RandomState or a sequence of RandomState instances.'
-    return _rng_backup(fn, rngs=rngs)

brainstate-0.2.1.dist-info/RECORD

@@ -1,111 +0,0 @@
-brainstate/__init__.py,sha256=QXLW5UOeT0Q25j6AyLkejFrGkBAQWA2I9O7TVycjV_w,6193
-brainstate/_compatible_import.py,sha256=JUAHEeZBFTN6KW9ZZoHW5M1KtxQA-XyJJB6v_f83ndM,11377
-brainstate/_compatible_import_test.py,sha256=ynjx_ErAzE7jZJkJe-mnBjVVuBSURiCV3HzZ-RYAT1E,23443
-brainstate/_deprecation.py,sha256=sg5xBy15QQd7wo9Ew7QH89vH6ZSs96ayCi5cEdyrpqE,8591
-brainstate/_deprecation_test.py,sha256=5KPrCC8o5xqLg9BhmvcLn-UT97YPLSsSaXhla2jVQQA,90644
-brainstate/_error.py,sha256=d7aaHVHLfh8aynhduvBmZSB8aK69aGHGm17D55xhTEY,1567
-brainstate/_state.py,sha256=ngz227ueaFV1E48LvlHrzjz7cRPKXntPyjYCSlQN2T8,59748
-brainstate/_state_test.py,sha256=LzfGfM2efKMjlX5EpIwFsi_tbBWM9HQkCPzeiSZ7PMg,1641
-brainstate/_utils.py,sha256=jgxW7Ouj_r6nKcG7UHjX4UYtyOX4QOMD5Q4CYrnswtU,1660
-brainstate/environ.py,sha256=WxCkG_9-RORunpVHU2zJebyxShf_XotD2HHzatFeUos,43664
-brainstate/environ_test.py,sha256=vgEcFtZXLYBnVwOPGLQJehz1qODcJc_C2esLL9ORD2M,43641
-brainstate/mixin.py,sha256=Iv149vfepDmYYFp8_wBVXc41JwP6WvPpTIG1p70N7Uk,46068
-brainstate/mixin_test.py,sha256=Vi9p3BVU6wrNVsJvQfooS3riJTml8tYNq1UM-wO4N8Q,35735
-brainstate/typing.py,sha256=cKFeLRa8Aowjq_-eWXkt1xO9cWX6zvyVn_IWDcm9fTA,27138
-brainstate/typing_test.py,sha256=nJc325IwKrB1DaxJX1fq4VIdLpmvT7x1vYtNcGWnZmA,27020
-brainstate/graph/__init__.py,sha256=KZOn3B0vuzYl55XQm6N5pCNnDTyHdrEggfp_h_-Ml7Q,868
-brainstate/graph/_node.py,sha256=ybMR1I-5ztLeNJ9GgxuiB-kgSSisYTp4xJ6XUt-qv40,6647
-brainstate/graph/_node_test.py,sha256=BIUOuXRYLj0KVPC5Slb98wRcO158T63T53zHnPnhpZk,19276
-brainstate/graph/_operation.py,sha256=UAQe2z73rOHX45PNFAC2X0ltkh9NVQ3daIk395Jxn70,55720
-brainstate/graph/_operation_test.py,sha256=_Z7ZJLdOsZAjorvc4PCu38hDQ7-2DQnJsgmwgHdePrg,40266
-brainstate/nn/__init__.py,sha256=3l3uT6UBQCta6MUMgIAyZIOtjhfi6KI6zT7gy2GiCX0,4896
-brainstate/nn/_activations.py,sha256=6GriLZzgEFcv_13f0EiaG7r3xXWun7X749vriPYHdLM,28017
-brainstate/nn/_activations_test.py,sha256=rjDx7Pn2gg5vk1RyMNfnK47kB2YDilY4J7ygmOEgutU,13792
-brainstate/nn/_collective_ops.py,sha256=db988H8e5B8PR8_BQJvspeOoKl-pSdb-hP2ZBRTbbR8,21913
-brainstate/nn/_collective_ops_test.py,sha256=82V4egN9GdQeph9T3BG4YH4I43qPmB98hgkBU--UAlk,26374
-brainstate/nn/_common.py,sha256=-uBCsn0B5Z1LMwJI4K9FuK_8l7UiU88H2S84Nv6jveA,7410
-brainstate/nn/_common_test.py,sha256=bqngm_1syVwozgDLhzkdmEaKtEExKghCo2Nqdq6R5MM,6120
-brainstate/nn/_conv.py,sha256=lpX75tqVDnA_AJRV7Mp0OpX662wTUzdiLhTjfTKSONw,84116
-brainstate/nn/_conv_test.py,sha256=srZDdlG4TiYQfhDwdQNOEyb_fH2FT79Loe8voZ48ods,31022
-brainstate/nn/_delay.py,sha256=SLeH0Qv953tVJwZv5GhRIwIVX1WkasaN4dDTTNFqEkE,22906
-brainstate/nn/_delay_test.py,sha256=RX1Ih0p2PzOOqEbVTePLITobKBN8AveUJ5shXGCfdgk,10619
-brainstate/nn/_dropout.py,sha256=kWOMzv6peNYc4K5sczennI10-dEPCMXiDicAWEBcUdA,23173
-brainstate/nn/_dropout_test.py,sha256=_Ym1iYz5XSjMfGk_YRIK7qZ7s2dFwJkP-8JIWizF1Jo,21938
-brainstate/nn/_dynamics.py,sha256=1UO9k077gePxnxMyZMMdB-2N0LalvxaPIaqII3H5E1w,45377
-brainstate/nn/_dynamics_test.py,sha256=Qz4sl9TpZygGzYxGnIbtYXNO8aJf-PHvWnZ_dByBYp4,2423
-brainstate/nn/_elementwise.py,sha256=I8FW5oWagFZzI-LhRMxxUf334Ai5t8elAsaeZKx1Lcc,36156
-brainstate/nn/_elementwise_test.py,sha256=2_lpkkJjAo-GxKNkxPo4sbSvsDbs0qbRMlHSXlFVVpk,28125
-brainstate/nn/_embedding.py,sha256=rEahcqYYq1oWpKL-chxzQMIectQdbFu_weVEULynVMg,15350
-brainstate/nn/_embedding_test.py,sha256=FnXhmYWtbzZvoyIFoJkOTsvP3hTOtfxB6LCHyaSbqQI,6398
-brainstate/nn/_event_fixedprob.py,sha256=ngR9LKqxf-YDMtHAm14eHrORRFagc4ez4qcpQJFgpBE,9628
-brainstate/nn/_event_fixedprob_test.py,sha256=XM-4XfqSkFQoSTChm6EIRThOHBPVRtdb-vJfoNp72PI,3951
-brainstate/nn/_event_linear.py,sha256=LfccbsMwaYMDqcKQVREvTTmKtuVEqGBllMjeTDd-lNs,2637
-brainstate/nn/_event_linear_test.py,sha256=g7-FqtJL32Me_WVA02XyoFZoJWsc_W5OX1sMqFt7zXM,3886
-brainstate/nn/_exp_euler.py,sha256=-o1gBs2bPOraWu83v-cSBG86PFeqKu5EC3pZVAC0GMM,8889
-brainstate/nn/_exp_euler_test.py,sha256=iLFgSDSF8sfke4GOmxZpzhH-OhwbiOiqbDa75yWAKH8,13391
-brainstate/nn/_linear.py,sha256=m5hggIYbgLYPbbfpQYgBrb0g8scu5ue7wnu3VW1djGc,24791
-brainstate/nn/_linear_test.py,sha256=YTVpbVP6ZzuwcheUuq9kHLK0JRQVIL9mvUg2mVkk12k,18351
-brainstate/nn/_metrics.py,sha256=-jl7fVYGmhGS4VMQUuf3iLzMboRRj9AfprZ1EfBXAxQ,37551
-brainstate/nn/_metrics_test.py,sha256=z5Z8BpHh6qTsS2qZkUupcZt5sRQ9QR1C-8E24SfBkr0,24715
-brainstate/nn/_module.py,sha256=1e8ZMdq4Nl2im_ZpDDr2TZu9Xi-4sO4BJedBtXg9Aj8,13159
-brainstate/nn/_module_test.py,sha256=XH1i8R15_uKkKkotfOAZK_2WNcXj-gLL6gmSgnFfHks,1454
-brainstate/nn/_normalizations.py,sha256=udyWdudXg5ZlOxwzsILocRDoHq2tv_SwK0heTvqEAoM,51519
-brainstate/nn/_normalizations_test.py,sha256=XPW5icjqiMyfC398-0TRelg9ZyFu1SD-wJJLVjSXPUU,24049
-brainstate/nn/_paddings.py,sha256=i7cLWHwwx9y5EOkxeDilbxr84xW4U8ZwBeECCGcS14E,33435
-brainstate/nn/_paddings_test.py,sha256=z0Tyz0rE4I1Y0C4wbZAqb3v8YUQl5Y8vBalzTm-YSlU,28131
-brainstate/nn/_poolings.py,sha256=WbA8helFy-PbBYfUphVXDrXwcYymhxp9vbosU7PZpjU,88284
-brainstate/nn/_poolings_test.py,sha256=bouz1sMq0xz5EmwdLuFFX2c26apSW1NW8dDxiWruFpc,35116
-brainstate/nn/_rnns.py,sha256=iKsJT7KmNYsVQGN23JvnrwvhGBAzSpg1yfJ8P18nwBY,33893
-brainstate/nn/_rnns_test.py,sha256=VGOtLh-C5iMUDQUppYWPDXyFCGy74sdberdeMIjUOP4,22704
-brainstate/nn/_utils.py,sha256=1iRQEBCXwel0E72c7XPw_SQolSRk7BZsEm6_A2b4ZkQ,7619
-brainstate/nn/_utils_test.py,sha256=fqxfbJM03N8K4qmjtoHiG6h8UlzwdDrd7YiEK13HWSA,14667
-brainstate/nn/init.py,sha256=UenFiGcB7KDa0f67eabsHZHl38-AYM0eoblrXAxiB4I,26465
-brainstate/nn/init_test.py,sha256=HkjWcNkLP9qVy4qLTazEaKgULsmBo-wyaBvHj77vbuI,6380
-brainstate/random/__init__.py,sha256=K_JLJNW0H3RZKW1G3sGnlZ9pK9l6ANx5DOxw5P0RlyY,8653
-brainstate/random/_rand_funs.py,sha256=ac2rA2hWNfZX6QznG8endpRcJv7umiq61b36kbQx0Xg,139277
-brainstate/random/_rand_funs_test.py,sha256=QETjwkq_PlVE2mOptQeNCNbfZ4pQ_8MsDVJLFW9N5GA,23680
-brainstate/random/_rand_seed.py,sha256=LBi4rCtR5iQ-GilU-3E3eHPc8d1rlHA8147Fa3I6R60,25602
-brainstate/random/_rand_seed_test.py,sha256=BIQdvb9TsfIHs4UF4N9qvU3wJKdnkqnDK5CkWYCo3QA,1539
-brainstate/random/_rand_state.py,sha256=NbnihT3BupYQOmiU9njHXN2HSKjugAa2xnxDI_0rj7k,55208
-brainstate/random/_rand_state_test.py,sha256=XB1xq3WXAFLSGmxo-vBJulw9WoeBnJq5AFuL5Xty5SQ,19778
-brainstate/transform/__init__.py,sha256=CjHzuRQHh0EDYIxugZ0T819X1P2IXIC-89Q_gQ9izCw,2300
-brainstate/transform/_ad_checkpoint.py,sha256=248gdnHJHoa0HMIyqC2MiS5dNVC3xVI5ViCxskneTtc,7088
-brainstate/transform/_ad_checkpoint_test.py,sha256=KLfpBozd6Doc3ck6an3VGunViKpzHEc2f-3gcT189C4,1746
-brainstate/transform/_autograd.py,sha256=rC_1zUVLgZMDTA3EAYpZYsifcMl3P2tgWf-2VDsFuSk,41016
-brainstate/transform/_autograd_test.py,sha256=jXX_w-iXjOrETKvSvRau6aTP98SOGEaOpqMcqPAEVGU,45401
-brainstate/transform/_conditions.py,sha256=eMIPgg5N_Olq-iaek4ELkFGG0zBkThNJC9aZLC5sdIk,11715
-brainstate/transform/_conditions_test.py,sha256=IpAzs8W5zmXcC_wYFwI4WFoNRbR9VmGdHUg8zw-S8BA,8713
-brainstate/transform/_error_if.py,sha256=Pe5h2sMeDmZqfuhBsmiaNQrFgbj0m6LLP2y8VYso5PY,2774
-brainstate/transform/_error_if_test.py,sha256=H4SUYIPwAVA9xMZsBYr_Kj79sndQzT26ypRck4iG7qk,1900
-brainstate/transform/_eval_shape.py,sha256=g00ue2Ijo3-jTvMjHzhJql0DR_u337nLKsDaLAov8R0,5500
-brainstate/transform/_eval_shape_test.py,sha256=BEmtqyudsKEz0Mg7SdXSm55dalL5W6ZXwhvs4oUJZEY,1393
-brainstate/transform/_jit.py,sha256=msAGXVHIBx5h0jdnfYXhDDYIqCDJBQuKe-iR56ELCd4,15820
-brainstate/transform/_jit_test.py,sha256=g97Kc5eYfe05opPTlr0N1JXetM6WilbwMfbH1t4xBjo,4156
-brainstate/transform/_loop_collect_return.py,sha256=UtAlwjD2NAD7Wd71SNqKKRf6tNEmsywXy64nixcgzGU,26574
-brainstate/transform/_loop_collect_return_test.py,sha256=vMptPxrsxWeVSBD8tcOhsO6B9aCjYMcz1QBEgccHNqQ,1811
-brainstate/transform/_loop_no_collection.py,sha256=bmxd2YWc63lVZhDq12mZPDgGghYqB6ObN63m0IzjqL0,10439
-brainstate/transform/_loop_no_collection_test.py,sha256=OlYGQEIXrHTr0ZylLBBx02f4FGZCjXYRixkvGCwH-oo,1421
-brainstate/transform/_make_jaxpr.py,sha256=dHE4TRBALtzVsCg8DwTfTKWJTpim5OMKVp8SdTGPxZw,75343
-brainstate/transform/_make_jaxpr_test.py,sha256=neK2iu1HK0dAk_XPERE6DUquuEkXBFnbsKiltXJBThE,50638
-brainstate/transform/_mapping.py,sha256=EtpGI_aABQNk5i9cKuEDr7pCG-YN97cBL_XkvsPH9zs,22099
-brainstate/transform/_mapping_test.py,sha256=RALLjERRZarsRt-XgCD8_vEvnbSbNnz25H-u9fuvumo,7841
-brainstate/transform/_progress_bar.py,sha256=z5SIRj3R8IUVMHPnOsXHQDFsEe3wYXIBM_b9-Zk9Weg,9216
-brainstate/transform/_random.py,sha256=kYa1Xiu_LGHofoPcV9JVBcobTNoOeC9zygJZ6UMeLcI,5430
-brainstate/transform/_unvmap.py,sha256=DLalab8FY6tzFiha4M6bZ4i4DRDJsaQpzmOxj0SkvhI,6575
-brainstate/transform/_util.py,sha256=6Dm4QE1zPS9qwxVop8gzzKKOWKXTt_FqjQY1rl6GtKo,11590
-brainstate/util/__init__.py,sha256=2lKyFfMScZjRh2Zl6PheaOnAJBZPNUPj8yZ82c3hiV0,1181
-brainstate/util/_others.py,sha256=giUSIzIspCAIuyZL9g6ZFRo4DH6lk0aslTt2ADagAzc,31861
-brainstate/util/_others_test.py,sha256=0oCdhvFjEoc5de2SkjSmmyoMweJEDABz3pnFzy_hoI0,31497
-brainstate/util/_pretty_pytree.py,sha256=aACAe6gEwO5dFia_O-cg96biflcxONVakvFrlz99Pyk,48116
-brainstate/util/_pretty_pytree_test.py,sha256=b0Zg-0p3oaaVSrsdvZZXDWbiPejrQMtIlppQYUFbAtQ,23037
-brainstate/util/_pretty_repr.py,sha256=QE8BzOhCCMmpJIHsbtdx9p-unfVdya_eVDAin8-7dQU,14420
-brainstate/util/_pretty_repr_test.py,sha256=dpIYgjtVT3Bp8U7RyHCFKFsvQ2LYpJfdLNrjgzgL4jw,23818
-brainstate/util/filter.py,sha256=94NfApx2qCMVhoW1Pw0opTCkgTWjPx3pkWT7QA2C9h8,28355
-brainstate/util/filter_test.py,sha256=vLnxNrCEK77GNo0IOMKaNRP63TkbB9MwpjOZxIPqRxU,32564
-brainstate/util/struct.py,sha256=wmIvV-QF3qRrr6YP-7uw0ONhDJnXXHlkABY00DTo8_A,25415
-brainstate/util/struct_test.py,sha256=UCBqTzmsZowHHqy5Dc1mcoJfO6RMzHtp8zTV7nf-_DI,17128
-brainstate-0.2.1.dist-info/licenses/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.2.1.dist-info/METADATA,sha256=fIjy8AFu5OWak7iKP83FBotT9u-sKDCq8EGwuDi_qLw,4529
-brainstate-0.2.1.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-brainstate-0.2.1.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.2.1.dist-info/RECORD,,