Leveraging Taichi Lang to Customize Brain Dynamics Operators

braintaichi leverages Taichi Lang to customize brain dynamics operators.

Quick Start

import taichi as ti

import braintaichi as bti


# define the custom kernel

@ti.kernel
def transpose_bool_homo_kernel(
    values: ti.types.ndarray(ndim=1),
    indices: ti.types.ndarray(ndim=1),
    indptr: ti.types.ndarray(ndim=1),
    events: ti.types.ndarray(ndim=1),
    out: ti.types.ndarray(ndim=1)
):
    value = values[0]
    ti.loop_config(serialize=True)
    for row_i in range(indptr.shape[0] - 1):
        if events[row_i]:
            for j in range(indptr[row_i], indptr[row_i + 1]):
                out[indices[j]] += value


kernel = bti.XLACustomOp(
    cpu_kernel=transpose_bool_homo_kernel,
    gpu_kernel=transpose_bool_homo_kernel,
)

# run with the sample data

import numpy as np
import jax
import jax.numpy as jnp
from scipy.sparse import csr_matrix

csr = csr_matrix((np.random.rand(10, 10) < 0.5).astype(float))
events = np.random.rand(10) < 0.5

out = kernel(
    jnp.array(csr.data),
    jnp.array(csr.indices),
    jnp.array(csr.indptr),
    events,
    outs=[jax.ShapeDtypeStruct([10], dtype=jnp.float32)]
)
print(out)

Installation

You can install braintaichi via pip:

pip install braintaichi --upgrade

Documentation

The official documentation is hosted on Read the Docs: https://braintaichi.readthedocs.io

See also the BDP ecosystem

We are building the brain dynamics programming ecosystem: https://ecosystem-for-brain-dynamics.readthedocs.io/

Citation

If you think braintaichi is significant in your work, please consider to cite the following pubilication:

@inproceedings{wang2024brainpy,
    title={A differentiable brain simulator bridging brain simulation and brain-inspired computing},
    author={Wang, Chaoming and Zhang, Tianqiu and He, Sichao and Gu, Hongyaoxing and Li, Shangyang and Wu, Si},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024}
}