brainstate 0.1.1__py2.py3-none-any.whl → 0.1.3__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (73) hide show
  1. brainstate/__init__.py +1 -1
  2. brainstate/_compatible_import.py +12 -9
  3. brainstate/_state.py +1 -1
  4. brainstate/augment/_autograd_test.py +132 -133
  5. brainstate/augment/_eval_shape_test.py +7 -9
  6. brainstate/augment/_mapping_test.py +75 -76
  7. brainstate/compile/_ad_checkpoint_test.py +6 -8
  8. brainstate/compile/_conditions_test.py +35 -36
  9. brainstate/compile/_error_if_test.py +10 -13
  10. brainstate/compile/_loop_collect_return_test.py +7 -9
  11. brainstate/compile/_loop_no_collection_test.py +7 -8
  12. brainstate/compile/_make_jaxpr.py +29 -14
  13. brainstate/compile/_make_jaxpr_test.py +20 -20
  14. brainstate/functional/_activations_test.py +61 -61
  15. brainstate/graph/_graph_node_test.py +16 -18
  16. brainstate/graph/_graph_operation_test.py +154 -156
  17. brainstate/init/_random_inits_test.py +20 -21
  18. brainstate/init/_regular_inits_test.py +4 -5
  19. brainstate/mixin.py +1 -14
  20. brainstate/nn/__init__.py +81 -17
  21. brainstate/nn/_collective_ops_test.py +8 -8
  22. brainstate/nn/{_interaction/_conv.py → _conv.py} +1 -1
  23. brainstate/nn/{_interaction/_conv_test.py → _conv_test.py} +31 -33
  24. brainstate/nn/{_dynamics/_state_delay.py → _delay.py} +6 -2
  25. brainstate/nn/{_elementwise/_dropout.py → _dropout.py} +1 -1
  26. brainstate/nn/{_elementwise/_dropout_test.py → _dropout_test.py} +9 -9
  27. brainstate/nn/{_dynamics/_dynamics_base.py → _dynamics.py} +139 -18
  28. brainstate/nn/{_dynamics/_dynamics_base_test.py → _dynamics_test.py} +14 -15
  29. brainstate/nn/{_elementwise/_elementwise.py → _elementwise.py} +1 -1
  30. brainstate/nn/_elementwise_test.py +169 -0
  31. brainstate/nn/{_interaction/_embedding.py → _embedding.py} +1 -1
  32. brainstate/nn/_exp_euler_test.py +5 -6
  33. brainstate/nn/{_event/_fixedprob_mv.py → _fixedprob_mv.py} +1 -1
  34. brainstate/nn/{_event/_fixedprob_mv_test.py → _fixedprob_mv_test.py} +0 -1
  35. brainstate/nn/{_dyn_impl/_inputs.py → _inputs.py} +4 -5
  36. brainstate/nn/{_interaction/_linear.py → _linear.py} +2 -5
  37. brainstate/nn/{_event/_linear_mv.py → _linear_mv.py} +1 -1
  38. brainstate/nn/{_event/_linear_mv_test.py → _linear_mv_test.py} +0 -1
  39. brainstate/nn/{_interaction/_linear_test.py → _linear_test.py} +15 -17
  40. brainstate/nn/{_event/__init__.py → _ltp.py} +7 -5
  41. brainstate/nn/_module_test.py +34 -37
  42. brainstate/nn/{_dyn_impl/_dynamics_neuron.py → _neuron.py} +2 -2
  43. brainstate/nn/{_dyn_impl/_dynamics_neuron_test.py → _neuron_test.py} +18 -19
  44. brainstate/nn/{_interaction/_normalizations.py → _normalizations.py} +1 -1
  45. brainstate/nn/{_interaction/_normalizations_test.py → _normalizations_test.py} +10 -12
  46. brainstate/nn/{_interaction/_poolings.py → _poolings.py} +1 -1
  47. brainstate/nn/{_interaction/_poolings_test.py → _poolings_test.py} +20 -22
  48. brainstate/nn/{_dynamics/_projection_base.py → _projection.py} +35 -3
  49. brainstate/nn/{_dyn_impl/_rate_rnns.py → _rate_rnns.py} +2 -2
  50. brainstate/nn/{_dyn_impl/_rate_rnns_test.py → _rate_rnns_test.py} +6 -7
  51. brainstate/nn/{_dyn_impl/_readout.py → _readout.py} +3 -3
  52. brainstate/nn/{_dyn_impl/_readout_test.py → _readout_test.py} +9 -10
  53. brainstate/nn/_stp.py +236 -0
  54. brainstate/nn/{_dyn_impl/_dynamics_synapse.py → _synapse.py} +17 -206
  55. brainstate/nn/{_dyn_impl/_dynamics_synapse_test.py → _synapse_test.py} +9 -10
  56. brainstate/nn/_synaptic_projection.py +133 -0
  57. brainstate/nn/{_dynamics/_synouts.py → _synouts.py} +4 -1
  58. brainstate/nn/{_dynamics/_synouts_test.py → _synouts_test.py} +4 -5
  59. brainstate/optim/_lr_scheduler_test.py +3 -3
  60. brainstate/optim/_optax_optimizer_test.py +8 -9
  61. brainstate/random/_rand_funs_test.py +183 -184
  62. brainstate/random/_rand_seed_test.py +10 -12
  63. {brainstate-0.1.1.dist-info → brainstate-0.1.3.dist-info}/METADATA +1 -1
  64. brainstate-0.1.3.dist-info/RECORD +131 -0
  65. brainstate/nn/_dyn_impl/__init__.py +0 -42
  66. brainstate/nn/_dynamics/__init__.py +0 -37
  67. brainstate/nn/_elementwise/__init__.py +0 -22
  68. brainstate/nn/_elementwise/_elementwise_test.py +0 -171
  69. brainstate/nn/_interaction/__init__.py +0 -41
  70. brainstate-0.1.1.dist-info/RECORD +0 -133
  71. {brainstate-0.1.1.dist-info → brainstate-0.1.3.dist-info}/LICENSE +0 -0
  72. {brainstate-0.1.1.dist-info → brainstate-0.1.3.dist-info}/WHEEL +0 -0
  73. {brainstate-0.1.1.dist-info → brainstate-0.1.3.dist-info}/top_level.txt +0 -0
