brainstate 0.0.2.post20241009__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

brainstate/compile/_conditions.py

@@ -0,0 +1,259 @@
+# 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 collections.abc import Callable, Sequence
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brainstate._utils import set_module_as
+from ._error_if import jit_error_if
+from ._make_jaxpr import StatefulFunction
+from ._util import wrap_single_fun_in_multi_branches, write_back_state_values
+
+__all__ = [
+    'cond', 'switch', 'ifelse',
+]
+
+
+@set_module_as('brainstate.compile')
+def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
+    """
+    Conditionally apply ``true_fun`` or ``false_fun``.
+
+    Provided arguments are correctly typed, ``cond()`` has equivalent
+    semantics to this Python implementation, where ``pred`` must be a
+    scalar type::
+
+      def cond(pred, true_fun, false_fun, *operands):
+        if pred:
+          return true_fun(*operands)
+        else:
+          return false_fun(*operands)
+
+
+    In contrast with :func:`jax.lax.select`, using ``cond`` indicates that only one of
+    the two branches is executed (up to compiler rewrites and optimizations).
+    However, when transformed with :func:`~jax.vmap` to operate over a batch of
+    predicates, ``cond`` is converted to :func:`~jax.lax.select`.
+
+    Args:
+      pred: Boolean scalar type, indicating which branch function to apply.
+      true_fun: Function (A -> B), to be applied if ``pred`` is True.
+      false_fun: Function (A -> B), to be applied if ``pred`` is False.
+      operands: Operands (A) input to either branch depending on ``pred``. The
+        type can be a scalar, array, or any pytree (nested Python tuple/list/dict)
+        thereof.
+
+    Returns:
+      Value (B) of either ``true_fun(*operands)`` or ``false_fun(*operands)``,
+      depending on the value of ``pred``. The type can be a scalar, array, or any
+      pytree (nested Python tuple/list/dict) thereof.
+    """
+    if not (callable(true_fun) and callable(false_fun)):
+        raise TypeError("true_fun and false_fun arguments should be callable.")
+
+    if pred is None:
+        raise TypeError("cond predicate is None")
+    if isinstance(pred, Sequence) or np.ndim(pred) != 0:
+        raise TypeError(f"Pred must be a scalar, got {pred} of " +
+                        (f"type {type(pred)}" if isinstance(pred, Sequence)
+                         else f"shape {np.shape(pred)}."))
+
+    # check pred
+    try:
+        pred_dtype = jax.dtypes.result_type(pred)
+    except TypeError as err:
+        raise TypeError("Pred type must be either boolean or number, got {}.".format(pred)) from err
+    if pred_dtype.kind != 'b':
+        if pred_dtype.kind in 'iuf':
+            pred = pred != 0
+        else:
+            raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
+
+    # not jit
+    if jax.config.jax_disable_jit and isinstance(jax.core.get_aval(pred), jax.core.ConcreteArray):
+        if pred:
+            return true_fun(*operands)
+        else:
+            return false_fun(*operands)
+
+    # evaluate jaxpr
+    with jax.ensure_compile_time_eval():
+        stateful_true = StatefulFunction(true_fun).make_jaxpr(*operands)
+        stateful_false = StatefulFunction(false_fun).make_jaxpr(*operands)
+
+    # state trace and state values
+    state_trace = stateful_true.get_state_trace() + stateful_false.get_state_trace()
+    read_state_vals = state_trace.get_read_state_values(True)
+    write_state_vals = state_trace.get_write_state_values(True)
+
+    # wrap the functions
+    true_fun = wrap_single_fun_in_multi_branches(stateful_true, state_trace, read_state_vals, True)
+    false_fun = wrap_single_fun_in_multi_branches(stateful_false, state_trace, read_state_vals, True)
+
+    # cond
+    write_state_vals, out = jax.lax.cond(pred, true_fun, false_fun, write_state_vals, *operands)
+
+    # assign the written state values and restore the read state values
+    write_back_state_values(state_trace, read_state_vals, write_state_vals)
+    return out
+
+
+@set_module_as('brainstate.compile')
+def switch(index, branches: Sequence[Callable], *operands):
+    """
+    Apply exactly one of ``branches`` given by ``index``.
+
+    If ``index`` is out of bounds, it is clamped to within bounds.
+
+    Has the semantics of the following Python::
+
+      def switch(index, branches, *operands):
+        index = clamp(0, index, len(branches) - 1)
+        return branches[index](*operands)
+
+    Internally this wraps XLA's `Conditional
+    <https://www.tensorflow.org/xla/operation_semantics#conditional>`_
+    operator. However, when transformed with :func:`~jax.vmap` to operate over a
+    batch of predicates, ``cond`` is converted to :func:`~jax.lax.select`.
+
+    Args:
+      index: Integer scalar type, indicating which branch function to apply.
+      branches: Sequence of functions (A -> B) to be applied based on ``index``.
+      operands: Operands (A) input to whichever branch is applied.
+
+    Returns:
+      Value (B) of ``branch(*operands)`` for the branch that was selected based
+      on ``index``.
+    """
+    # check branches
+    if not all(callable(branch) for branch in branches):
+        raise TypeError("branches argument should be a sequence of callables.")
+
+    # check index
+    if len(np.shape(index)) != 0:
+        raise TypeError(f"Branch index must be scalar, got {index} of shape {np.shape(index)}.")
+    try:
+        index_dtype = jax.dtypes.result_type(index)
+    except TypeError as err:
+        msg = f"Index type must be an integer, got {index}."
+        raise TypeError(msg) from err
+    if index_dtype.kind not in 'iu':
+        raise TypeError(f"Index type must be an integer, got {index} as {index_dtype}")
+
+    # format branches
+    branches = tuple(branches)
+    if len(branches) == 0:
+        raise ValueError("Empty branch sequence")
+    elif len(branches) == 1:
+        return branches[0](*operands)
+
+    # format index
+    index = jax.lax.convert_element_type(index, np.int32)
+    lo = np.array(0, np.int32)
+    hi = np.array(len(branches) - 1, np.int32)
+    index = jax.lax.clamp(lo, index, hi)
+
+    # not jit
+    if jax.config.jax_disable_jit and isinstance(jax.core.core.get_aval(index), jax.core.ConcreteArray):
+        return branches[int(index)](*operands)
+
+    # evaluate jaxpr
+    with jax.ensure_compile_time_eval():
+        wrapped_branches = [StatefulFunction(branch) for branch in branches]
+        for wrapped_branch in wrapped_branches:
+            wrapped_branch.make_jaxpr(*operands)
+
+    # wrap the functions
+    state_trace = wrapped_branches[0].get_state_trace() + wrapped_branches[1].get_state_trace()
+    state_trace.merge(*[wrapped_branch.get_state_trace() for wrapped_branch in wrapped_branches[2:]])
+    read_state_vals = state_trace.get_read_state_values(True)
+    write_state_vals = state_trace.get_write_state_values(True)
+    branches = [
+        wrap_single_fun_in_multi_branches(wrapped_branch, state_trace, read_state_vals, True)
+        for wrapped_branch in wrapped_branches
+    ]
+
+    # switch
+    write_state_vals, out = jax.lax.switch(index, branches, write_state_vals, *operands)
+
+    # write back state values or restore them
+    write_back_state_values(state_trace, read_state_vals, write_state_vals)
+    return out
+
+
+@set_module_as('brainstate.compile')
+def ifelse(conditions, branches, *operands, check_cond: bool = True):
+    """
+    ``If-else`` control flows looks like native Pythonic programming.
+
+    Examples
+    --------
+
+    >>> import brainstate as bst
+    >>> def f(a):
+    >>>    return bst.compile.ifelse(conditions=[a > 10, a > 5, a > 2, a > 0],
+    >>>                               branches=[lambda: 1,
+    >>>                                         lambda: 2,
+    >>>                                         lambda: 3,
+    >>>                                         lambda: 4,
+    >>>                                         lambda: 5])
+    >>> f(1)
+    4
+    >>> f(0)
+    5
+
+    Parameters
+    ----------
+    conditions: bool, sequence of bool, Array
+      The boolean conditions.
+    branches: Any
+      The branches, at least has two elements. Elements can be functions,
+      arrays, or numbers. The number of ``branches`` and ``conditions`` has
+      the relationship of `len(branches) == len(conditions) + 1`.
+      Each branch should receive one arguement for ``operands``.
+    *operands: optional, Any
+      The operands for each branch.
+    check_cond: bool
+      Whether to check the conditions. Default is True.
+
+    Returns
+    -------
+    res: Any
+      The results of the control flow.
+    """
+    # check branches
+    if not all(callable(branch) for branch in branches):
+        raise TypeError("branches argument should be a sequence of callables.")
+
+    # format branches
+    branches = tuple(branches)
+    if len(branches) == 0:
+        raise ValueError("Empty branch sequence")
+    elif len(branches) == 1:
+        return branches[0](*operands)
+    if len(conditions) != len(branches):
+        raise ValueError("The number of conditions should be equal to the number of branches.")
+
+    # format index
+    conditions = jnp.asarray(conditions, np.int32)
+    if check_cond:
+        jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {}.", err_arg=conditions)
+    index = jnp.where(conditions, size=1, fill_value=len(conditions) - 1)[0][0]
+    return switch(index, branches, *operands)

