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

brainstate/compile/_make_jaxpr_test.py

@@ -1,156 +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 unittest
-
-import jax
-import jax.numpy as jnp
-import pytest
-
-import brainstate
-from brainstate._compatible_import import jaxpr_as_fun
-
-
-class TestMakeJaxpr(unittest.TestCase):
-    def test_compar_jax_make_jaxpr(self):
-        def func4(arg):  # Arg is a pair
-            temp = arg[0] + jnp.sin(arg[1]) * 3.
-            c = brainstate.random.rand_like(arg[0])
-            return jnp.sum(temp + c)
-
-        key = brainstate.random.DEFAULT.value
-        jaxpr = jax.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
-        print(jaxpr)
-        self.assertTrue(len(jaxpr.in_avals) == 2)
-        self.assertTrue(len(jaxpr.consts) == 1)
-        self.assertTrue(len(jaxpr.out_avals) == 1)
-        self.assertTrue(jnp.allclose(jaxpr.consts[0], key))
-
-        brainstate.random.seed(1)
-        print(brainstate.random.DEFAULT.value)
-
-        jaxpr2, states = brainstate.compile.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
-        print(jaxpr2)
-        self.assertTrue(len(jaxpr2.in_avals) == 3)
-        self.assertTrue(len(jaxpr2.out_avals) == 2)
-        self.assertTrue(len(jaxpr2.consts) == 0)
-        print(brainstate.random.DEFAULT.value)
-
-    def test_StatefulFunction_1(self):
-        def func4(arg):  # Arg is a pair
-            temp = arg[0] + jnp.sin(arg[1]) * 3.
-            c = brainstate.random.rand_like(arg[0])
-            return jnp.sum(temp + c)
-
-        fun = brainstate.compile.StatefulFunction(func4).make_jaxpr((jnp.zeros(8), jnp.ones(8)))
-        print(fun.get_states())
-        print(fun.get_jaxpr())
-
-    def test_StatefulFunction_2(self):
-        st1 = brainstate.State(jnp.ones(10))
-
-        def f1(x):
-            st1.value = x + st1.value
-
-        def f2(x):
-            jaxpr = brainstate.compile.make_jaxpr(f1)(x)
-            c = 1. + x
-            return c
-
-        def f3(x):
-            jaxpr = brainstate.compile.make_jaxpr(f1)(x)
-            c = 1.
-            return c
-
-        print()
-        jaxpr = brainstate.compile.make_jaxpr(f1)(jnp.zeros(1))
-        print(jaxpr)
-        jaxpr = jax.make_jaxpr(f2)(jnp.zeros(1))
-        print(jaxpr)
-        jaxpr = jax.make_jaxpr(f3)(jnp.zeros(1))
-        print(jaxpr)
-        jaxpr, _ = brainstate.compile.make_jaxpr(f3)(jnp.zeros(1))
-        print(jaxpr)
-        self.assertTrue(jnp.allclose(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
-                                     f3(jnp.zeros(1))))
-
-    def test_compare_jax_make_jaxpr2(self):
-        st1 = brainstate.State(jnp.ones(10))
-
-        def fa(x):
-            st1.value = x + st1.value
-
-        def ffa(x):
-            jaxpr, states = brainstate.compile.make_jaxpr(fa)(x)
-            c = 1. + x
-            return c
-
-        jaxpr, states = brainstate.compile.make_jaxpr(ffa)(jnp.zeros(1))
-        print()
-        print(jaxpr)
-        print(states)
-        print(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
-        jaxpr = jax.make_jaxpr(ffa)(jnp.zeros(1))
-        print(jaxpr)
-        print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
-
-    def test_compare_jax_make_jaxpr3(self):
-        def fa(x):
-            return 1.
-
-        jaxpr, states = brainstate.compile.make_jaxpr(fa)(jnp.zeros(1))
-        print()
-        print(jaxpr)
-        print(states)
-        # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
-        jaxpr = jax.make_jaxpr(fa)(jnp.zeros(1))
-        print(jaxpr)
-        # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
-
-    def test_static_argnames(self):
-        def func4(a, b):  # Arg is a pair
-            temp = a + jnp.sin(b) * 3.
-            c = brainstate.random.rand_like(a)
-            return jnp.sum(temp + c)
-
-        jaxpr, states = brainstate.compile.make_jaxpr(func4, static_argnames='b')(jnp.zeros(8), 1.)
-        print()
-        print(jaxpr)
-        print(states)
-
-    def test_state_in(self):
-        def f(a):
-            return a.value
-
-        with pytest.raises(ValueError):
-            brainstate.compile.StatefulFunction(f).make_jaxpr(brainstate.State(1.))
-
-    def test_state_out(self):
-        def f(a):
-            return brainstate.State(a)
-
-        with pytest.raises(ValueError):
-            brainstate.compile.StatefulFunction(f).make_jaxpr(1.)
-
-    def test_return_states(self):
-        a = brainstate.State(jnp.ones(3))
-
-        @brainstate.compile.jit
-        def f():
-            return a
-
-        with pytest.raises(ValueError):
-            f()

brainstate/compile/_util.py

@@ -1,147 +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.
-# ==============================================================================
-
-from functools import wraps
-from typing import Sequence, Tuple
-
-from brainstate._state import StateTraceStack
-from brainstate.typing import PyTree
-from ._make_jaxpr import StatefulFunction
-
-
-def write_back_state_values(
-    state_trace: StateTraceStack,
-    read_state_vals: Sequence[PyTree],
-    write_state_vals: Sequence[PyTree],
-):
-    assert len(state_trace.states) == len(state_trace.been_writen) == len(read_state_vals) == len(write_state_vals)
-    for st, write, val_r, val_w in zip(state_trace.states, state_trace.been_writen, read_state_vals, write_state_vals):
-        if write:
-            st.value = val_w
-        else:
-            st.restore_value(val_r)
-
-
-def wrap_single_fun_in_multi_branches(
-    stateful_fun: StatefulFunction,
-    merged_state_trace: StateTraceStack,
-    read_state_vals: Sequence[PyTree | None],
-    return_states: bool = True
-):
-    state_ids_belong_to_this_fun = {id(st): st for st in stateful_fun.get_states()}
-
-    @wraps(stateful_fun.fun)
-    def wrapped_branch(write_state_vals, *operands):
-        # "write_state_vals" should have the same length as "merged_state_trace.states"
-        assert len(merged_state_trace.states) == len(write_state_vals) == len(read_state_vals)
-
-        # get all state values needed for this function, which is a subset of "write_state_vals"
-        st_vals_for_this_fun = []
-        for write, st, val_w, val_r in zip(merged_state_trace.been_writen,
-                                           merged_state_trace.states,
-                                           write_state_vals,
-                                           read_state_vals):
-            if id(st) in state_ids_belong_to_this_fun:
-                st_vals_for_this_fun.append(val_w if write else val_r)
-
-        # call this function
-        new_state_vals, out = stateful_fun.jaxpr_call(st_vals_for_this_fun, *operands)
-        assert len(new_state_vals) == len(st_vals_for_this_fun)
-
-        if return_states:
-            # get all written state values
-            new_state_vals = {id(st): val for st, val in zip(stateful_fun.get_states(), new_state_vals)}
-            write_state_vals = tuple([
-                (new_state_vals[id(st)] if id(st) in state_ids_belong_to_this_fun else w_val)
-                if write else None
-                for write, st, w_val in zip(merged_state_trace.been_writen,
-                                            merged_state_trace.states,
-                                            write_state_vals)
-            ])
-            return write_state_vals, out
-        return out
-
-    return wrapped_branch
-
-
-def wrap_single_fun_in_multi_branches_while_loop(
-    stateful_fun: StatefulFunction,
-    merged_state_trace: StateTraceStack,
-    read_state_vals: Sequence[PyTree | None],
-    return_states: bool = True
-):
-    state_ids_belong_to_this_fun = {id(st): st for st in stateful_fun.get_states()}
-
-    @wraps(stateful_fun.fun)
-    def wrapped_branch(init_val):
-        write_state_vals, init_val = init_val
-        # "write_state_vals" should have the same length as "merged_state_trace.states"
-        assert len(merged_state_trace.states) == len(write_state_vals) == len(read_state_vals)
-
-        # get all state values needed for this function, which is a subset of "write_state_vals"
-        st_vals_for_this_fun = []
-        for write, st, val_w, val_r in zip(merged_state_trace.been_writen,
-                                           merged_state_trace.states,
-                                           write_state_vals,
-                                           read_state_vals):
-            if id(st) in state_ids_belong_to_this_fun:
-                st_vals_for_this_fun.append(val_w if write else val_r)
-
-        # call this function
-        new_state_vals, out = stateful_fun.jaxpr_call(st_vals_for_this_fun, init_val)
-        assert len(new_state_vals) == len(st_vals_for_this_fun)
-
-        if return_states:
-            # get all written state values
-            new_state_vals = {id(st): val for st, val in zip(stateful_fun.get_states(), new_state_vals)}
-            write_state_vals = tuple([
-                (new_state_vals[id(st)] if id(st) in state_ids_belong_to_this_fun else w_val)
-                if write else None
-                for write, st, w_val in zip(merged_state_trace.been_writen,
-                                            merged_state_trace.states,
-                                            write_state_vals)
-            ])
-            return write_state_vals, out
-        return out
-
-    return wrapped_branch
-
-
-def wrap_single_fun(
-    stateful_fun: StatefulFunction,
-    been_writen: Tuple[bool],
-    read_state_vals: Tuple[PyTree | None],
-):
-    @wraps(stateful_fun.fun)
-    def wrapped_fun(new_carry, inputs):
-        writen_state_vals, carry = new_carry
-        assert len(been_writen) == len(writen_state_vals) == len(read_state_vals)
-
-        # collect all written and read states
-        state_vals = [
-            written_val if written else read_val
-            for written, written_val, read_val in zip(been_writen, writen_state_vals, read_state_vals)
-        ]
-
-        # call the jaxpr
-        state_vals, (carry, out) = stateful_fun.jaxpr_call(state_vals, carry, inputs)
-
-        # only return the written states
-        writen_state_vals = tuple([val if written else None for written, val in zip(been_writen, state_vals)])
-
-        # return
-        return (writen_state_vals, carry), out
-
-    return wrapped_fun

brainstate/functional/__init__.py

@@ -1,27 +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.
-# ==============================================================================
-
-from brainstate.nn._activations import *
-from brainstate.nn._activations import __all__ as act_all
-from brainstate.nn._normalizations import weight_standardization
-from brainstate.nn._others import clip_grad_norm
-
-__all__ = ['weight_standardization', 'clip_grad_norm'] + act_all
-del act_all
-
-if __name__ == '__main__':
-    relu
-    clip_grad_norm
-    weight_standardization

brainstate/graph/_graph_node.py

@@ -1,244 +0,0 @@
-# The file is adapted from the Flax library (https://github.com/google/flax).
-# The credit should go to the Flax authors.
-#
-# Copyright 2024 The Flax Authors.
-#
-# 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.
-
-from abc import ABCMeta
-from copy import deepcopy
-from typing import Any, Callable, Type, TypeVar, Tuple, TYPE_CHECKING, Mapping, Iterator, Sequence
-
-import brainunit as u
-import jax
-import numpy as np
-
-from brainstate._state import State, TreefyState
-from brainstate.typing import Key
-from brainstate.util.pretty_pytree import PrettyObject
-from ._graph_operation import register_graph_node_type
-
-__all__ = [
-    'Node', 'Dict', 'List', 'Sequential',
-]
-
-G = TypeVar('G', bound='Node')
-A = TypeVar('A')
-
-
-class GraphNodeMeta(ABCMeta):
-    if not TYPE_CHECKING:
-        def __call__(cls, *args: Any, **kwargs: Any) -> Any:
-            node = cls.__new__(cls, *args, **kwargs)
-            node.__init__(*args, **kwargs)
-            return node
-
-
-class Node(PrettyObject, metaclass=GraphNodeMeta):
-    """
-    Base class for all graph nodes.
-
-    This class provides the following functionalities:
-    - Register the node type with the graph tool.
-    - Prevent mutation of the node from different trace level.
-    - Provide a pretty repr for the node.
-    - Provide a treescope repr for the node.
-    - Deepcopy the node.
-
-    """
-
-    graph_invisible_attrs = ()
-
-    def __init_subclass__(cls) -> None:
-        super().__init_subclass__()
-
-        register_graph_node_type(
-            type=cls,
-            flatten=_node_flatten,
-            set_key=_node_set_key,
-            pop_key=_node_pop_key,
-            create_empty=_node_create_empty,
-            clear=_node_clear,
-        )
-
-    def __deepcopy__(self: G, memo=None) -> G:
-        """
-        Deepcopy the object.
-        """
-        from ._graph_operation import treefy_split, treefy_merge
-
-        graphdef, state = treefy_split(self)
-        graphdef = deepcopy(graphdef)
-        state = deepcopy(state)
-        return treefy_merge(graphdef, state)
-
-
-class String:
-    def __init__(self, msg):
-        self.msg = msg
-
-    def __repr__(self):
-        return self.msg
-
-
-def _to_shape_dtype(value):
-    if isinstance(value, State):
-        return value.replace(jax.tree.map(_to_shape_dtype, value.value))
-    elif isinstance(value, (np.ndarray, jax.Array)):
-        return String(f'Array(shape={value.shape}, dtype={value.dtype.name})')
-    elif isinstance(value, u.Quantity):
-        return String(f'Quantity(mantissa=Array(shape={value.shape}, dtype={value.dtype.name}), unit={value.unit})')
-    return value
-
-
-# -------------------------------
-# Graph Definition
-# -------------------------------
-
-
-def _node_flatten(
-    node: Node
-) -> Tuple[Tuple[Tuple[str, Any], ...], Tuple[Type]]:
-    # graph_invisible_attrs = getattr(node, 'graph_invisible_attrs', ())
-    # graph_invisible_attrs = tuple(graph_invisible_attrs) + ('_trace_state',)
-    graph_invisible_attrs = ('_trace_state',)
-    nodes = sorted(
-        (key, value) for key, value in vars(node).items()
-        if (key not in graph_invisible_attrs)
-    )
-    return nodes, (type(node),)
-
-
-def _node_set_key(
-    node: Node,
-    key: Key,
-    value: Any
-) -> None:
-    if not isinstance(key, str):
-        raise KeyError(f'Invalid key: {key!r}')
-    elif (
-        hasattr(node, key)
-        and isinstance(state := getattr(node, key), State)
-        and isinstance(value, TreefyState)
-    ):
-        state.update_from_ref(value)
-    else:
-        setattr(node, key, value)
-
-
-def _node_pop_key(
-    node: Node,
-    key: Key
-):
-    if not isinstance(key, str):
-        raise KeyError(f'Invalid key: {key!r}')
-    return vars(node).pop(key)
-
-
-def _node_create_empty(
-    static: tuple[Type[G],]
-) -> G:
-    node_type, = static
-    node = object.__new__(node_type)
-    return node
-
-
-def _node_clear(node: Node):
-    module_state = node._trace_state
-    module_vars = vars(node)
-    module_vars.clear()
-    module_vars['_trace_state'] = module_state
-
-
-class Dict(Node, Mapping[str, A]):
-    """
-    A dictionary node.
-    """
-
-    def __init__(self, *args, **kwargs):
-        for name, value in dict(*args, **kwargs).items():
-            setattr(self, name, value)
-
-    def __getitem__(self, key) -> A:
-        return getattr(self, key)
-
-    def __setitem__(self, key, value):
-        setattr(self, key, value)
-
-    def __getattr__(self, key) -> A:
-        return super().__getattribute__(key)
-
-    def __setattr__(self, key, value):
-        super().__setattr__(key, value)
-
-    def __iter__(self) -> Iterator[str]:
-        return (k for k in vars(self) if k != '_object__state')
-
-    def __len__(self) -> int:
-        return len(vars(self))
-
-
-class List(Node):
-    """
-    A list node.
-    """
-
-    def __init__(self, seq=()):
-        vars(self).update({str(i): item for i, item in enumerate(seq)})
-
-    def __getitem__(self, idx):
-        return getattr(self, str(idx))
-
-    def __setitem__(self, idx, value):
-        setattr(self, str(idx), value)
-
-    def __iter__(self):
-        return iter(vars(self).values())
-
-    def __len__(self):
-        return len(vars(self))
-
-    def __add__(self, other: Sequence[A]) -> 'List[A]':
-        return List(list(self) + list(other))
-
-    def append(self, value):
-        self[len(vars(self))] = value
-
-    def extend(self, values):
-        for value in values:
-            self.append(value)
-
-
-class Sequential(Node):
-    def __init__(self, *fns: Callable[..., Any]):
-        self.layers = list(fns)
-
-    def __call__(self, *args, **kwargs) -> Any:
-        output: Any = None
-
-        for i, f in enumerate(self.layers):
-            if not callable(f):
-                raise TypeError(f'Sequence[{i}] is not callable: {f}')
-            if i > 0:
-                if isinstance(output, tuple):
-                    args = output
-                    kwargs = {}
-                elif isinstance(output, dict):
-                    args = ()
-                    kwargs = output
-                else:
-                    args = (output,)
-                    kwargs = {}
-            output = f(*args, **kwargs)
-
-        return output

brainstate/graph/_graph_node_test.py

@@ -1,73 +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 unittest
-
-import brainstate
-
-
-class TestSequential(unittest.TestCase):
-    def test1(self):
-        s = brainstate.graph.Sequential(brainstate.nn.Linear(1, 2),
-                                        brainstate.nn.Linear(2, 3))
-        graphdef, states = brainstate.graph.treefy_split(s)
-        print(states)
-        self.assertTrue(len(states.to_flat()) == 2)
-
-
-class TestStateRetrieve(unittest.TestCase):
-    def test_list_of_states_1(self):
-        class Model(brainstate.graph.Node):
-            def __init__(self):
-                self.a = [1, 2, 3]
-                self.b = [brainstate.State(1), brainstate.State(2), brainstate.State(3)]
-
-        m = Model()
-        graphdef, states = brainstate.graph.treefy_split(m)
-        print(states.to_flat())
-        self.assertTrue(len(states.to_flat()) == 3)
-
-    def test_list_of_states_2(self):
-        class Model(brainstate.graph.Node):
-            def __init__(self):
-                self.a = [1, 2, 3]
-                self.b = [brainstate.State(1), [brainstate.State(2), brainstate.State(3)]]
-
-        m = Model()
-        graphdef, states = brainstate.graph.treefy_split(m)
-        print(states.to_flat())
-        self.assertTrue(len(states.to_flat()) == 3)
-
-    def test_list_of_node_1(self):
-        class Model(brainstate.graph.Node):
-            def __init__(self):
-                self.a = [1, 2, 3]
-                self.b = [brainstate.nn.Linear(1, 2), brainstate.nn.Linear(2, 3)]
-
-        m = Model()
-        graphdef, states = brainstate.graph.treefy_split(m)
-        print(states.to_flat())
-        self.assertTrue(len(states.to_flat()) == 2)
-
-    def test_list_of_node_2(self):
-        class Model(brainstate.graph.Node):
-            def __init__(self):
-                self.a = [1, 2, 3]
-                self.b = [brainstate.nn.Linear(1, 2), [brainstate.nn.Linear(2, 3)], (brainstate.nn.Linear(3, 4), brainstate.nn.Linear(4, 5))]
-
-        m = Model()
-        graphdef, states = brainstate.graph.treefy_split(m)
-        print(states.to_flat())
-        self.assertTrue(len(states.to_flat()) == 4)