brainstate 0.0.1__py2.py3-none-any.whl

brainstate/transform/_controls_test.py

@@ -0,0 +1,220 @@
+# 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 brainstate as bc
+
+
+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))
+
+  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))
+
+    def true_fun(x):
+      st1.value = st2.value @ st4.value + x
+
+    def false_fun(x):
+      st3.value = (st3.value + 1.) * x
+
+    bc.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)
+    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 bc.transform.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 bc.transform.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 bc.transform.switch(i + 1, branches1, x)
+
+    def fun2(x, i):
+      return bc.transform.switch(i + 1, branches2, x)
+
+    def fun3(x, i):
+      return bc.transform.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: bc.transform.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 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,
+                                           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: 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,
+                                         lambda _: 2,
+                                         lambda _: 3,
+                                         lambda _: 4,
+                                         lambda _: 5, ],
+                                        a)
+      return jax.vmap(f)(operands)
+
+    r = f(bc.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)
+
+    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,
+                                  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/transform/_jit.py

@@ -0,0 +1,239 @@
+# 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 collections.abc import Iterable, Sequence
+from typing import (Any, Callable, Union)
+
+import jax
+from jax._src import sharding_impls
+from jax.lib import xla_client as xc
+
+from ._make_jaxpr import StatefulFunction, _ensure_index_tuple, _assign_state_values
+from brainstate._utils import set_module_as
+
+__all__ = ['jit']
+
+
+class JittedFunction(Callable):
+  """
+  A wrapped version of ``fun``, set up for just-in-time compilation.
+  """
+  origin_fun: Callable   # the original function
+  stateful_fun: StatefulFunction  # the stateful function for extracting states
+  jitted_fun: jax.stages.Wrapped  # the jitted function
+  clear_cache: Callable # clear the cache of the jitted function
+
+
+def _get_jitted_fun(
+    fun: Callable,
+    in_shardings,
+    out_shardings,
+    static_argnums,
+    donate_argnums,
+    donate_argnames,
+    keep_unused,
+    device,
+    backend,
+    inline,
+    abstracted_axes,
+    **kwargs
+) -> JittedFunction:
+  static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+  # TODO: add to cache stack for clear_cache
+  fun = StatefulFunction(fun, static_argnums=static_argnums, abstracted_axes=abstracted_axes, cache_type='jit')
+  jit_fun = jax.jit(fun.jaxpr_call,
+                    static_argnums=tuple(i + 1 for i in static_argnums),
+                    donate_argnums=donate_argnums,
+                    donate_argnames=donate_argnames,
+                    keep_unused=keep_unused,
+                    device=device,
+                    backend=backend,
+                    inline=inline,
+                    in_shardings=in_shardings,
+                    out_shardings=out_shardings,
+                    abstracted_axes=abstracted_axes,
+                    **kwargs)
+
+  @functools.wraps(fun.fun)
+  def jitted_fun(*args, **params):
+    if jax.config.jax_disable_jit:
+      return fun.fun(*args, **params)
+    states = fun.compile_and_get_states_by_static_args(*args, **kwargs)
+    state_vals, outs = jit_fun([st.value for st in states], *args, **params)
+    _assign_state_values(states, state_vals)
+    return outs
+
+  def clear_cache():
+    # clear the cache of the stateful function
+    fun.clear_cache()
+    # clear the cache of the jitted function
+    jit_fun.clear_cache()
+
+  jitted_fun: JittedFunction
+  # the original function
+  jitted_fun.origin_fun = fun.fun
+  # the stateful function for extracting states
+  jitted_fun.stateful_fun = fun
+  # the jitted function
+  jitted_fun.jitted_fun = jit_fun
+  # clear cache
+  jitted_fun.clear_cache = clear_cache
+
+  return jitted_fun
+
+
+@set_module_as('brainstate.transform')
+def jit(
+  fun: Callable = None,
+  in_shardings=sharding_impls.UNSPECIFIED,
+  out_shardings=sharding_impls.UNSPECIFIED,
+  static_argnums: int | Sequence[int] | None = None,
+  donate_argnums: int | Sequence[int] | None = None,
+  donate_argnames: str | Iterable[str] | None = None,
+  keep_unused: bool = False,
+  device: xc.Device | None = None,
+  backend: str | None = None,
+  inline: bool = False,
+  abstracted_axes: Any | None = None,
+  **kwargs
+) -> Union[JittedFunction, Callable[[Callable], JittedFunction]]:
+  """
+  Sets up ``fun`` for just-in-time compilation with XLA.
+
+  Does not support setting ``static_argnames`` as in ``jax.jit()``.
+
+
+  Args:
+    fun: Function to be jitted.
+    in_shardings: Pytree of structure matching that of arguments to ``fun``,
+      with all actual arguments replaced by resource assignment specifications.
+      It is also valid to specify a pytree prefix (e.g. one value in place of a
+      whole subtree), in which case the leaves get broadcast to all values in
+      that subtree.
+
+      The ``in_shardings`` argument is optional. JAX will infer the shardings
+      from the input :py:class:`jax.Array`'s and defaults to replicating the input
+      if the sharding cannot be inferred.
+
+      The valid resource assignment specifications are:
+        - :py:class:`XLACompatibleSharding`, which will decide how the value
+            will be partitioned. With this, using a mesh context manager is not
+            required.
+        - :py:obj:`None`, will give JAX the freedom to choose whatever sharding
+          it wants.
+          For in_shardings, JAX will mark is as replicated but this behavior
+          can change in the future.
+          For out_shardings, we will rely on the XLA GSPMD partitioner to
+          determine the output shardings.
+
+      The size of every dimension has to be a multiple of the total number of
+      resources assigned to it. This is similar to pjit's in_shardings.
+    out_shardings: Like ``in_shardings``, but specifies resource
+      assignment for function outputs. This is similar to pjit's
+      out_shardings.
+
+      The ``out_shardings`` argument is optional. If not specified, :py:func:`jax.jit`
+      will use GSPMD's sharding propagation to figure out what the sharding of the
+      output(s) should be.
+    static_argnums: An optional int or collection of ints that specify which
+      positional arguments to treat as static (compile-time constant).
+      Operations that only depend on static arguments will be constant-folded in
+      Python (during tracing), and so the corresponding argument values can be
+      any Python object.
+
+      Static arguments should be hashable, meaning both ``__hash__`` and
+      ``__eq__`` are implemented, and immutable. Calling the jitted function
+      with different values for these constants will trigger recompilation.
+      Arguments that are not arrays or containers thereof must be marked as
+      static.
+
+      If neither ``static_argnums`` nor ``static_argnames`` is provided, no
+      arguments are treated as static. If ``static_argnums`` is not provided but
+      ``static_argnames`` is, or vice versa, JAX uses
+      :code:`inspect.signature(fun)` to find any positional arguments that
+      correspond to ``static_argnames``
+      (or vice versa). If both ``static_argnums`` and ``static_argnames`` are
+      provided, ``inspect.signature`` is not used, and only actual
+      parameters listed in either ``static_argnums`` or ``static_argnames`` will
+      be treated as static.
+    donate_argnums: Specify which positional argument buffers are "donated" to
+      the computation. It is safe to donate argument buffers if you no longer
+      need them once the computation has finished. In some cases XLA can make
+      use of donated buffers to reduce the amount of memory needed to perform a
+      computation, for example recycling one of your input buffers to store a
+      result. You should not reuse buffers that you donate to a computation, JAX
+      will raise an error if you try to. By default, no argument buffers are
+      donated.
+
+      If neither ``donate_argnums`` nor ``donate_argnames`` is provided, no
+      arguments are donated. If ``donate_argnums`` is not provided but
+      ``donate_argnames`` is, or vice versa, JAX uses
+      :code:`inspect.signature(fun)` to find any positional arguments that
+      correspond to ``donate_argnames``
+      (or vice versa). If both ``donate_argnums`` and ``donate_argnames`` are
+      provided, ``inspect.signature`` is not used, and only actual
+      parameters listed in either ``donate_argnums`` or ``donate_argnames`` will
+      be donated.
+
+      For more details on buffer donation see the
+      `FAQ <https://jax.readthedocs.io/en/latest/faq.html#buffer-donation>`_.
+    donate_argnames: An optional string or collection of strings specifying
+      which named arguments are donated to the computation. See the
+      comment on ``donate_argnums`` for details. If not
+      provided but ``donate_argnums`` is set, the default is based on calling
+      ``inspect.signature(fun)`` to find corresponding named arguments.
+    keep_unused: If `False` (the default), arguments that JAX determines to be
+      unused by `fun` *may* be dropped from resulting compiled XLA executables.
+      Such arguments will not be transferred to the device nor provided to the
+      underlying executable. If `True`, unused arguments will not be pruned.
+    device: This is an experimental feature and the API is likely to change.
+      Optional, the Device the jitted function will run on. (Available devices
+      can be retrieved via :py:func:`jax.devices`.) The default is inherited
+      from XLA's DeviceAssignment logic and is usually to use
+      ``jax.devices()[0]``.
+    backend: This is an experimental feature and the API is likely to change.
+      Optional, a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or
+      ``'tpu'``.
+    inline: Specify whether this function should be inlined into enclosing
+      jaxprs (rather than being represented as an application of the xla_call
+      primitive with its own subjaxpr). Default False.
+    abstracted_axes:
+
+  Returns:
+    A wrapped version of ``fun``, set up for just-in-time compilation.
+    The returned object is a :py:class:`JittedFunction` that can be called with the same arguments
+    and has the following attributes and methods:
+
+    - ``stateful_fun``: the stateful function for extracting states, an instance of :py:class:`StatefulFunction`.
+    - ``origin_fun(*args, **kwargs)``: the original function
+    - ``jitted_fun(*args, **kwargs)``: the jitted function
+    - ``clear_cache(*args, **kwargs)``: clear the cache of the jitted function
+
+  """
+
+  if fun is None:
+    def wrapper(fun_again: Callable) -> JittedFunction:
+      return _get_jitted_fun(fun_again, in_shardings, out_shardings, static_argnums,
+                             donate_argnums, donate_argnames, keep_unused,
+                             device, backend, inline, abstracted_axes, **kwargs)
+    return wrapper
+
+  else:
+    return _get_jitted_fun(fun, in_shardings, out_shardings, static_argnums,
+                           donate_argnums, donate_argnames, keep_unused,
+                           device, backend, inline, abstracted_axes, **kwargs)

