brainstate 0.1.0__py2.py3-none-any.whl → 0.1.0.post20241125__py2.py3-none-any.whl

brainstate/event/_xla_custom_op.py

@@ -0,0 +1,309 @@
+# -*- coding: utf-8 -*-
+
+import ctypes
+import functools
+import importlib.util
+from functools import partial
+from typing import Callable, Sequence, Tuple, Protocol
+
+import jax
+import numpy as np
+from jax import tree_util
+from jax.core import Primitive
+from jax.interpreters import batching, ad
+from jax.interpreters import xla, mlir
+from jax.lib import xla_client
+from jaxlib.hlo_helpers import custom_call
+
+numba_installed = importlib.util.find_spec('numba') is not None
+
+__all__ = [
+    'defjvp',
+    'XLACustomOp',
+]
+
+#                                         [void* pointer,
+#                                          const char *name,
+#                                          PyCapsule_Destructor destructor]
+ctypes.pythonapi.PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
+ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object
+
+
+def defjvp(primitive, *jvp_rules):
+    """Define JVP rules for any JAX primitive.
+
+    This function is similar to ``jax.interpreters.ad.defjvp``.
+    However, the JAX one only supports primitive with ``multiple_results=False``.
+    ``brainpy.math.defjvp`` enables to define the independent JVP rule for
+    each input parameter no matter ``multiple_results=False/True``.
+
+    For examples, please see ``test_ad_support.py``.
+
+    Args:
+      primitive: Primitive, XLACustomOp.
+      *jvp_rules: The JVP translation rule for each primal.
+    """
+    assert isinstance(primitive, Primitive)
+    if primitive.multiple_results:
+        ad.primitive_jvps[primitive] = partial(_standard_jvp, jvp_rules, primitive)
+    else:
+        ad.primitive_jvps[primitive] = partial(ad.standard_jvp, jvp_rules, primitive)
+
+
+def _standard_jvp(jvp_rules, primitive: Primitive, primals, tangents, **params):
+    assert primitive.multiple_results
+    val_out = tuple(primitive.bind(*primals, **params))
+    tree = tree_util.tree_structure(val_out)
+    tangents_out = []
+    for rule, t in zip(jvp_rules, tangents):
+        if rule is not None and type(t) is not ad.Zero:
+            r = tuple(rule(t, *primals, **params))
+            tangents_out.append(r)
+            assert tree_util.tree_structure(r) == tree
+    r = functools.reduce(
+        _add_tangents,
+        tangents_out,
+        tree_util.tree_map(
+            # compatible with JAX 0.4.34
+            lambda a: ad.Zero.from_primal_value(a) if jax.__version__ >= '0.4.34' else ad.Zero.from_value(a),
+            val_out
+        )
+    )
+    return val_out, r
+
+
+def _add_tangents(xs, ys):
+    return tree_util.tree_map(ad.add_tangents, xs, ys, is_leaf=lambda a: isinstance(a, ad.Zero))
+
+
+def _shape_to_layout(shape):
+    return tuple(range(len(shape) - 1, -1, -1))
+
+
+def _numba_mlir_cpu_translation_rule(
+    kernel,
+    debug: bool,
+    ctx,
+    *ins,
+    **kwargs
+):
+    if not numba_installed:
+        raise ImportError('Numba is required to compile the CPU kernel for the custom operator.')
+
+    from numba import types, carray, cfunc  # pylint: disable=import-error
+    from numba.core.dispatcher import Dispatcher  # pylint: disable=import-error
+
+    if not isinstance(kernel, Dispatcher):
+        kernel = kernel(**kwargs)
+    assert isinstance(kernel, Dispatcher), f'The kernel should be a Numba dispatcher. But we got {kernel}'
+
+    # output information
+    outs = ctx.avals_out
+    output_shapes = tuple([out.shape for out in outs])
+    output_dtypes = tuple([out.dtype for out in outs])
+    output_layouts = tuple([_shape_to_layout(out.shape) for out in outs])
+    result_types = [mlir.aval_to_ir_type(out) for out in outs]
+
+    # input information
+    avals_in = ctx.avals_in
+    input_layouts = [_shape_to_layout(a.shape) for a in avals_in]
+    input_dtypes = tuple(inp.dtype for inp in avals_in)
+    input_shapes = tuple(inp.shape for inp in avals_in)
+
+    # compiling function
+    code_scope = dict(func_to_call=kernel,
+                      input_shapes=input_shapes,
+                      input_dtypes=input_dtypes,
+                      output_shapes=output_shapes,
+                      output_dtypes=output_dtypes, carray=carray)
+    args_in = [f'in{i} = carray(input_ptrs[{i}], input_shapes[{i}], dtype=input_dtypes[{i}])'
+               for i in range(len(input_shapes))]
+    if len(output_shapes) > 1:
+        args_out = [f'out{i} = carray(output_ptrs[{i}], output_shapes[{i}], dtype=output_dtypes[{i}])'
+                    for i in range(len(output_shapes))]
+        sig = types.void(types.CPointer(types.voidptr), types.CPointer(types.voidptr))
+    else:
+        args_out = [f'out0 = carray(output_ptrs, output_shapes[0], dtype=output_dtypes[0])']
+        sig = types.void(types.voidptr, types.CPointer(types.voidptr))
+    args_call = [f'in{i}' for i in range(len(input_shapes))] + [f'out{i}' for i in range(len(output_shapes))]
+    code_string = '''
+def numba_cpu_custom_call_target(output_ptrs, input_ptrs):
+    {args_in}
+    {args_out}
+    func_to_call({args_call})
+      '''.format(args_in="\n    ".join(args_in),
+                 args_out="\n    ".join(args_out),
+                 args_call=", ".join(args_call))
+    if debug:
+        print(code_string)
+    exec(compile(code_string.strip(), '', 'exec'), code_scope)
+    new_f = code_scope['numba_cpu_custom_call_target']
+
+    # register
+    xla_c_rule = cfunc(sig)(new_f)
+    target_name = f'numba_custom_call_{str(xla_c_rule.address)}'
+    capsule = ctypes.pythonapi.PyCapsule_New(xla_c_rule.address, b"xla._CUSTOM_CALL_TARGET", None)
+    xla_client.register_custom_call_target(target_name, capsule, "cpu")
+
+    # call
+    return custom_call(
+        call_target_name=target_name,
+        operands=ins,
+        operand_layouts=list(input_layouts),
+        result_layouts=list(output_layouts),
+        result_types=list(result_types),
+        has_side_effect=False,
+    ).results
+
+
+def register_numba_mlir_cpu_translation_rule(
+    primitive: jax.core.Primitive,
+    cpu_kernel: Callable,
+    debug: bool = False
+):
+    rule = partial(_numba_mlir_cpu_translation_rule, cpu_kernel, debug)
+    mlir.register_lowering(primitive, rule, platform='cpu')
+
+
+class ShapeDtype(Protocol):
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        ...
+
+    @property
+    def dtype(self) -> np.dtype:
+        ...
+
+
+class XLACustomOp:
+    """Creating a XLA custom call operator.
+
+    Args:
+      cpu_kernel_generator: Callable. The function defines the computation on CPU backend.
+      gpu_kernel_generator: Callable. The function defines the computation on GPU backend.
+      batching_translation: Callable. The batching translation rule of JAX.
+      jvp_translation: Callable. The JVP translation rule of JAX.
+      transpose_translation: Callable. The transpose translation rule of JAX.
+      name: str. The primitive name.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        cpu_kernel_generator: Callable,
+        gpu_kernel_generator: Callable = None,
+        batching_translation: Callable = None,
+        jvp_translation: Callable = None,
+        transpose_translation: Callable = None,
+    ):
+        # set cpu_kernel and gpu_kernel
+        self.cpu_kernel = cpu_kernel_generator
+
+        # primitive
+        self.primitive = jax.core.Primitive(name)
+        self.primitive.multiple_results = True
+
+        # abstract evaluation
+        self.primitive.def_impl(partial(xla.apply_primitive, self.primitive))
+        self.primitive.def_abstract_eval(self._abstract_eval)
+
+        # cpu kernel
+        if cpu_kernel_generator is not None:
+            self.def_cpu_kernel(cpu_kernel_generator)
+        if gpu_kernel_generator is not None:
+            self.def_gpu_kernel(gpu_kernel_generator)
+
+        # batching rule
+        if batching_translation is not None:
+            batching.primitive_batchers[self.primitive] = batching_translation
+
+        # jvp rule
+        if jvp_translation is not None:
+            ad.primitive_jvps[self.primitive] = jvp_translation
+
+        # transpose rule
+        if transpose_translation is not None:
+            ad.primitive_transposes[self.primitive] = transpose_translation
+
+    def _abstract_eval(self, *ins, outs: Sequence[ShapeDtype], **kwargs):
+        return tuple(outs)
+
+    def __call__(self, *ins, outs: Sequence[ShapeDtype], **kwargs):
+        assert isinstance(outs, (tuple, list)), 'The `outs` should be a tuple or list of shape-dtype pairs.'
+        outs = jax.tree.map(_transform_to_shapedarray, outs)
+        return self.primitive.bind(*ins, **kwargs, outs=tuple(outs))
+
+    def def_cpu_kernel(self, kernel_generator: Callable):
+        """
+        Define the CPU kernel using Numba.
+        """
+        register_numba_mlir_cpu_translation_rule(self.primitive, kernel_generator)
+
+    def def_gpu_kernel(self, kernel_generator: Callable):
+        """
+        Define the GPU kernel using the JAX Pallas language.
+        """
+        lower = mlir.lower_fun(
+            lambda *args, **kwargs: kernel_generator(**kwargs)(*args),
+            multiple_results=True
+        )
+        mlir.register_lowering(self.primitive, lower, platform='cuda')
+        mlir.register_lowering(self.primitive, lower, platform='tpu')
+
+    def def_batching_rule(self, fun):
+        """Define the batching rule.
+
+        Args:
+          fun: The batching rule.
+        """
+        batching.primitive_batchers[self.primitive] = fun
+
+    def def_jvp_rule(self, fun):
+        """Define the JVP rule.
+
+        Args:
+          fun: The JVP rule.
+        """
+        ad.primitive_jvps[self.primitive] = fun
+
+    def defjvp(self, *jvp_rules):
+        """
+        Define the JVP rule. Similar to ``jax.interpreters.ad.defjvp``,
+        but supports the Primitive with multiple results.
+
+        Args:
+          jvp_rules: The JVP rules.
+        """
+        defjvp(self.primitive, *jvp_rules)
+
+    def def_transpose_rule(self, fun):
+        """Define the transpose rule.
+
+        Args:
+          fun: The transpose rule.
+        """
+        ad.primitive_transposes[self.primitive] = fun
+
+    def def_xla_translation(self, platform, fun):
+        """Define the XLA translation rule.
+
+        Args:
+          platform: str. The computing platform.
+          fun: The XLA translation rule.
+        """
+        xla.backend_specific_translations[platform][self.primitive] = fun
+
+    def def_mlir_lowering(self, platform, fun):
+        """
+        Define the MLIR lowering rule.
+
+        Args:
+          platform: str. The computing platform.
+          fun: The lowering rule.
+        """
+        mlir.register_lowering(self.primitive, fun, platform)
+
+
+def _transform_to_shapedarray(a):
+    return jax.core.ShapedArray(a.shape, a.dtype)

brainstate/event/_xla_custom_op_test.py

@@ -0,0 +1,55 @@
+# 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.numpy as jnp
+from jax.experimental import pallas as pl
+
+import brainstate as bst
+
+
+def test1():
+    import numba
+    def add_vectors_kernel(x_ref, y_ref, o_ref):
+        x, y = x_ref[...], y_ref[...]
+        o_ref[...] = x + y
+
+    def cpu_kernel(**kwargs):
+        @numba.njit
+        def add_kernel_numba(x, y, out):
+            out[...] = x + y
+
+        return add_kernel_numba
+
+    def gpu_kernel(x_info):
+        return pl.pallas_call(
+            add_vectors_kernel,
+            out_shape=[jax.ShapeDtypeStruct(x_info.shape, x_info.dtype)],
+            interpret=jax.default_backend() == 'cpu',
+        )
+
+    prim = bst.event.XLACustomOp(
+        'add',
+        cpu_kernel,
+        gpu_kernel,
+    )
+
+    a = bst.random.rand(64)
+    b = bst.random.rand(64)
+    x_info = jax.ShapeDtypeStruct(a.shape, a.dtype)
+    r1 = prim(a, b, outs=[jax.ShapeDtypeStruct((64,), jax.numpy.float32)], x_info=x_info)
+    r2 = gpu_kernel(x_info)(a, b)
+
+    assert jnp.allclose(r1[0], r2[0])

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -112,10 +112,8 @@ class STP(Synapse):
         self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
 
     def update(self, pre_spike):
-        du = lambda u: self.U - u / self.tau_f
-        dx = lambda x: (1 - x) / self.tau_d
-        u = exp_euler_step(du, self.u.value)
-        x = exp_euler_step(dx, self.x.value)
+        u = exp_euler_step(lambda u: self.U - u / self.tau_f, self.u.value)
+        x = exp_euler_step(lambda x: (1 - x) / self.tau_d, self.x.value)
 
         # --- original code:
         #   if pre_spike.dtype == jax.numpy.bool_:
@@ -131,7 +129,7 @@ class STP(Synapse):
 
         self.u.value = u
         self.x.value = x
-        return u * x
+        return u * x * pre_spike
 
 
 class STD(Synapse):
@@ -167,8 +165,7 @@ class STD(Synapse):
         self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
 
     def update(self, pre_spike):
-        dx = lambda x: (1 - x) / self.tau
-        x = exp_euler_step(dx, self.x.value)
+        x = exp_euler_step(lambda x: (1 - x) / self.tau, self.x.value)
 
         # --- original code:
         # self.x.value = bm.where(pre_spike, x - self.U * self.x, x)
@@ -176,7 +173,7 @@ class STD(Synapse):
         # --- simplified code:
         self.x.value = x - pre_spike * self.U * self.x.value
 
-        return self.x.value
+        return self.x.value * pre_spike
 
 
 class AMPA(Synapse):
@@ -315,6 +312,4 @@ class GABAa(AMPA):
         T: ArrayLike = 1.0 * u.mM,
         T_dur: ArrayLike = 1.0 * u.ms,
     ):
-        super().__init__(alpha=alpha, beta=beta, T=T,
-                         T_dur=T_dur, name=name,
-                         in_size=in_size)
+        super().__init__(alpha=alpha, beta=beta, T=T, T_dur=T_dur, name=name, in_size=in_size)

brainstate/nn/_dyn_impl/_rate_rnns.py

@@ -23,7 +23,7 @@ import jax.numpy as jnp
 
 from brainstate import random, init, functional
 from brainstate._state import HiddenState, ParamState
-from brainstate.nn._interaction._connections import Linear
+from brainstate.nn._interaction._linear import Linear
 from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
 

brainstate/nn/_dynamics/_projection_base.py

@@ -207,7 +207,7 @@ class AlignPostProj(Interaction):
 
 
 class DeltaProj(Interaction):
-    """Full-chain of the synaptic projection for the Delta synapse model.
+    r"""Full-chain of the synaptic projection for the Delta synapse model.
 
     The synaptic projection requires the input is the spiking data, otherwise
     the synapse is not the Delta synapse model.

brainstate/nn/_exp_euler.py

@@ -32,7 +32,7 @@ __all__ = [
 def exp_euler_step(
     fn: Callable, *args, **kwargs
 ):
-    """
+    r"""
     One-step Exponential Euler method for solving ODEs.
 
     Examples

brainstate/nn/_interaction/__init__.py

@@ -13,20 +13,29 @@
 # limitations under the License.
 # ==============================================================================
 
-from ._connections import *
-from ._connections import __all__ as connections_all
+from ._conv import *
+from ._conv import __all__ as conv_all
 from ._embedding import *
 from ._embedding import __all__ as embed_all
+from ._linear import *
+from ._linear import __all__ as linear_all
 from ._normalizations import *
 from ._normalizations import __all__ as normalizations_all
 from ._poolings import *
 from ._poolings import __all__ as poolings_all
 
 __all__ = (
-    connections_all +
+    conv_all +
+    linear_all +
     normalizations_all +
     poolings_all +
     embed_all
 )
 
-del connections_all, normalizations_all, poolings_all, embed_all
+del (
+    conv_all,
+    linear_all,
+    normalizations_all,
+    poolings_all,
+    embed_all
+)