一个为 TypeScript（及 JavaScript）打造的、快速、支持网络连接且可微分的张量库。基于 bun 和 flashlight 构建，适用于软件工程师和研究人员。

⚠️ 本项目为实验性软件！ ⚠️

• 使用方法
• 安装
• 从源码构建
• 基准测试
• 贡献

快速入门

安装 Bun 和 ArrayFire

curl https://bun.sh/install | bash
brew install arrayfire

curl https://bun.sh/install | bash
sudo apt install -y gnupg2 ca-certificates
sudo apt-key adv --fetch-key https://repo.arrayfire.com/GPG-PUB-KEY-ARRAYFIRE-2020.PUB
echo "deb https://repo.arrayfire.com/debian all main" | sudo tee /etc/apt/sources.list.d/arrayfire.list
sudo apt update
sudo apt install -y arrayfire-cpu3-dev arrayfire-cpu3-openblas

curl https://bun.sh/install | bash
sudo apt remove libarrayfire-dev libarrayfire-cpu3 libarrayfire-cpu-dev
sudo apt install -y libblas-dev liblapack-dev liblapacke-dev libfftw3-dev libboost-all-dev cmake make g++
cd /tmp
sudo rm -rf arrayfire
git clone https://github.com/arrayfire/arrayfire.git
cd arrayfire
cmake -Bbuild -DAF_BUILD_EXAMPLES=OFF -DCMAKE_BUILD_TYPE=Release -DAF_BUILD_UNIFIED=OFF -DAF_TEST_WITH_MTX_FILES=OFF -DBUILD_TESTING=OFF
make -j4 -Cbuild
sudo make install -Cbuild

然后运行：

bun install @shumai/shumai

目前仅支持 macOS 和 Linux 系统。Linux 默认使用 CUDA 进行 GPU 计算，而 macOS 则使用 CPU。详细的安装说明请参见 下方。

安装仍在开发中：如果您遇到任何问题，请提交 issue。

使用方法

shumai 会优先尝试使用连接的 GPU 或加速器；不过，CPU 计算将使用 ArrayFire 的 CPU 后端，该后端尚未经过充分优化。

我们计划尽快支持 ArrayFire 的 OpenCL 后端以及其他非 ArrayFire 的张量后端。

如果 shumai 运行速度异常缓慢，请提交 issue！

标准数组工具：

import * as sm from "@shumai/shumai"

// 创建一个 1024x1024 的随机正态分布张量
let X = sm.randn([1024, 1024])
let W = sm.identity(1024)
let Y = X.matmul(W)
console.log(Y.shape)

与原生 JavaScript 数组之间的转换：

const data : Float32Array = new Float32Array(128)
for (let i = 0; i < 128; ++i) {
  data[i] = Math.random()
}

const X : Tensor = sm.tensor(data)
const pi = sm.scalar(3.14)
const Y = X.mul(pi)

// 张量可以转换回原生 JavaScript
const Y_data = Y.toFloat32Array()

// 标量张量可以转换为 JavaScript 数值
const total : number = X.sum().toFloat32()

梯度：

const W = sm.randn([128, 128])
W.requires_grad = true

const X = sm.randn([128, 128])
const diff = X.sub(W)
const mse = diff.mul(diff).sum()
mse.backward()

W.grad // 此时梯度已被填充

// 复制 W 但不允许更新梯度
const Y = W.detach()
Y.sum().backward() // 不会发生任何变化

更多示例请参阅 这里。

支持的操作列表请参见 此处。

安装

安装流程仍在开发中！ 如果您在构建或安装过程中遇到任何问题，我们非常感谢您提交 issue。请务必告知我们您的平台/操作系统信息。

先决条件：

