NuGet Package Manager

    @shumai/shumai

    0.0.4 • Public • Published

    shumai_logo_light shumai_logo_dark

    A fast differentiable tensor library for research in TypeScript and JavaScript. Built with bun + flashlight.

    ⚠️ This is experimental software! ⚠️

    docs build tests npm Discord GitHub commit activity GitHub

    Quickstart

    Install Bun and ArrayFire, then run:

    bun install @shumai/shumai
    

    Only macOS and Linux are supported. Linux installs default to GPU computation with CUDA, and macOS to CPU. Detailed install instructions below.

    Install is work in progress: please file an issue if you run into problems.

    Why build this?

    With Shumai, we hope to make

    • Creating datasets
      • JavaScript, with native typed arrays and a JIT compiler, is perfect for twiddling with data before it can be made into big, flat GPU-compatible arrays.
    • 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 facilitates better application development

    Usage

    shumai will always attempt to use an attached GPU or accelerator; although CPU computation will use the ArrayFire CPU backend, which is not well-optimized.

    We hope to support the ArrayFire OpenCL backend and other non-ArrayFire tensor backends soon.

    If shumai seems unusually 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

    The install procedure is a work in progress! If you have any problems building or installing, we would greatly appreciate filed issues. Please tell us about your platform/OS when you do.

    Prerequisites:

    • Ensure you have bun installed (https://bun.sh).
    • Install ArrayFire. macOS users should install ArrayFire's CPU backend; Linux users should install the CUDA backend^.
      • macOS --- ArrayFire can easily be installed with Homebrew:
      brew install arrayfire
      
    • Linux --- instructions can be found here. On Ubuntu, ArrayFire can be installed via package managers (e.g. apt).

    Once bun and ArrayFire are installed, install the package and backing libs with bun:

    bun install @shumai/shumai

    ^Linux users can use the CPU backend by swapping the required package.json dependency from @shumai/linux_x64_shumai_flashlight to @shumai/linux_x64_shumai_flashlight_cpu, i.e. running:

    sed -i "s/linux_x64_shumai_flashlight/linux_x64_shumai_flashlight_cpu/g" package.json

    Building Native Libraries from Source

    Note: not required when developing TypeScript/Javascript library components locally.

    From source build instructions for:

    This process will build the dependent ffi libraries (libflashlight and libflashlight_binding) and pack them using 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.

    Building on macOS from Source

    First, install ArrayFire CPU with brew install arrayfire.

    Build and install Flashlight:

    mkdir -p $HOME/usr/ # installing flashlight here
    git clone --recursive --depth 1 https://github.com/flashlight/flashlight.git
    cd flashlight
    mkdir -p build
    cd build
    cmake .. \
      -DCMAKE_BUILD_TYPE=RelWithDebInfo \ # or as specified
      -DFL_ARRAYFIRE_USE_CPU=ON \
      -DFL_ARRAYFIRE_USE_CUDA=OFF \
      -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 Flashlight bindings for Shumai:

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

    Profiling

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

    Building on Linux from Source

    First install ArrayFire. The Linux build for shumai uses the CUDA backend, but from source, you can build the CPU backend as well (OpenCL support coming soon).

    Build and install Flashlight:

    mkdir -p $HOME/usr/ # installing flashlight here
    git clone --recursive --depth 1 https://github.com/flashlight/flashlight.git
    cd flashlight
    mkdir -p build
    cd build
    cmake .. \
      -DCMAKE_BUILD_TYPE=RelWithDebInfo \ # or as specified
      -DFL_ARRAYFIRE_USE_CPU=OFF \
      \ # swap with the above to build for 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

    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 # if built from source, else not needed
    make -j$(nproc)

    Contributing

    If you'd like to make changes to the core bindings or ffi, 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. 😁

    License

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

    Install

    npm i @shumai/shumai

    DownloadsWeekly Downloads

    25

    Version

    0.0.4

    License

    none

    Unpacked Size

    129 kB

    Total Files

    28

    Last publish

    Collaborators

    • shumaiml
    • jacobkahn
    • bwasti