braintaichi 0.0.5

Leveraging Taichi Lang to customize brain dynamics operators

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}
}