@shumai/shumai_cpu v0.0.1

Fast machine learning in JavaScript with bun + flashlight. ⚠️ This is experimental software! ⚠️

• Usage
• Install
• Build from source
• Contributing
• Supported operations

Why build this?

Here are some of the areas that shumai hopes to tackle:

• Creating datasets
  • JavaScript, with native typed arrays and a JIT compiler, is perfect for twiddling with data before it can be made into a big flat GPU-compatible array
• Training small models
  • FFI bindings in Bun are crazy fast (~3ns), so JS gets out of the way when training small models
• Advanced/fine-grained training/inference logic
  • Bun uses the JSC JIT compiler, meaning you can confidently write complex training logic without needing a native C++ implementation
• Building applications
  • JavaScript has a large HUGE ecosystem, which makes application development a lot easier

Usage

shumai will always attempt to use an attached GPU or accelerator, so it should be quite fast. If it seems slow, please file an issue!

Standard array utilities:

import * as sm from "@shumai/shumai"

// create a 1024 by 1024 tensor, randomly filled with normal distribution
let X = sm.randn([1024, 1024])
let W = sm.identity(1024)
let Y = X.matmul(W)
console.log(Y.shape)

Conversion to and from JavaScript native arrays:

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

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

// tensors can be converted back to native JavaScript
const Y_data = Y.toFloat32Array()

// scalar tensors can be converted to JavaScript numbers
const total : number = 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()

W.grad // this gradient is now populated

// copy W without allowing gradient updates
const Y = W.detach()
Y.sum().backward() // nothing changes

Some more examples can be found here.

Supported operators can be found here.

Install

This is a current work in progress! If you have any problems building or installing, we would greatly appreciate filed issues.

Ensure you have bun installed (https://bun.sh).

On MacOS:

brew install arrayfire
bun install @shumai/shumai

On Linux (Ubuntu/Debian):

sudo apt install arrayfire-cuda3-cuda-11-6
bun install @shumai/shumai

Installing local build from source

Note: not required when developing the library locally

This process will require building the FFI for bun and then running npm pack to generate a @shumai/shumai_*.tgz package. You can then use npm install $PATH_TO_SOURCE/@shumai/shumai-*.tgz to install the package where you'd like.

MacOS from source

Install flashlight

mkdir -p $HOME/usr/ # installing flashlight here
brew install arrayfire
git clone --recursive --depth 1 https://github.com/bwasti/flashlight.git
cd flashlight
mkdir -p build
cd build
cmake .. -DFL_ARRAYFIRE_USE_CPU=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

Build bindings for shumai

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

(you can record perf stuff with xcrun xctrace record --template "Time Profiler" --launch $(which bun) train.js)

Linux from source

First, install flashlight.

Then, build bindings for shumai:

mkdir -p build && cd build
cmake .. \
    -DBUILD_SHARED_LIBS=ON \
    -DCMAKE_BUILD_TYPE=RelWithDebInfo \ # or as specified
    -Dflashlight_DIR=${FLASHLIGHT_INSTALL_PREFIX}/share/flashlight/cmake \
    -DArrayFire_DIR=${ARRAYFIRE_INSTALL_PREFIX}/share/ArrayFire/cmake
make -j$(nproc)

Contributing

If you'd like to make changes to the code, first build from source.

All files ending in *.inl or *_gen.ts are generated. These can be modified by editing scripts/gen_binding.py and running ./scripts/gen_all_binding.sh.

See the CONTRIBUTING file for style guidance and more info on how to help out. 😁

Supported operations

Some operations are supported as both static functions and methods on existing tensors.

License

shumai is MIT licensed, as found in the LICENSE file.