brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/compile/_make_jaxpr_test.py

@@ -0,0 +1,134 @@
+# 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 __future__ import annotations
+
+import unittest
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+import brainstate as bst
+
+
+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 = bst.random.rand_like(arg[0])
+            return jnp.sum(temp + c)
+
+        key = bst.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))
+
+        bst.random.seed(1)
+        print(bst.random.DEFAULT.value)
+
+        jaxpr2, states = bst.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(bst.random.DEFAULT.value)
+
+    def test_StatefulFunction_1(self):
+        def func4(arg):  # Arg is a pair
+            temp = arg[0] + jnp.sin(arg[1]) * 3.
+            c = bst.random.rand_like(arg[0])
+            return jnp.sum(temp + c)
+
+        fun = bst.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 = bst.State(jnp.ones(10))
+
+        def f1(x):
+            st1.value = x + st1.value
+
+        def f2(x):
+            jaxpr = bst.compile.make_jaxpr(f1)(x)
+            c = 1. + x
+            return c
+
+        def f3(x):
+            jaxpr = bst.compile.make_jaxpr(f1)(x)
+            c = 1.
+            return c
+
+        print()
+        jaxpr = bst.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, _ = bst.compile.make_jaxpr(f3)(jnp.zeros(1))
+        print(jaxpr)
+        self.assertTrue(jnp.allclose(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
+                                     f3(jnp.zeros(1))))
+
+    def test_compar_jax_make_jaxpr2(self):
+        st1 = bst.State(jnp.ones(10))
+
+        def fa(x):
+            st1.value = x + st1.value
+
+        def ffa(x):
+            jaxpr, states = bst.compile.make_jaxpr(fa)(x)
+            c = 1. + x
+            return c
+
+        jaxpr, states = bst.compile.make_jaxpr(ffa)(jnp.zeros(1))
+        print()
+        print(jaxpr)
+        print(states)
+        print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
+        jaxpr = jax.make_jaxpr(ffa)(jnp.zeros(1))
+        print(jaxpr)
+        print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+
+    def test_compar_jax_make_jaxpr3(self):
+        def fa(x):
+            return 1.
+
+        jaxpr, states = bst.compile.make_jaxpr(fa)(jnp.zeros(1))
+        print()
+        print(jaxpr)
+        print(states)
+        # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        jaxpr = jax.make_jaxpr(fa)(jnp.zeros(1))
+        print(jaxpr)
+        # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+
+
+def test_return_states():
+    import jax.numpy
+
+    a = bst.State(jax.numpy.ones(3))
+
+    @bst.compile.jit
+    def f():
+        return a
+
+    with pytest.raises(ValueError):
+        f()

brainstate/compile/_progress_bar.py

@@ -0,0 +1,111 @@
+# 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 __future__ import annotations
+
+import copy
+from typing import Optional
+
+import jax
+
+try:
+    from tqdm.auto import tqdm
+except (ImportError, ModuleNotFoundError):
+    tqdm = None
+
+__all__ = [
+    'ProgressBar',
+]
+
+
+class ProgressBar(object):
+    __module__ = "brainstate.compile"
+
+    def __init__(self, freq: Optional[int] = None, count: Optional[int] = None, **kwargs):
+        self.print_freq = freq
+        self.print_count = count
+        if self.print_freq is not None and self.print_count is not None:
+            raise ValueError("Cannot specify both count and freq.")
+        for kwarg in ("total", "mininterval", "maxinterval", "miniters"):
+            kwargs.pop(kwarg, None)
+        self.kwargs = kwargs
+        if tqdm is None:
+            raise ImportError("tqdm is not installed.")
+
+    def init(self, n: int):
+        kwargs = copy.copy(self.kwargs)
+        freq = self.print_freq
+        count = self.print_count
+        if count is not None:
+            freq, remainder = divmod(n, count)
+            if freq == 0:
+                raise ValueError(f"Count {count} is too large for n {n}.")
+        elif freq is None:
+            if n > 20:
+                freq = int(n / 20)
+            else:
+                freq = 1
+            remainder = n % freq
+        else:
+            if freq < 1:
+                raise ValueError(f"Print rate should be > 0 got {freq}")
+            elif freq > n:
+                raise ValueError("Print rate should be less than the "
+                                 f"number of steps {n}, got {freq}")
+            remainder = n % freq
+        desc = kwargs.pop("desc", f"Running for {n:,} iterations")
+        message = kwargs.pop("message", desc)
+        return ProgressBarRunner(n, message, freq, remainder, **kwargs)
+
+
+class ProgressBarRunner(object):
+    __module__ = "brainstate.compile"
+
+    def __init__(self, n: int, message, print_freq: int, remainder: int, **kwargs):
+        self.tqdm_bars = {}
+        self.kwargs = kwargs
+        self.n = n
+        self.print_freq = print_freq
+        self.remainder = remainder
+        self.message = message
+
+    def _define_tqdm(self):
+        self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
+        self.tqdm_bars[0].set_description(self.message, refresh=False)
+
+    def _update_tqdm(self):
+        self.tqdm_bars[0].update(self.print_freq)
+
+    def _close_tqdm(self):
+        if self.remainder > 0:
+            self.tqdm_bars[0].update(self.remainder)
+        self.tqdm_bars[0].close()
+
+    def __call__(self, iter_num, *args, **kwargs):
+        _ = jax.lax.cond(
+            iter_num == 0,
+            lambda: jax.debug.callback(self._define_tqdm),
+            lambda: None,
+        )
+        _ = jax.lax.cond(
+            (iter_num + 1) % self.print_freq == 0,
+            lambda: jax.debug.callback(self._update_tqdm),
+            lambda: None,
+        )
+        _ = jax.lax.cond(
+            iter_num == self.n - 1,
+            lambda: jax.debug.callback(self._close_tqdm),
+            lambda: None,
+        )