brainstate/compile/_conditions_test.py

@@ -0,0 +1,221 @@
+# 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 brainstate as bst
+
+
+class TestCond(unittest.TestCase):
+    def test1(self):
+        bst.random.seed(1)
+        bst.compile.cond(True, lambda: bst.random.random(10), lambda: bst.random.random(10))
+        bst.compile.cond(False, lambda: bst.random.random(10), lambda: bst.random.random(10))
+
+    def test2(self):
+        st1 = bst.State(bst.random.rand(10))
+        st2 = bst.State(bst.random.rand(2))
+        st3 = bst.State(bst.random.rand(5))
+        st4 = bst.State(bst.random.rand(2, 10))
+
+        def true_fun(x):
+            st1.value = st2.value @ st4.value + x
+
+        def false_fun(x):
+            st3.value = (st3.value + 1.) * x
+
+        bst.compile.cond(True, true_fun, false_fun, 2.)
+        assert not isinstance(st1.value, jax.core.Tracer)
+        assert not isinstance(st2.value, jax.core.Tracer)
+        assert not isinstance(st3.value, jax.core.Tracer)
+        assert not isinstance(st4.value, jax.core.Tracer)
+
+
+class TestSwitch(unittest.TestCase):
+    def testSwitch(self):
+        def branch(x):
+            y = jax.lax.mul(2, x)
+            return y, jax.lax.mul(2, y)
+
+        branches = [lambda x: (x, x),
+                    branch,
+                    lambda x: (x, -x)]
+
+        def fun(x):
+            if x <= 0:
+                return branches[0](x)
+            elif x == 1:
+                return branches[1](x)
+            else:
+                return branches[2](x)
+
+        def cfun(x):
+            return bst.compile.switch(x, branches, x)
+
+        self.assertEqual(fun(-1), cfun(-1))
+        self.assertEqual(fun(0), cfun(0))
+        self.assertEqual(fun(1), cfun(1))
+        self.assertEqual(fun(2), cfun(2))
+        self.assertEqual(fun(3), cfun(3))
+
+        cfun = jax.jit(cfun)
+
+        self.assertEqual(fun(-1), cfun(-1))
+        self.assertEqual(fun(0), cfun(0))
+        self.assertEqual(fun(1), cfun(1))
+        self.assertEqual(fun(2), cfun(2))
+        self.assertEqual(fun(3), cfun(3))
+
+    def testSwitchMultiOperands(self):
+        branches = [jax.lax.add, jax.lax.mul]
+
+        def fun(x):
+            i = 0 if x <= 0 else 1
+            return branches[i](x, x)
+
+        def cfun(x):
+            return bst.compile.switch(x, branches, x, x)
+
+        self.assertEqual(fun(-1), cfun(-1))
+        self.assertEqual(fun(0), cfun(0))
+        self.assertEqual(fun(1), cfun(1))
+        self.assertEqual(fun(2), cfun(2))
+        cfun = jax.jit(cfun)
+        self.assertEqual(fun(-1), cfun(-1))
+        self.assertEqual(fun(0), cfun(0))
+        self.assertEqual(fun(1), cfun(1))
+        self.assertEqual(fun(2), cfun(2))
+
+    def testSwitchResidualsMerge(self):
+        def get_conds(fun):
+            jaxpr = jax.make_jaxpr(jax.grad(fun))(0., 0)
+            return [eqn for eqn in jaxpr.jaxpr.eqns if eqn.primitive.name == 'cond']
+
+        def branch_invars_len(cond_eqn):
+            lens = [len(jaxpr.jaxpr.invars) for jaxpr in cond_eqn.params['branches']]
+            assert len(set(lens)) == 1
+            return lens[0]
+
+        def branch_outvars_len(cond_eqn):
+            lens = [len(jaxpr.jaxpr.outvars) for jaxpr in cond_eqn.params['branches']]
+            assert len(set(lens)) == 1
+            return lens[0]
+
+        branches1 = [lambda x: jnp.sin(x),
+                     lambda x: jnp.cos(x)]  # branch residuals overlap, should be reused
+        branches2 = branches1 + [lambda x: jnp.sinh(x)]  # another overlapping residual, expect reuse
+        branches3 = branches2 + [lambda x: jnp.sin(x) + jnp.cos(x)]  # requires one more residual slot
+
+        def fun1(x, i):
+            return bst.compile.switch(i + 1, branches1, x)
+
+        def fun2(x, i):
+            return bst.compile.switch(i + 1, branches2, x)
+
+        def fun3(x, i):
+            return bst.compile.switch(i + 1, branches3, x)
+
+        fwd1, bwd1 = get_conds(fun1)
+        fwd2, bwd2 = get_conds(fun2)
+        fwd3, bwd3 = get_conds(fun3)
+
+        fwd1_num_out = branch_outvars_len(fwd1)
+        fwd2_num_out = branch_outvars_len(fwd2)
+        fwd3_num_out = branch_outvars_len(fwd3)
+        assert fwd1_num_out == fwd2_num_out
+        assert fwd3_num_out == fwd2_num_out + 1
+
+        bwd1_num_in = branch_invars_len(bwd1)
+        bwd2_num_in = branch_invars_len(bwd2)
+        bwd3_num_in = branch_invars_len(bwd3)
+        assert bwd1_num_in == bwd2_num_in
+        assert bwd3_num_in == bwd2_num_in + 1
+
+    def testOneBranchSwitch(self):
+        branch = lambda x: -x
+        f = lambda i, x: bst.compile.switch(i, [branch], x)
+        x = 7.
+        self.assertEqual(f(-1, x), branch(x))
+        self.assertEqual(f(0, x), branch(x))
+        self.assertEqual(f(1, x), branch(x))
+        cf = jax.jit(f)
+        self.assertEqual(cf(-1, x), branch(x))
+        self.assertEqual(cf(0, x), branch(x))
+        self.assertEqual(cf(1, x), branch(x))
+        cf = jax.jit(f, static_argnums=0)
+        self.assertEqual(cf(-1, x), branch(x))
+        self.assertEqual(cf(0, x), branch(x))
+        self.assertEqual(cf(1, x), branch(x))
+
+
+class TestIfElse(unittest.TestCase):
+    def test1(self):
+        def f(a):
+            return bst.compile.ifelse(conditions=[a < 0,
+                                                  a >= 0 and a < 2,
+                                                  a >= 2 and a < 5,
+                                                  a >= 5 and a < 10,
+                                                  a >= 10],
+                                      branches=[lambda: 1,
+                                                lambda: 2,
+                                                lambda: 3,
+                                                lambda: 4,
+                                                lambda: 5])
+
+        self.assertTrue(f(3) == 3)
+        self.assertTrue(f(1) == 2)
+        self.assertTrue(f(-1) == 1)
+
+    def test_vmap(self):
+        def f(operands):
+            f = lambda a: bst.compile.ifelse([a > 10,
+                                              jnp.logical_and(a <= 10, a > 5),
+                                              jnp.logical_and(a <= 5, a > 2),
+                                              jnp.logical_and(a <= 2, a > 0),
+                                              a <= 0],
+                                             [lambda _: 1,
+                                              lambda _: 2,
+                                              lambda _: 3,
+                                              lambda _: 4,
+                                              lambda _: 5, ],
+                                             a)
+            return jax.vmap(f)(operands)
+
+        r = f(bst.random.randint(-20, 20, 200))
+        self.assertTrue(r.size == 200)
+
+    def test_grad1(self):
+        def F2(x):
+            return bst.compile.ifelse((x >= 10, x < 10),
+                                      [lambda x: x, lambda x: x ** 2, ],
+                                      x)
+
+        self.assertTrue(jax.grad(F2)(9.0) == 18.)
+        self.assertTrue(jax.grad(F2)(11.0) == 1.)
+
+    def test_grad2(self):
+        def F3(x):
+            return bst.compile.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
+                                      [lambda x: x,
+                                       lambda x: x ** 2,
+                                       lambda x: x ** 4, ],
+                                      x)
+
+        self.assertTrue(jax.grad(F3)(9.0) == 18.)
+        self.assertTrue(jax.grad(F3)(11.0) == 1.)

