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

From PyPI

You can install braintaichi via pip:

pip install braintaichi --upgrade

From Source

To compile and install from source code:

Prerequisites

1. Install dependencies:

   pip install taichi==1.7.3 pybind11
   

Platform-specific Requirements

Linux with GPU support:

• NVIDIA GPU drivers
• CUDA Toolkit (version 12.8 recommended, see ci/linux/gpu/prepare_build_wheel_on_cuda.sh)
• Development packages:

  # Example for RHEL/CentOS
  yum install -y python3-devel cuda-cudart-devel
  

Windows:

• Visual Studio with C++ development tools (Installation Guide)

Build and Install

git clone https://github.com/chaoming0625/braintaichi.git
cd braintaichi
pip install -e .

Or alternatively:

python setup.py install

Documentation

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

See also the ecosystem

We are building the brain modeling ecosystem: https://brainmodeling.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}
}