brainstate 0.2.0__tar.gz → 0.2.2__tar.gz

{brainstate-0.2.0 → brainstate-0.2.2}/PKG-INFO

@@ -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 → brainstate-0.2.2}/brainstate/__init__.py

@@ -17,8 +17,8 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.2.0"
-__versio_info__ = (0, 2, 0)
+__version__ = "0.2.2"
+__versio_info__ = (0, 2, 2)
 
 from . import environ
 from . import graph
@@ -45,12 +45,10 @@ _augment_apis = {
     'jacobian': 'brainstate.transform._autograd',
     'jacrev': 'brainstate.transform._autograd',
     'jacfwd': 'brainstate.transform._autograd',
-    'abstract_init': 'brainstate.transform._eval_shape',
     'vmap': 'brainstate.transform._mapping',
     'pmap': 'brainstate.transform._mapping',
     'map': 'brainstate.transform._mapping',
     'vmap_new_states': 'brainstate.transform._mapping',
-    'restore_rngs': 'brainstate.transform._random',
 }
 
 augment = create_deprecated_module_proxy(

{brainstate-0.2.0 → brainstate-0.2.2}/brainstate/_deprecation_test.py

@@ -51,8 +51,8 @@ class TestDeprecatedAugmentModule(unittest.TestCase):
         # Check that expected APIs are available
         expected_apis = [
             'GradientTransform', 'grad', 'vector_grad', 'hessian', 'jacobian',
-            'jacrev', 'jacfwd', 'abstract_init', 'vmap', 'pmap', 'map',
-            'vmap_new_states', 'restore_rngs'
+            'jacrev', 'jacfwd', 'vmap', 'pmap', 'map',
+            'vmap_new_states',
         ]
 
         for api in expected_apis:
@@ -1379,12 +1379,10 @@ class TestDeprecatedAugment(unittest.TestCase):
             'jacobian',
             'jacrev',
             'jacfwd',
-            'abstract_init',
             'vmap',
             'pmap',
             'map',
             'vmap_new_states',
-            'restore_rngs',
         ]
 
         for func_name in augment_funcs:
@@ -1503,26 +1501,6 @@ class TestDeprecatedAugment(unittest.TestCase):
             vmap_new_states = brainstate.augment.vmap_new_states
             self.assertIsNotNone(vmap_new_states)
 
-    def test_abstract_init(self):
-        """Test abstract_init function."""
-        with warnings.catch_warnings(record=True):
-            warnings.simplefilter("always")
-            import brainstate
-
-            # Test abstract_init
-            abstract_init = brainstate.augment.abstract_init
-            self.assertIsNotNone(abstract_init)
-
-    def test_restore_rngs(self):
-        """Test restore_rngs function."""
-        with warnings.catch_warnings(record=True):
-            warnings.simplefilter("always")
-            import brainstate
-
-            # Test restore_rngs
-            restore_rngs = brainstate.augment.restore_rngs
-            self.assertIsNotNone(restore_rngs)
-
     def test_module_attributes(self):
         """Test module-level attributes."""
         with warnings.catch_warnings(record=True):