brainstate 0.1.0.post20250105__py2.py3-none-any.whl → 0.1.0.post20250126__py2.py3-none-any.whl

brainstate/event/_linear_mv_test.py

@@ -1,117 +0,0 @@
-# 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 jax
-import jax.numpy as jnp
-from absl.testing import parameterized
-
-import brainstate as bst
-from brainstate.event._linear_mv import Linear
-
-
-class TestEventLinear(parameterized.TestCase):
-    @parameterized.product(
-        homo_w=[True, False],
-        bool_x=[True, False],
-    )
-    def test1(self, homo_w, bool_x):
-        x = bst.random.rand(20) < 0.1
-        if not bool_x:
-            x = jnp.asarray(x, dtype=float)
-        m = Linear(20, 40, 1.5 if homo_w else bst.init.KaimingUniform(),  float_as_event=bool_x)
-        y = m(x)
-        print(y)
-
-        self.assertTrue(jnp.allclose(y, (x.sum() * m.weight.value) if homo_w else (x @ m.weight.value)))
-
-    def test_grad_bool(self):
-        n_in = 20
-        n_out = 30
-        x = bst.random.rand(n_in) < 0.3
-        fn = Linear(n_in, n_out, bst.init.KaimingUniform())
-
-        with self.assertRaises(TypeError):
-            print(jax.grad(lambda x: fn(x).sum())(x))
-
-    @parameterized.product(
-        bool_x=[True, False],
-        homo_w=[True, False]
-    )
-    def test_vjp(self, bool_x, homo_w):
-        n_in = 20
-        n_out = 30
-        if bool_x:
-            x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
-        else:
-            x = bst.random.rand(n_in)
-
-        fn = Linear(n_in, n_out, 1.5 if homo_w else bst.init.KaimingUniform(), float_as_event=bool_x)
-        w = fn.weight.value
-
-        def f(x, w):
-            fn.weight.value = w
-            return fn(x).sum()
-
-        r1 = jax.grad(f, argnums=(0, 1))(x, w)
-
-        # -------------------
-        # TRUE gradients
-
-        def f2(x, w):
-            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
-            return y.sum()
-
-        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
-        self.assertTrue(jnp.allclose(r1[0], r2[0]))
-
-        if not jnp.allclose(r1[1], r2[1]):
-            print(r1[1] - r2[1])
-
-        self.assertTrue(jnp.allclose(r1[1], r2[1]))
-
-    @parameterized.product(
-        bool_x=[True, False],
-        homo_w=[True, False]
-    )
-    def test_jvp(self, bool_x, homo_w):
-        n_in = 20
-        n_out = 30
-        if bool_x:
-            x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
-        else:
-            x = bst.random.rand(n_in)
-
-        fn = Linear(n_in, n_out, 1.5 if homo_w else bst.init.KaimingUniform(),
-                    float_as_event=bool_x)
-        w = fn.weight.value
-
-        def f(x, w):
-            fn.weight.value = w
-            return fn(x)
-
-        o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-
-        # -------------------
-        # TRUE gradients
-
-        def f2(x, w):
-            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
-            return y
-
-        o2, r2 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-        self.assertTrue(jnp.allclose(o1, o2))
-        self.assertTrue(jnp.allclose(r1, r2))

brainstate/event/_misc.py

@@ -1,34 +0,0 @@
-# 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 typing import Union
-
-import numpy as np
-
-__all__ = [
-    'FloatScalar',
-    'IntScalar',
-]
-
-FloatScalar = Union[
-    np.number,  # NumPy scalar types
-    float,  # Python scalar types
-]
-
-IntScalar = Union[
-    np.number,  # NumPy scalar types
-    int,  # Python scalar types
-]

brainstate/event/_xla_custom_op.py

@@ -1,313 +0,0 @@
-# -*- 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 jaxlib.hlo_helpers import custom_call
-
-if jax.__version_info__ < (0, 4, 35):
-    from jax.lib import xla_client
-else:
-    import jax.extend as je
-
-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)
-    if jax.__version_info__ < (0, 4, 35):
-        xla_client.register_custom_call_target(target_name, capsule, "cpu")
-    else:
-        je.ffi.register_ffi_target(target_name, capsule, "cpu", api_version=0)
-
-    # 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_or_generator: Callable. The function defines the computation on CPU backend.
-      gpu_kernel_or_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_or_generator: Callable,
-        gpu_kernel_or_generator: Callable = None,
-        batching_translation: Callable = None,
-        jvp_translation: Callable = None,
-        transpose_translation: Callable = None,
-    ):
-        # 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_or_generator is not None:
-            self.def_cpu_kernel(cpu_kernel_or_generator)
-        if gpu_kernel_or_generator is not None:
-            self.def_gpu_kernel(gpu_kernel_or_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

@@ -1,55 +0,0 @@
-# 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])