brainstate/graph/_graph_node_test.py CHANGED
@@ -13,63 +13,61 @@
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
15
 
16
- from __future__ import annotations
17
-
18
16
  import unittest
19
17
 
20
- import brainstate as bst
18
+ import brainstate
21
19
 
22
20
 
23
21
  class TestSequential(unittest.TestCase):
24
22
  def test1(self):
25
- s = bst.graph.Sequential(bst.nn.Linear(1, 2),
26
- bst.nn.Linear(2, 3))
27
- graphdef, states = bst.graph.treefy_split(s)
23
+ s = brainstate.graph.Sequential(brainstate.nn.Linear(1, 2),
24
+ brainstate.nn.Linear(2, 3))
25
+ graphdef, states = brainstate.graph.treefy_split(s)
28
26
  print(states)
29
27
  self.assertTrue(len(states.to_flat()) == 2)
30
28
 
31
29
 
32
30
  class TestStateRetrieve(unittest.TestCase):
33
31
  def test_list_of_states_1(self):
34
- class Model(bst.graph.Node):
32
+ class Model(brainstate.graph.Node):
35
33
  def __init__(self):
36
34
  self.a = [1, 2, 3]
37
- self.b = [bst.State(1), bst.State(2), bst.State(3)]
35
+ self.b = [brainstate.State(1), brainstate.State(2), brainstate.State(3)]
38
36
 
39
37
  m = Model()
40
- graphdef, states = bst.graph.treefy_split(m)
38
+ graphdef, states = brainstate.graph.treefy_split(m)
41
39
  print(states.to_flat())
42
40
  self.assertTrue(len(states.to_flat()) == 3)
43
41
 
44
42
  def test_list_of_states_2(self):
45
- class Model(bst.graph.Node):
43
+ class Model(brainstate.graph.Node):
46
44
  def __init__(self):
47
45
  self.a = [1, 2, 3]
48
- self.b = [bst.State(1), [bst.State(2), bst.State(3)]]
46
+ self.b = [brainstate.State(1), [brainstate.State(2), brainstate.State(3)]]
49
47
 
50
48
  m = Model()
51
- graphdef, states = bst.graph.treefy_split(m)
49
+ graphdef, states = brainstate.graph.treefy_split(m)
52
50
  print(states.to_flat())
53
51
  self.assertTrue(len(states.to_flat()) == 3)
54
52
 
55
53
  def test_list_of_node_1(self):
56
- class Model(bst.graph.Node):
54
+ class Model(brainstate.graph.Node):
57
55
  def __init__(self):
58
56
  self.a = [1, 2, 3]
59
- self.b = [bst.nn.Linear(1, 2), bst.nn.Linear(2, 3)]
57
+ self.b = [brainstate.nn.Linear(1, 2), brainstate.nn.Linear(2, 3)]
60
58
 
61
59
  m = Model()
62
- graphdef, states = bst.graph.treefy_split(m)
60
+ graphdef, states = brainstate.graph.treefy_split(m)
63
61
  print(states.to_flat())
64
62
  self.assertTrue(len(states.to_flat()) == 2)
65
63
 
66
64
  def test_list_of_node_2(self):
67
- class Model(bst.graph.Node):
65
+ class Model(brainstate.graph.Node):
68
66
  def __init__(self):
69
67
  self.a = [1, 2, 3]
70
- self.b = [bst.nn.Linear(1, 2), [bst.nn.Linear(2, 3)], (bst.nn.Linear(3, 4), bst.nn.Linear(4, 5))]
68
+ self.b = [brainstate.nn.Linear(1, 2), [brainstate.nn.Linear(2, 3)], (brainstate.nn.Linear(3, 4), brainstate.nn.Linear(4, 5))]
71
69
 
72
70
  m = Model()
73
- graphdef, states = bst.graph.treefy_split(m)
71
+ graphdef, states = brainstate.graph.treefy_split(m)
74
72
  print(states.to_flat())
75
73
  self.assertTrue(len(states.to_flat()) == 4)