brainstate/compile/_unvmap.py

@@ -0,0 +1,159 @@
+# 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 __future__ import annotations
+
+import jax
+import jax.core
+import jax.interpreters.batching as batching
+import jax.interpreters.mlir as mlir
+import jax.numpy as jnp
+
+from brainstate._utils import set_module_as
+
+__all__ = [
+    "unvmap",
+]
+
+
+@set_module_as('brainstate.augment')
+def unvmap(x, op: str = 'any'):
+    if op == 'all':
+        return unvmap_all(x)
+    elif op == 'any':
+        return unvmap_any(x)
+    elif op == 'none':
+        return _without_vmap(x)
+    elif op == 'max':
+        return unvmap_max(x)
+    else:
+        raise ValueError(f'Do not support type: {op}')
+
+
+# unvmap_all
+
+unvmap_all_p = jax.core.Primitive("unvmap_all")
+
+
+def unvmap_all(x):
+    """As `jnp.all`, but ignores batch dimensions."""
+    return unvmap_all_p.bind(x)
+
+
+def _unvmap_all_impl(x):
+    return jnp.all(x)
+
+
+def _unvmap_all_abstract_eval(x):
+    return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype)  # pyright: ignore
+
+
+def _unvmap_all_batch(x, batch_axes):
+    (x,) = x
+    return unvmap_all(x), batching.not_mapped
+
+
+unvmap_all_p.def_impl(_unvmap_all_impl)
+unvmap_all_p.def_abstract_eval(_unvmap_all_abstract_eval)
+batching.primitive_batchers[unvmap_all_p] = _unvmap_all_batch  # pyright: ignore
+mlir.register_lowering(
+    unvmap_all_p,
+    mlir.lower_fun(_unvmap_all_impl, multiple_results=False),
+)
+
+# unvmap_any
+
+unvmap_any_p = jax.core.Primitive("unvmap_any")
+
+
+def unvmap_any(x):
+    """As `jnp.any`, but ignores batch dimensions."""
+    return unvmap_any_p.bind(x)
+
+
+def _unvmap_any_impl(x):
+    return jnp.any(x)
+
+
+def _unvmap_any_abstract_eval(x):
+    return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype)  # pyright: ignore
+
+
+def _unvmap_any_batch(x, batch_axes):
+    (x,) = x
+    return unvmap_any(x), batching.not_mapped
+
+
+unvmap_any_p.def_impl(_unvmap_any_impl)
+unvmap_any_p.def_abstract_eval(_unvmap_any_abstract_eval)
+batching.primitive_batchers[unvmap_any_p] = _unvmap_any_batch  # pyright: ignore
+mlir.register_lowering(
+    unvmap_any_p,
+    mlir.lower_fun(_unvmap_any_impl, multiple_results=False),
+)
+
+# unvmap_max
+
+unvmap_max_p = jax.core.Primitive("unvmap_max")
+
+
+def unvmap_max(x):
+    """As `jnp.max`, but ignores batch dimensions."""
+    return unvmap_max_p.bind(x)
+
+
+def _unvmap_max_impl(x):
+    return jnp.max(x)
+
+
+def _unvmap_max_abstract_eval(x):
+    return jax.core.ShapedArray(shape=(), dtype=x.dtype)
+
+
+def _unvmap_max_batch(x, batch_axes):
+    (x,) = x
+    return unvmap_max(x), batching.not_mapped
+
+
+unvmap_max_p.def_impl(_unvmap_max_impl)
+unvmap_max_p.def_abstract_eval(_unvmap_max_abstract_eval)
+batching.primitive_batchers[unvmap_max_p] = _unvmap_max_batch  # pyright: ignore
+mlir.register_lowering(
+    unvmap_max_p,
+    mlir.lower_fun(_unvmap_max_impl, multiple_results=False),
+)
+
+
+def _without_vmap(x):
+    return _no_vmap_prim.bind(x)
+
+
+def _without_vmap_imp(x):
+    return x
+
+
+def _without_vmap_abs(x):
+    return x
+
+
+def _without_vmap_batch(x, batch_axes):
+    (x,) = x
+    return _without_vmap(x), batching.not_mapped
+
+
+_no_vmap_prim = jax.core.Primitive('no_vmap')
+_no_vmap_prim.def_impl(_without_vmap_imp)
+_no_vmap_prim.def_abstract_eval(_without_vmap_abs)
+batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch
+mlir.register_lowering(_no_vmap_prim, mlir.lower_fun(_without_vmap_imp, multiple_results=False))

brainstate/compile/_util.py

@@ -0,0 +1,147 @@
+# 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 __future__ import annotations
+
+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.write_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