brainstate 0.0.2.post20240824__py2.py3-none-any.whl → 0.0.2.post20240826__py2.py3-none-any.whl

brainstate/transform/{_controls_test.py → _conditions_test.py}

@@ -18,20 +18,20 @@ import unittest
 import jax
 import jax.numpy as jnp
 
-import brainstate as bc
+import brainstate as bst
 
 
 class TestCond(unittest.TestCase):
   def test1(self):
-    bc.random.seed(1)
-    bc.transform.cond(True, lambda: bc.random.random(10), lambda: bc.random.random(10))
-    bc.transform.cond(False, lambda: bc.random.random(10), lambda: bc.random.random(10))
+    bst.random.seed(1)
+    bst.transform.cond(True, lambda: bst.random.random(10), lambda: bst.random.random(10))
+    bst.transform.cond(False, lambda: bst.random.random(10), lambda: bst.random.random(10))
 
   def test2(self):
-    st1 = bc.State(bc.random.rand(10))
-    st2 = bc.State(bc.random.rand(2))
-    st3 = bc.State(bc.random.rand(5))
-    st4 = bc.State(bc.random.rand(2, 10))
+    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
@@ -39,7 +39,7 @@ class TestCond(unittest.TestCase):
     def false_fun(x):
       st3.value = (st3.value + 1.) * x
 
-    bc.transform.cond(True, true_fun, false_fun, 2.)
+    bst.transform.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)
@@ -65,7 +65,7 @@ class TestSwitch(unittest.TestCase):
         return branches[2](x)
 
     def cfun(x):
-      return bc.transform.switch(x, branches, x)
+      return bst.transform.switch(x, branches, x)
 
     self.assertEqual(fun(-1), cfun(-1))
     self.assertEqual(fun(0), cfun(0))
@@ -89,7 +89,7 @@ class TestSwitch(unittest.TestCase):
       return branches[i](x, x)
 
     def cfun(x):
-      return bc.transform.switch(x, branches, x, x)
+      return bst.transform.switch(x, branches, x, x)
 
     self.assertEqual(fun(-1), cfun(-1))
     self.assertEqual(fun(0), cfun(0))
@@ -122,13 +122,13 @@ class TestSwitch(unittest.TestCase):
     branches3 = branches2 + [lambda x: jnp.sin(x) + jnp.cos(x)]  # requires one more residual slot
 
     def fun1(x, i):
-      return bc.transform.switch(i + 1, branches1, x)
+      return bst.transform.switch(i + 1, branches1, x)
 
     def fun2(x, i):
-      return bc.transform.switch(i + 1, branches2, x)
+      return bst.transform.switch(i + 1, branches2, x)
 
     def fun3(x, i):
-      return bc.transform.switch(i + 1, branches3, x)
+      return bst.transform.switch(i + 1, branches3, x)
 
     fwd1, bwd1 = get_conds(fun1)
     fwd2, bwd2 = get_conds(fun2)
@@ -148,7 +148,7 @@ class TestSwitch(unittest.TestCase):
 
   def testOneBranchSwitch(self):
     branch = lambda x: -x
-    f = lambda i, x: bc.transform.switch(i, [branch], x)
+    f = lambda i, x: bst.transform.switch(i, [branch], x)
     x = 7.
     self.assertEqual(f(-1, x), branch(x))
     self.assertEqual(f(0, x), branch(x))
@@ -166,12 +166,12 @@ class TestSwitch(unittest.TestCase):
 class TestIfElse(unittest.TestCase):
   def test1(self):
     def f(a):
-      return bc.transform.ifelse(conditions=[a < 0,
-                                             a >= 0 and a < 2,
-                                             a >= 2 and a < 5,
-                                             a >= 5 and a < 10,
-                                             a >= 10],
-                                 branches=[lambda: 1,
+      return bst.transform.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,
@@ -183,38 +183,38 @@ class TestIfElse(unittest.TestCase):
 
   def test_vmap(self):
     def f(operands):
-      f = lambda a: bc.transform.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,
+      f = lambda a: bst.transform.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)
+                                         a)
       return jax.vmap(f)(operands)
 
-    r = f(bc.random.randint(-20, 20, 200))
+    r = f(bst.random.randint(-20, 20, 200))
     self.assertTrue(r.size == 200)
 
   def test_grad1(self):
     def F2(x):
-      return bc.transform.ifelse((x >= 10, x < 10),
-                                 [lambda x: x, lambda x: x ** 2, ],
-                                 x)
+      return bst.transform.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 bc.transform.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
-                                 [lambda x: x,
+      return bst.transform.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
+                                  [lambda x: x,
                                   lambda x: x ** 2,
                                   lambda x: x ** 4, ],
-                                 x)
+                                  x)
 
     self.assertTrue(jax.grad(F3)(9.0) == 18.)
     self.assertTrue(jax.grad(F3)(11.0) == 1.)

brainstate/transform/_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.transform')
+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/transform/{_jit_error_test.py → _error_if_test.py}

@@ -25,18 +25,18 @@ import brainstate as bst
 class TestJitError(unittest.TestCase):
   def test1(self):
     with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-      bst.transform.jit_error(True, 'error')
+      bst.transform.jit_error_if(True, 'error')
 
     def err_f(x):
       raise ValueError(f'error: {x}')
 
     with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-      bst.transform.jit_error(True, err_f, 1.)
+      bst.transform.jit_error_if(True, err_f, 1.)
 
   def test_vmap(self):
 
     def f(x):
-      bst.transform.jit_error(x, 'error: {x}', x=x)
+      bst.transform.jit_error_if(x, 'error: {x}', x=x)
 
     jax.vmap(f)(jnp.array([False, False, False]))
     with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
@@ -45,7 +45,7 @@ class TestJitError(unittest.TestCase):
   def test_vmap_vmap(self):
 
     def f(x):
-      bst.transform.jit_error(x, 'error: {x}', x=x)
+      bst.transform.jit_error_if(x, 'error: {x}', x=x)
 
     jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
                                      [False, False, False]]))