brainstate/transform/_jit_error.py

@@ -0,0 +1,158 @@
+# 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, partial
+from typing import Callable, Union
+
+import jax
+from jax import numpy as jnp
+from jax.core import Primitive, ShapedArray
+from jax.interpreters import batching, mlir, xla
+from jax.lax import cond
+
+from brainstate._utils import set_module_as
+
+__all__ = [
+  'jit_error',
+]
+
+
+@set_module_as('brainstate.transform')
+def remove_vmap(x, op='any'):
+  if op == 'any':
+    return _any_without_vmap(x)
+  elif op == 'all':
+    return _all_without_vmap(x)
+  else:
+    raise ValueError(f'Do not support type: {op}')
+
+
+_any_no_vmap_prim = Primitive('any_no_vmap')
+
+
+def _any_without_vmap(x):
+  return _any_no_vmap_prim.bind(x)
+
+
+def _any_without_vmap_imp(x):
+  return jnp.any(x)
+
+
+def _any_without_vmap_abs(x):
+  return ShapedArray(shape=(), dtype=jnp.bool_)
+
+
+def _any_without_vmap_batch(x, batch_axes):
+  (x,) = x
+  return _any_without_vmap(x), batching.not_mapped
+
+
+_any_no_vmap_prim.def_impl(_any_without_vmap_imp)
+_any_no_vmap_prim.def_abstract_eval(_any_without_vmap_abs)
+batching.primitive_batchers[_any_no_vmap_prim] = _any_without_vmap_batch
+if hasattr(xla, "lower_fun"):
+  xla.register_translation(_any_no_vmap_prim,
+                           xla.lower_fun(_any_without_vmap_imp, multiple_results=False, new_style=True))
+mlir.register_lowering(_any_no_vmap_prim, mlir.lower_fun(_any_without_vmap_imp, multiple_results=False))
+
+_all_no_vmap_prim = Primitive('all_no_vmap')
+
+
+def _all_without_vmap(x):
+  return _all_no_vmap_prim.bind(x)
+
+
+def _all_without_vmap_imp(x):
+  return jnp.all(x)
+
+
+def _all_without_vmap_abs(x):
+  return ShapedArray(shape=(), dtype=jnp.bool_)
+
+
+def _all_without_vmap_batch(x, batch_axes):
+  (x,) = x
+  return _all_without_vmap(x), batching.not_mapped
+
+
+_all_no_vmap_prim.def_impl(_all_without_vmap_imp)
+_all_no_vmap_prim.def_abstract_eval(_all_without_vmap_abs)
+batching.primitive_batchers[_all_no_vmap_prim] = _all_without_vmap_batch
+if hasattr(xla, "lower_fun"):
+  xla.register_translation(_all_no_vmap_prim,
+                           xla.lower_fun(_all_without_vmap_imp, multiple_results=False, new_style=True))
+mlir.register_lowering(_all_no_vmap_prim, mlir.lower_fun(_all_without_vmap_imp, multiple_results=False))
+
+
+def _err_jit_true_branch(err_fun, x):
+  jax.debug.callback(err_fun, x)
+  return
+
+
+def _err_jit_false_branch(x):
+  return
+
+
+def _cond(err_fun, pred, err_arg):
+  @wraps(err_fun)
+  def true_err_fun(*arg):
+    err_fun(*arg)
+
+  cond(pred,
+       partial(_err_jit_true_branch, true_err_fun),
+       _err_jit_false_branch,
+       err_arg)
+
+
+def _error_msg(msg, *arg):
+  if len(arg) == 0:
+    raise ValueError(msg)
+  else:
+    raise ValueError(msg.format(arg))
+
+
+@set_module_as('brainstate.transform')
+def jit_error(pred, err_fun: Union[Callable, str], err_arg=None, scope: str = 'any'):
+  """Check errors in a jit function.
+
+  >>> def error(arg):
+  >>>    raise ValueError(f'error {arg}')
+  >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
+  >>> jit_error(x.sum() < 5., error, err_arg=x)
+
+  Parameters
+  ----------
+  pred: bool, Array
+    The boolean prediction.
+  err_fun: callable
+    The error function, which raise errors.
+  err_arg: any
+    The arguments which passed into `err_f`.
+  scope: str
+    The scope of the error message. Can be None, 'all' or 'any'.
+  """
+  if isinstance(err_fun, str):
+    err_fun = partial(_error_msg, err_fun)
+  if scope is None:
+    pred = pred
+  elif scope == 'all':
+    pred = remove_vmap(pred, 'all')
+  elif scope == 'any':
+    pred = remove_vmap(pred, 'any')
+  else:
+    raise ValueError(f"Unknown scope: {scope}")
+  _cond(err_fun, pred, err_arg)