Node-API bindings for MLX
npm install @frost-beta/mlxA machine learning framework for Node.js, based on
MLX.
This project is not affiliated with Apple, you can support the development by
sponsoring me.
GPU support:
- Macs with Apple Silicon
CPU support:
- x64 Macs
- x64/arm64 Linux
* llm.js - Large language models
implemented with JavaScript, vision models included.
* not-a-vector-database -
Node.js module for storing and searching embeddings.
* sisi - Semantic Image Search CLI.
* clip - Node.js module for the CLIP
model.
* train-model-with-js -
Train a simple text generation model with JavaScript.
* train-llama3-js - Train a
tiny Llama3 model with parquet datasets.
* train-japanese-llama3-js -
Train a Japanese language model.
* fine-tune-decoder-js -
Fine-tune a decoder-only model.
``typescript
import mlx from '@frost-beta/mlx';
const {core: mx, nn} = mlx;
const model = new nn.Sequential(
new nn.Sequential(new nn.Linear(2, 10), nn.relu),
new nn.Sequential(new nn.Linear(10, 10), new nn.ReLU()),
new nn.Linear(10, 1),
mx.sigmoid,
);
const y = model.forward(mx.random.normal([32, 2]));
console.log(y);
`
There is currently no documentations for JavaScript APIs, please check the
TypeScript definitions for available APIs, and MLX's official website
for documentations.
The JavaScript APIs basically duplicate the official Python APIs by converting
the API names from snake_case to camelCase. For example the mx.not_equalmx.notEqual
Python API is renamed to in JavaScript.
There are a few exceptions due to limitations of JavaScript:
* JavaScript numbers are always floating-point values, so the default dtype
of mx.array(42) is mx.float32 instead of mx.int32.mx.var
* The API is renamed to mx.variance.mx.add(a, b)
* Operator overloading does not work, use instead of a + b.[]
* Indexing via operator does not work, use array.item andarray.itemPut_
methods instead (the _ suffix means inplace operation).delete array
* does nothing, you must wait for garbage collection to get themx.dispose
array's memory freed or use .Module
* The instances can not be used as functions, the forward method must
be used instead.
Some features are not supported yet and will be implemented in future:
* The distributed module has not been implemented.mx.custom_function
* The API has not been implemented.export
* The APIs have not been implemented.mx.vmap
* The custom Metal kernel has not been implemented.
* It is not supported using JavaScript Array as index.
* The function passed to must have all parameters being mx.array.inputs
* The captured /outputs parameters of mx.compile has not beenmx.array
implemented.
* When creating a from JavaScript Array, the Array must only includemx.uniform(0, 1, [2, 2])
primitive values.
* The APIs only accept plain parameters, e.g. . Namedmx.uniform({shape: [2, 2]})
parameter calls like has not been implemented..npz
* The tensor format is not supported yet.
There are a few new APIs in node-mlx, for solving
JavaScript-only problems.
#### Constructor of mx.array
You can pass following JavaScript types to mx.array:
* number/boolean/mx.ComplexArray
* new Array(length)
* - Creates a 1D array filled with 0.Buffer
* - Same with UInt8Array.Int8Array
* /Uint8Array/Int16Array/Uint16Array/Int32Array/Uint32Array/Float32Array
#### mx.array.toTypedArray
While it is possible to use the tolist() method of mx.array to convert theArray
array to JavaScript , a copy of data is invovled.
For 1D arrays whose dtype has a representation in JavaScript
TypedArray,
you can use the toTypedArray() method to expose the data to JavaScript without
copying.
`typescript`
const buffer: Uint8Array = mx.array([1, 2, 3, 4], mx.uint8);
#### mx.asyncEval
The mx.eval API is synchronous that the main thread would be blocked waiting
for the result, which breaks the assumption of Node.js that nothing should block
in main thread, and results in a hanging process that not responding to
anything, including tries to end the process with Ctrl+C.
The mx.asyncEval API is the asynchronous version that returns a Promise that
can be awaited for. It is useful when the program runs a large model and you
want to make the app alive to user interactions while doing computations.
`typescript`
const y = model.forward(x);
await mx.asyncEval(y);
#### mx.tidy
This is the same with tf.tidy
API of TensorFlow.js, which cleans up all intermediate tensors allocated in the
passed functions except for the returned ones.
`typescript`
let result = mx.tidy(() => {
return model.forward(x);
});
In addition, it also works with async functions.
`typescript`
await mx.tidy(async () => { ... });
This is the same with tf.dispose
API of TensorFlow.js, which cleans up all the tensors found in the object.
`typescript`
mx.dispose({ a: mx.array([1, 2, 3, 4]) });
mx.dispose(mx.array([1]), mx.array([2]));
There is no built-in complex numbers in JavaScript, and we use objects to
represent them:
`typescript`
interface Complex {
re: number;
im: number;
}
You can also use the mx.Complex(real, imag?) helper to create complex numbers.
Slice in JavaScript is represented as object:
`typescript`
interface Slice {
start: number | null;
stop: number | null;
step: number | null;
}
You can also use the mx.Slice(start?, stop?, step?) helper to create slices.
The JavaScript standard does not allow using ... as values. To use ellipsis as"..."
index, use string instead.
When using arrays as indices, make sure a integer dtype is specified because
the default dtype is float32, for examplea.index(mx.array([ 1, 2, 3 ], mx.uint32)).
Here are some examples of translating Python indexing code to JavaScript:
#### Getters
| Python | JavaScript |
|--------------------------------------|----------------------------------------------------|
| array[None] | array.index(null) |array[Ellipsis, ...]
| | array.index('...', '...') |array[1, 2]
| | array.index(1, 2) |array[True, False]
| | array.index(true, false) |array[1::2]
| | array.index(mx.Slice(1, None, 2)) |array[mx.array([1, 2])]
| | array.index(mx.array([1, 2], mx.int32)) |array[..., 0, True, 1::2]
| | array.index('...', 0, true, mx.Slice(1, null, 2) |
#### Setters
| Python | JavaScript |
|--------------------------------------|--------------------------------------------------------------|
| array[None] = 1 | array.indexPut_(null, 1) |array[Ellipsis, ...] = 1
| | array.indexPut_(['...', '...'], 1) |array[1, 2] = 1
| | array.indexPut_([1, 2], 1) |array[True, False] = 1
| | array.indexPut_([true, false], 1) |array[1::2] = 1
| | array.indexPut_(mx.Slice(1, null, 2), 1) |array[mx.array([1, 2])] = 1
| | array.indexPut_(mx.array([1, 2], mx.int32), 1) |array[..., 0, True, 1::2] = 1
| | array.indexPut_(['...', 0, true, mx.Slice(1, null, 2)], 1) |
#### Translating between Python/JavaScript index types
| Python | JavaScript |
|----------------------|------------------------------|
| None | null |Ellipsis
| | "..." |...
| | "..." |123
| | 123 |True
| | true |False
| | false |:
| or :: | mx.Slice() |1:
| or 1:: | mx.Slice(1) |:3
| or :3: | mx.Slice(null, 3) |::2
| | mx.Slice(null, null, 2) |1:3
| | mx.Slice(1, 3) |1::2
| | mx.Slice(1, null, 2) |:3:2
| | mx.Slice(null, 3, 2) |1:3:2
| | mx.Slice(1, 3, 2) |mx.array([1, 2])
| | mx.array([1, 2], mx.int32) |
For building on platforms other than Macs with Apple Silicon, you must have blas
installed.
`bash`Linux
sudo apt-get install -y libblas-dev liblapack-dev liblapacke-devx64 Mac
brew install openblas
This project is mixed with C++ and TypeScript code, and uses
cmake-js to build the native code.
`bash`
git clone --recursive https://github.com/frost-beta/node-mlx.git
cd node-mlx
npm install
npm run build -- -p 8
npm run test
The prebuilt binaries are uploaded to the GitHub Releases, when installing
node-mlx from npm registry, the prebuilt binaries will always be downloaded and
there is no fallback for building from source code.
The version string is always 0.0.1-dev in package.json, which means local
development, and npm package can only be published via GitHub workflow by
pushing a new tag.
Before matching the features and stability of the official Python APIs, this
project will stay on 0.0.x for npm versions.
The tests and most TypeScript source code were converted from the Python code
of the official MLX project, when updating the deps/mlx` submodule, review
every new commit and make sure changes to Python APIs/tests/implementations are
also reflected in this repo.
Check out the contribution guidelines for more information
on contributing to node-mlx.