• 确保已安装 bun (https://bun.sh)。
• 安装 ArrayFire。macOS 用户应安装 ArrayFire 的 CPU 后端；Linux 用户则应安装 CUDA 后端^。
  • macOS --- 可以通过 Homebrew 轻松安装 ArrayFire：

  brew install arrayfire
  

• Linux --- 安装说明请参见 此处。在 Ubuntu 上，可以通过包管理器（如 apt）安装 ArrayFire。

一旦 bun 和 ArrayFire 安装完毕，即可使用 bun 安装包及其依赖库：

bun install @shumai/shumai

Windows 支持

虽然未正式支持，但有用户成功利用 Docker + WSL2 + Linux 环境来运行，甚至包括 CUDA 支持。

从源码构建原生库

注意： 在本地开发 TypeScript/Javascript 库组件时无需执行此步骤。

以下是针对以下平台的源码构建说明：

• macOS
• Linux

此过程将构建依赖的 FFI 库（libflashlight 和 libflashlight_binding），并使用 npm pack 将其打包成 @shumai/shumai_*.tgz 包。随后，您可以使用 npm install $PATH_TO_SOURCE/@shumai/shumai-*.tgz 在任意位置安装该包。

在 macOS 上从源码构建

首先，使用 brew install arrayfire 安装 ArrayFire CPU 版本。

构建并安装 Flashlight：

mkdir -p $HOME/usr/ # 将 Flashlight 安装到这里
git clone --recursive --depth 1 https://github.com/flashlight/flashlight.git
cd flashlight
mkdir -p build
cd build
cmake .. \
  -DCMAKE_BUILD_TYPE=Release \
  -DBUILD_SHARED_LIBS=ON  \
  -DCMAKE_INSTALL_PREFIX=$HOME/usr \
  -DFL_USE_ARRAYFIRE=ON \
  -DFL_ARRAYFIRE_USE_CPU=ON \
  -DFL_USE_ONEDNN=OFF \
  -DFL_BUILD_DISTRIBUTED=OFF \
  -DFL_BUILD_TESTS=OFF \
  -DFL_BUILD_EXAMPLES=OFF
make -j$(nproc)
make install

为 Shumai 构建 Flashlight 绑定：

cd shumai
mkdir -p build
cd build
cmake .. -Dflashlight_DIR=$HOME/usr/share/flashlight/cmake/
make -j$(nproc)

性能分析

在 macOS 上，你可以使用 xcrun xctrace record --template "Time Profiler" --launch $(which bun) train.js 来录制性能分析数据。

在 Linux 上从源码构建

首先安装 ArrayFire。Shumai 的 Linux 构建默认使用 CUDA 后端，但如果你从源码编译，也可以构建 CPU 后端（OpenCL 支持即将推出）。

构建并安装 Flashlight：

mkdir -p $HOME/usr/ # 将 Flashlight 安装到这里
git clone --recursive --depth 1 https://github.com/flashlight/flashlight.git
cd flashlight
mkdir -p build
cd build
cmake .. \
  -DCMAKE_BUILD_TYPE=RelWithDebInfo \ # 或根据需求指定
  -DFL_ARRAYFIRE_USE_CPU=OFF \
  \ # 切换为上述选项以构建 CPU 版本
  -DFL_ARRAYFIRE_USE_CUDA=ON \ 
  -DFL_BUILD_DISTRIBUTED=OFF \
  -DFL_USE_ONEDNN=OFF \
  -DFL_BUILD_TESTS=OFF \
  -DFL_BUILD_EXAMPLES=OFF \
  -DFL_BUILD_SCRIPTS=OFF \
  -DCMAKE_INSTALL_PREFIX=$HOME/usr/
make -j$(nproc)
make install

构建 Shumai 的绑定：

mkdir -p build && cd build
cmake .. \
    -DBUILD_SHARED_LIBS=ON \
    -DCMAKE_BUILD_TYPE=RelWithDebInfo \ # 或根据需求指定
    -Dflashlight_DIR=${FLASHLIGHT_INSTALL_PREFIX}/share/flashlight/cmake \
    -DArrayFire_DIR=${ARRAYFIRE_INSTALL_PREFIX}/share/ArrayFire/cmake # 如果是从源码构建，则需要；否则不需要
make -j$(nproc)

为什么要构建这个？

借助 Shumai，我们希望实现以下目标：

• 更轻松地创建数据集
  • JavaScript 拥有原生的类型化数组和 JIT 编译器，非常适合在将数据转换为大型、扁平化的 GPU 兼容数组之前进行数据处理。
• 更快地训练小型模型
  • Bun 中的 FFI 绑定速度极快（约 3 纳秒），因此在训练小型模型时，JavaScript 不会成为瓶颈。
• 更灵活、更精细的训练/推理逻辑
  • Bun 使用 JSC JIT 编译器，这意味着你可以放心地编写复杂的训练逻辑，而无需依赖原生 C++ 实现。
• 开发更具吸引力的应用程序
  • JavaScript 拥有一个庞大巨大的生态系统，这有助于更好地进行应用程序开发。

基准测试

基准测试数据来自 https://github.com/shumai-org/benchmarks。

在 Apple M1 Pro 上：

在 Nvidia GP100 上：

内存使用

虽然开箱即用的内存管理在许多情况下已经足够，但通过调整内存使用策略，可以显著减少垃圾回收器带来的额外开销，从而提升性能。

import { util } from '@shumai/shumai'

util.memoryOptions({
  lowerBoundThreshold: 100e6, // 100MB
  upperBoundThreshold: 5e9, // 5GB
  delayBetweenGCs: 1000 // 1 秒
})

请特别注意 upperBoundThreshold 参数：一旦超过该阈值，每次分配张量时都会强制触发垃圾回收，此时 delayBetweenGCs 将被忽略。合理设置该参数以充分利用硬件资源，可以显著提升性能。

统计信息

graph TD
  OpA(操作A) --> statsA{{"统计A"}};
  OpB(操作B) --> statsA;
  statsA --> LoggerA{{"LoggerConsole A"}};
  LoggerA --> Stdout(("标准输出"));
  OpC(操作C) --> statsA;
  OpD(操作D) --> statsA;
  statsA --> LoggerB("自定义日志记录器B");
  LoggerB --> Disk(("磁盘"));

收集统计信息的基本用法非常简单，只需使用默认的 StatsLoggerConsole 添加一个收集器即可。

import { stats, StatsLoggerConsole, rand, matmul } from '@shumai/shumai'

stats.enabled = true // 之后的所有操作都会捕获统计信息

// 执行操作...

stats.enabled = false // 之后的所有操作将不再捕获统计信息

虽然上述示例对于简单的使用场景已经足够，但如果你希望在多个线程、进程和/或主机上捕获统计信息，StatsLoggerHttp 将能满足你的需求。

graph TD
  subgraph 主机C
    Processor("Http日志处理器")
    style Processor stroke:#222,stroke-width:4px,stroke-dasharray:5 5
  end
  subgraph 主机A
    OpA(操作A) --> statsA{{"统计A"}};
    OpB(操作B) --> statsA;
    statsA --> LoggerA{{"Http日志记录器A"}};
    LoggerA --> Processor;
  end
  subgraph 主机B
    OpC(操作C) --> statsB{{"统计B"}};
    OpD(操作D) --> statsB;
    statsB --> LoggerB{{"Http日志记录器B"}};
    LoggerB --> Processor;
  end

import { StatsLoggerHttp } from '@shumai/shumai'

stats.logger = new StatsLoggerHttp({ url: 'http://localhost:4242' })

对于更定制化的需求，你可以提供自己的日志记录器：

import { StatsLogger, StatsLoggerData } from '@shumai/shumai'

class CustomLogger implements StatsLogger {
  async process(data: StatsLoggerData): Promise<void> {
    const summary = data.collector.getSummary()
    console.log('收集器统计信息:', summary)
  }
}

stats.logger = new CustomLogger()

默认情况下，堆栈跟踪功能是禁用的，因为它会增加 50% 以上的开销。不过，可以通过设置 stats.collectStacks = true 来启用它。

作用域内的统计信息

如果你希望隔离统计分析，也可以这样做：

import { collectStats } from '@shumai/shumai'

const scopedStats = collectStats(() => {
  // 执行操作...
}/*, StatsCollectorOptions | StatsLogger */)
console.log(scopedStats.getSummary())

贡献

如果你想对核心绑定或 FFI 进行修改，请先从源代码构建。

所有以 *.inl 或 *_gen.ts 结尾的文件都是生成的。这些文件可以通过编辑 scripts/gen_binding.py 并运行 ./scripts/gen_all_binding.sh 来修改。

有关风格指南及如何参与贡献的更多信息，请参阅 CONTRIBUTING 文件。 😁

许可证

shumai 采用 MIT 许可证，具体内容请参见 LICENSE 文件。

Shumai 快速上手指南

Shumai 是一个基于 TypeScript/JavaScript 的高性能、可微分张量库，专为软件工程师和研究人员设计。它利用 Bun 运行时和 Flashlight 后端，提供快速的网络连通性和 GPU 加速能力。

⚠️ 注意：本项目目前处于实验阶段，主要支持 macOS 和 Linux 系统。

1. 环境准备

在开始之前，请确保你的系统满足以下要求：

• 操作系统：macOS 或 Linux (Windows 用户建议通过 WSL2 + Docker 使用)。
• 运行时：必须安装 Bun。
• 底层依赖：必须安装 ArrayFire。
  • macOS：默认使用 CPU 后端。
  • Linux：默认推荐使用 CUDA 后端（需 NVIDIA GPU），也可构建 CPU 后端。

安装前置依赖

macOS 用户

使用 Homebrew 一键安装 ArrayFire：

curl https://bun.sh/install | bash
brew install arrayfire

Linux 用户 (Ubuntu x86-64)

使用官方源安装 ArrayFire (CPU/OpenBLAS 版本)：

curl https://bun.sh/install | bash
sudo apt install -y gnupg2 ca-certificates
sudo apt-key adv --fetch-key https://repo.arrayfire.com/GPG-PUB-KEY-ARRAYFIRE-2020.PUB
echo "deb https://repo.arrayfire.com/debian all main" | sudo tee /etc/apt/sources.list.d/arrayfire.list
sudo apt update
sudo apt install -y arrayfire-cpu3-dev arrayfire-cpu3-openblas

注：若需使用 GPU 加速，请参考 ArrayFire 官方文档安装 CUDA 版本。

Linux 用户 (Ubuntu ARMv8)

需要从源码编译 ArrayFire：

curl https://bun.sh/install | bash
sudo apt remove libarrayfire-dev libarrayfire-cpu3 libarrayfire-cpu-dev
sudo apt install -y libblas-dev liblapack-dev liblapacke-dev libfftw3-dev libboost-all-dev cmake make g++
cd /tmp
sudo rm -rf arrayfire
git clone https://github.com/arrayfire/arrayfire.git
cd arrayfire
cmake -Bbuild -DAF_BUILD_EXAMPLES=OFF -DCMAKE_BUILD_TYPE=Release -DAF_BUILD_UNIFIED=OFF -DAF_TEST_WITH_MTX_FILES=OFF -DBUILD_TESTING=OFF
make -j4 -Cbuild
sudo make install -Cbuild

2. 安装步骤

完成上述依赖安装后，使用 Bun 安装 Shumai：

bun install @shumai/shumai

提示：如果安装过程中遇到问题，由于安装流程仍在完善中，建议在 GitHub 仓库提交 Issue 并附上你的平台信息。

3. 基本使用

Shumai 会自动检测并使用连接的 GPU 或加速器；若无可用设备，将回退到 ArrayFire CPU 后端。

创建张量与矩阵运算

import * as sm from "@shumai/shumai"

// 创建一个 1024x1024 的张量，填充正态分布随机数
let X = sm.randn([1024, 1024])
let W = sm.identity(1024)

// 矩阵乘法
let Y = X.matmul(W)

console.log(Y.shape)

与原生 JavaScript 数组互转

// 原生 JS 数组 -> Shumai Tensor
const data = new Float32Array(128)
for (let i = 0; i < 128; ++i) {
  data[i] = Math.random()
}

const X = sm.tensor(data)
const pi = sm.scalar(3.14)
const Y = X.mul(pi)

// Shumai Tensor -> 原生 JS 数组
const Y_data = Y.toFloat32Array()

// 标量张量 -> JS Number
const total = X.sum().toFloat32()

自动求导 (Gradients)

const W = sm.randn([128, 128])
W.requires_grad = true // 开启梯度追踪

const X = sm.randn([128, 128])
const diff = X.sub(W)
const mse = diff.mul(diff).sum()

mse.backward() // 反向传播

console.log(W.grad) // 梯度已计算并存入 .grad

// 停止梯度追踪
const Y = W.detach()
Y.sum().backward() // 不会改变 W 的梯度