brainstate/compile/_error_if.py

@@ -0,0 +1,94 @@
+# 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 functools
+from functools import partial
+from typing import Callable, Union
+
+import jax
+
+from brainstate._utils import set_module_as
+from ._unvmap import unvmap
+
+__all__ = [
+    'jit_error_if',
+]
+
+
+def _err_jit_true_branch(err_fun, args, kwargs):
+    jax.debug.callback(err_fun, *args, **kwargs)
+
+
+def _err_jit_false_branch(args, kwargs):
+    pass
+
+
+def _error_msg(msg, *arg, **kwargs):
+    if len(arg):
+        msg = msg % arg
+    if len(kwargs):
+        msg = msg.format(**kwargs)
+    raise ValueError(msg)
+
+
+@set_module_as('brainstate.compile')
+def jit_error_if(
+    pred,
+    error: Union[Callable, str],
+    *err_args,
+    **err_kwargs,
+):
+    """
+    Check errors in a jit function.
+
+    Examples
+    --------
+
+    It can give a function which receive arguments that passed from the JIT variables and raise errors.
+
+    >>> def error(x):
+    >>>    raise ValueError(f'error {x}')
+    >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
+    >>> jit_error_if(x.sum() < 5., error, x)
+
+    Or, it can be a simple string message.
+
+    >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
+    >>> jit_error_if(x.sum() < 5., "Error: the sum is less than 5. Got {s}", s=x.sum())
+
+
+    Parameters
+    ----------
+    pred: bool, Array
+      The boolean prediction.
+    error: callable, str
+      The error function, which raise errors, or a string indicating the error message.
+    err_args:
+      The arguments which passed into `err_f`.
+    err_kwargs:
+      The keywords which passed into `err_f`.
+    """
+    if isinstance(error, str):
+        error = partial(_error_msg, error)
+
+    jax.lax.cond(
+        unvmap(pred, op='any'),
+        partial(_err_jit_true_branch, error),
+        _err_jit_false_branch,
+        jax.tree.map(functools.partial(unvmap, op='none'), err_args),
+        jax.tree.map(functools.partial(unvmap, op='none'), err_kwargs),
+    )

brainstate/compile/_error_if_test.py

@@ -0,0 +1,54 @@
+# 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 jaxlib.xla_extension
+
+import brainstate as bst
+
+
+class TestJitError(unittest.TestCase):
+    def test1(self):
+        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            bst.compile.jit_error_if(True, 'error')
+
+        def err_f(x):
+            raise ValueError(f'error: {x}')
+
+        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            bst.compile.jit_error_if(True, err_f, 1.)
+
+    def test_vmap(self):
+        def f(x):
+            bst.compile.jit_error_if(x, 'error: {x}', x=x)
+
+        jax.vmap(f)(jnp.array([False, False, False]))
+        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            jax.vmap(f)(jnp.array([True, False, False]))
+
+    def test_vmap_vmap(self):
+        def f(x):
+            bst.compile.jit_error_if(x, 'error: {x}', x=x)
+
+        jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
+                                         [False, False, False]]))
+        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
+                                             [True, False, False]]))