https://github.com/chaobrain/braintaichi
Leveraging Taichi Lang to customize brain dynamics operators.
https://github.com/chaobrain/braintaichi
brain-dynamics event-driven-computation sparse-computation
Last synced: about 2 months ago
JSON representation
Leveraging Taichi Lang to customize brain dynamics operators.
- Host: GitHub
- URL: https://github.com/chaobrain/braintaichi
- Owner: chaobrain
- License: apache-2.0
- Created: 2024-08-24T19:58:48.000Z (8 months ago)
- Default Branch: main
- Last Pushed: 2025-02-06T08:49:08.000Z (2 months ago)
- Last Synced: 2025-02-06T09:36:15.549Z (2 months ago)
- Topics: brain-dynamics, event-driven-computation, sparse-computation
- Language: Python
- Homepage: https://braintaichi.readthedocs.io
- Size: 1.75 MB
- Stars: 0
- Watchers: 2
- Forks: 0
- Open Issues: 8
-
Metadata Files:
- Readme: README.md
- Changelog: changelog.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Security: SECURITY.md
Awesome Lists containing this project
- awesome-jax - braintaichi - Leveraging Taichi Lang to customize brain dynamics operators. <img src="https://img.shields.io/github/stars/chaobrain/braintaichi?style=social" align="center"> (Libraries / New Libraries)
README
# Leveraging Taichi Lang to Customize Brain Dynamics Operators
[``braintaichi``](https://github.com/chaoming0625/braintaichi) leverages Taichi Lang to customize brain dynamics operators.
## Quick Start
```python
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:
```bash
pip install braintaichi --upgrade
```
## Documentation
The official documentation is hosted on Read the Docs: [https://braintaichi.readthedocs.io](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:
```bibtex
@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}
}
```