brainstate 0.1.7__py2.py3-none-any.whl → 0.1.9__py2.py3-none-any.whl

brainstate/compile/_unvmap.py

@@ -1,159 +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.
-# ==============================================================================
-
-import jax
-import jax.core
-import jax.interpreters.batching as batching
-import jax.interpreters.mlir as mlir
-import jax.numpy as jnp
-
-from brainstate._compatible_import import Primitive
-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 = 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 = 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 = 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 = 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))
+# 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 jax
+import jax.core
+import jax.interpreters.batching as batching
+import jax.interpreters.mlir as mlir
+import jax.numpy as jnp
+
+from brainstate._compatible_import import Primitive
+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 = 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 = 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 = 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 = 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

@@ -1,147 +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 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
+# 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