keras-js manual

keras-js is the Keras 3 API for JavaScript. It runs on jax-js, which compiles array programs to WebGPU and WebAssembly kernels, the way Keras 3 runs on JAX in Python. This manual covers models, layers, training, saving, and how keras-js differs from Python Keras.

Getting started

Install

npm install @lamng3/keras-js @jax-js/jax
npm install h5wasm fflate   # optional: .keras and .weights.h5 files

Without a bundler, use an import map:

<script type="importmap">
  {
    "imports": {
      "@jax-js/jax": "https://cdn.jsdelivr.net/npm/@jax-js/jax@0.1.27/dist/index.js",
      "@lamng3/keras-js": "https://lamng3.com/keras-js/keras-js.js",
      "h5wasm": "https://esm.sh/h5wasm@0.10.3",
      "fflate": "https://esm.sh/fflate@0.8.3"
    }
  }
</script>
<script type="module">
  import * as keras from "@lamng3/keras-js";
  console.log(await keras.init()); // "webgpu" or "wasm"
</script>

First model

import * as keras from "@lamng3/keras-js";

await keras.init(); // WebGPU when available, else WebAssembly

const model = new keras.Sequential([
  keras.Input([784]),
  new keras.layers.Dense(128, { activation: "relu" }),
  new keras.layers.Dropout(0.2),
  new keras.layers.Dense(10),
]);

model.compile({
  optimizer: "adam",
  loss: new keras.losses.SparseCategoricalCrossentropy({ fromLogits: true }),
  metrics: ["accuracy"],
});

const history = await model.fit(xTrain, yTrain, { epochs: 5, batchSize: 64, validationSplit: 0.1 });
const [loss, accuracy] = await model.evaluate(xTest, yTest);
const logits = model.predict(xTest);

x and y can be jax-js arrays, nested JS arrays, or typed arrays. The rest reads like Keras: the argument names are the Keras names in camelCase (batch_size → batchSize).

Devices

keras.init() initializes jax-js and selects a device: WebGPU when the browser exposes navigator.gpu, otherwise WebAssembly. Pass a device to force one:

await keras.init("wasm"); // "webgpu" | "wasm" | "webgl" | "cpu"
keras.backend.device(); // the active device

In Node, keras-js uses WebAssembly.

Arrays and ownership

keras-js computes with jax-js arrays (keras.ops is the jax-js NumPy API). JavaScript has no destructors, so jax-js frees GPU memory by reference counting, and operations consume their arguments. Use .ref to keep an array alive after passing it in:

import { numpy as np } from "@jax-js/jax";

const x = np.array([1, 2, 3]);
const y = x.ref.mul(2); // x stays alive
const z = x.add(1); // x is consumed here

The rules in keras-js:

A model keeps exactly one reference to each weight; the test suite checks that fit() leaks nothing. See the jax-js README for the full rules.

Models

Sequential

const model = new keras.Sequential([
  keras.Input([28, 28, 1]),
  new keras.layers.Conv2D(32, 3, { activation: "relu" }),
  new keras.layers.MaxPooling2D(2),
  new keras.layers.Flatten(),
  new keras.layers.Dense(10),
]);
model.add(new keras.layers.Activation("softmax"));
model.summary();

Without keras.Input, the model builds on its first batch, as in Keras. model.layers excludes the input layer; model.pop() removes the last layer.

Functional API

keras.Input() returns a symbolic KerasTensor. Calling a layer on symbolic tensors records the graph. JavaScript objects are not callable, so layer(x) is written layer.apply(x):

const L = keras.layers;
const image = keras.Input([32, 32, 3], { name: "image" });
const meta = keras.Input([4], { name: "meta" });

let h = new L.Conv2D(16, 3, { activation: "relu" }).apply(image);
h = new L.GlobalAveragePooling2D().apply(h);
h = new L.Concatenate().apply([h, meta]);

const label = new L.Dense(10, { name: "label" }).apply(h);
const score = new L.Dense(1, { name: "score" }).apply(h);

const model = new keras.Model({ inputs: [image, meta], outputs: [label, score] });

Layers can be shared (apply one layer object twice and the weights are shared), and models nest inside other models. model.getLayer(name) looks a layer up by name.

Subclassing

class Residual extends keras.Model {
  inner = new keras.layers.Dense(64, { activation: "relu" });
  outer = new keras.layers.Dense(64);

  call(x, { training } = {}) {
    return this.outer.apply(this.inner.apply(x.ref, { training })).add(x);
  }
}

Layers stored on properties (or in arrays on properties) are tracked automatically, as Keras tracks attributes. call follows the ownership rules: it receives its input and must consume it exactly once, so x.ref is used for the first of two uses.

Layers

Constructors take Keras' required arguments positionally, then an options object: new Dense(units, opts), new Conv2D(filters, kernelSize, opts), new Dropout(rate, opts), new Embedding(inputDim, outputDim, opts), new MultiHeadAttention(numHeads, keyDim, opts), new LSTM(units, opts). Every layer accepts name, trainable and dtype.

Convolutions accept strides, padding ("valid", "same"), dilationRate, groups and dataFormat ("channels_last" by default). Recurrent layers accept returnSequences, returnState, goBackwards and zeroOutputForMask (see Recurrent states).

MultiHeadAttention takes its tensors as a list: mha.apply([query, value]), mha.apply([query, value, key]), or mha.apply(x) for self-attention. Options such as useCausalMask and attentionMask (a [B, T, S] boolean tensor) go in the second argument and are saved with the graph.

Masking

Masks mark which timesteps are real and which are padding, as in Keras. A mask is created by Embedding({ maskZero: true }) (token 0 is padding) or by Masking({ maskValue }) (a step whose features all equal maskValue), and it flows through the model by itself:

const tokens = keras.Input([null], { dtype: "int32" });
let h = new L.Embedding(10_000, 64, { maskZero: true }).apply(tokens);
h = new L.LSTM(64, { returnSequences: true }).apply(h); // skips padded steps
const attn = new L.MultiHeadAttention(4, 16).apply([h, h]); // padded keys are never attended to
h = new L.GlobalAveragePooling1D().apply(new L.Add().apply([h, attn])); // averages real steps only

The loss uses the output mask: the mean runs over unmasked entries only. As in Keras 3.15 (JAX backend), metrics given to compile() do not use the output mask; the parity tests check both behaviors.

Eagerly, pass a mask with layer.apply(x, { mask }) and read one with const [y, mask] = layer.applyWithMask(x, { mask }). A custom layer passes masks through by setting supportsMasking = true, or computes its own by overriding computeMask(inputs, mask), which returns a new mask the caller owns.

Recurrent states

returnState: true adds the final states to the output: [output, h], or [output, h, c] for LSTM, and [merged, ...forwardStates, ...backwardStates] for Bidirectional. The initialState option starts a layer from given states, which is how encoder-decoder models are written:

const [, h, c] = new L.LSTM(128, { returnState: true }).apply(encoderInputs);
const decoded = new L.LSTM(128, { returnSequences: true }).apply(decoderInputs, { initialState: [h, c] });

Tensor options such as initialState and attentionMask are graph inputs: they are saved in the config as Keras writes them (initial_state, attention_mask) and load in Python.

Activations and initializers

Activations are named as in Keras ("relu", "gelu", "silu", "softmax", "leaky_relu", …; camelCase also works) or passed as functions. GELU is the exact (erf) form and leaky_relu uses a slope of 0.2, matching Keras 3.

Initializers are strings ("glorot_uniform", "he_normal", "orthogonal", …) or objects: new keras.initializers.RandomNormal({ stddev: 0.02, seed: 1 }). Seeds make initialization reproducible; keras.utils.setRandomSeed(n) seeds everything else.

Custom layers

class Scale extends keras.Layer {
  build(inputShape) {
    this.w = this.addWeight({ name: "w", shape: [inputShape.at(-1)], initializer: "ones" });
  }
  call(x) {
    return x.mul(this.w.value);
  }
}

build runs once with the input shape. Create weights there with addWeight (it accepts initializer, regularizer and trainable). In call, read a weight with .value, which returns a fresh reference. Output shapes are inferred automatically by running the layer on zeros; override computeOutputShape(inputShape) to skip that. To make the layer loadable from saved files, implement getConfig() and register it with keras.saving.registerSerializable.

Training

compile

model.compile({
  optimizer: new keras.optimizers.AdamW({ learningRate: 3e-4, weightDecay: 0.01, globalClipnorm: 1 }),
  loss: "categorical_crossentropy", // or one loss per output: [loss1, loss2]
  lossWeights: [1, 0.5],
  metrics: ["accuracy"], // or one list per output
  weightedMetrics: ["accuracy"], // these also see sampleWeight
  jitCompile: true, // default
});

With jitCompile (the default), the train, evaluation and prediction steps are each compiled once per input shape, so the first batch is slower than the rest.

fit, evaluate, predict

For several inputs or outputs, pass lists in the order of inputs and outputs. Logs use Keras names: loss, accuracy, val_loss, and for several outputs <output>_<metric>.

Sample and class weights follow Keras exactly. sampleWeight has one value per sample ([N]) or per timestep ([N, T], for sequence outputs), or one array per output; the loss becomes sum(loss × weight) / count, so a weight of 0 removes an example and 2 counts it twice. classWeight: { 0: 1, 1: 5 } turns class indices (or one-hot targets) into sample weights. In validationData, a third element holds validation sample weights. Plain metrics ignore weights; weightedMetrics use them and are logged with a weighted_ prefix when the name collides, as in Keras.

await model.fit(x, y, { sampleWeight: confidence, classWeight: { 0: 1, 1: 10 }, epochs: 5 });

fit waits on the device every yieldEvery batches (default 8), which keeps a page responsive. Callbacks that want per-batch logs set wantsBatchLogs, which syncs every batch.

Losses and metrics

Losses: MeanSquaredError, MeanAbsoluteError, MeanAbsolutePercentageError, BinaryCrossentropy, CategoricalCrossentropy, SparseCategoricalCrossentropy (with fromLogits and labelSmoothing), Huber, LogCosh, KLDivergence, CosineSimilarity, Hinge, or their snake_case names ("mse", "sparse_categorical_crossentropy", …). Any (yTrue, yPred) => perSample function also works.

Metrics: "accuracy" resolves to binary, categorical or sparse accuracy from the loss, as in Keras. Also BinaryAccuracy, CategoricalAccuracy, SparseCategoricalAccuracy, MeanSquaredError, MeanAbsoluteError, RootMeanSquaredError, Precision, Recall. Metrics support the stateful API too: updateState, await result(), resetState.

Optimizers

SGD (momentum, Nesterov), Adam (AMSGrad), AdamW, Adamax, RMSprop (momentum, centered), Adagrad. Every optimizer accepts clipnorm, clipvalue, globalClipnorm and weightDecay, and the update rules follow Keras 3. Learning rates are numbers or schedules:

const lr = new keras.optimizers.CosineDecay(1e-3, 10_000, { warmupTarget: 3e-3, warmupSteps: 500 });
new keras.optimizers.Adam({ learningRate: lr });

Also ExponentialDecay, InverseTimeDecay and PiecewiseConstantDecay. Scalars such as the learning rate are passed into the compiled step as arrays, so a schedule or callback can change them without recompiling.

Callbacks

History (returned by fit), ProgbarLogger, EarlyStopping, ModelCheckpoint, LearningRateScheduler, ReduceLROnPlateau, TerminateOnNaN, and LambdaCallback:

await model.fit(x, y, {
  epochs: 100,
  callbacks: [
    new keras.callbacks.EarlyStopping({ monitor: "val_loss", patience: 5, restoreBestWeights: true }),
    new keras.callbacks.LambdaCallback({ onEpochEnd: (epoch, logs) => chart.push(logs.loss) }),
  ],
});

Subclass keras.callbacks.Callback for your own. Hooks may be async; fit awaits them.

ModelCheckpoint(filepath, { monitor, saveBestOnly, saveWeightsOnly, mode, verbose, onSave }) saves a .keras file (or .weights.h5) at the end of each epoch, or only when monitor improves. filepath may contain {epoch} or {epoch:02d}. Node writes the file; in every runtime the latest save is in checkpoint.bytes, so keras.saving.loadModel(checkpoint.bytes) restores the best model, and onSave(bytes, path) can store it elsewhere (IndexedDB, a download).

Utilities

Regularization and freezing

kernelRegularizer: new keras.regularizers.L2(1e-4) (or "l1", "l2", L1L2) adds a penalty to the loss. Setting layer.trainable = false freezes a layer; a frozen BatchNormalization runs in inference mode, as in Keras, which is what fine-tuning needs:

const base = await keras.saving.loadModel(bytes);
base.trainable = false;
const model = new keras.Sequential([base, new keras.layers.Dense(5)]);

Custom training loops

Every model can run as a pure function of its weights. This is what fit does internally, and it works with any jax-js transformation:

import { valueAndGrad } from "@jax-js/jax";

const vars = model.trainableWeights;
const opt = new keras.optimizers.Adam({ learningRate: 1e-3 });

const lossFn = (values, x, y) => {
  const [[pred]] = model.statelessCall(values, model.nonTrainableWeights.map((v) => v.value), [x]);
  return keras.ops.mean(keras.ops.square(pred.sub(y)));
};

const [loss, grads] = valueAndGrad(lossFn)(vars.map((v) => v.value), x.ref, y.ref);
opt.applyGradients(grads.map((g, i) => [g, vars[i]]));

statelessCall(trainable, nonTrainable, inputs, { training, key }) returns [outputs, newNonTrainable, regularizationLoss]. For lower-level control, new keras.StatelessScope(variables, values) binds variables to values while a function runs.

Saving and Python interop

keras-js uses Keras 3's file formats with the same weight layout, so files move both ways.

.keras and .h5 need the optional h5wasm and fflate packages, loaded on first use. Safetensors and JSON need nothing.

Optimizer state. A compiled model's .keras file includes the optimizer's step count, learning rate and slots (Adam's moments, RMSprop's velocities, …) in Keras' layout, and loadModel restores them, so training resumes exactly where it stopped, whichever side saved the file. Pass { includeOptimizer: false } to saveModel for a smaller file. saving.loadOptimizerH5(model, bytes) restores the state into an already-compiled model.

Keras ops in graphs. Keras records how masks are computed as op nodes (NotEqual, Any, LogicalOr, ExpandDims, …) in functional configs. keras-js replays these ops, along with basic arithmetic and shape ops written with keras.ops. An op it does not know raises an error naming it.

Python to the browser. Save in Python and load in JavaScript:

model.save("classifier.keras")

const bytes = new Uint8Array(await (await fetch("classifier.keras")).arrayBuffer());
const model = await keras.saving.loadModel(bytes); // compiled if it was saved compiled
const probs = model.predict(images);

Browser to Python. Save in JavaScript and load in Python:

const bytes = await keras.saving.saveModel(model);
// download it, upload it, or fs.writeFileSync("model.keras", bytes) in Node

model = keras.saving.load_model("model.keras")

Both directions are tested against Keras 3.15 (see Parity).

Examples from keras.io

Five examples from keras.io/examples are ported in examples/keras-io/, line for line, and run on the Examples page:

Each file exports main({ epochs }). In Node:

npm run build
node examples/run.js text_classification_with_transformer --epochs 2

What changes in translation is what this manual describes: layer(x) is layer.apply(x), custom layers pass { training } on to their sublayers, and arrays used twice take a .ref. keras.datasets has no JavaScript counterpart; examples/datasets.js loads the same data (IMDB already passed through pad_sequences, built by tools/make_example_data.py).

Performance

• WebGPU is the fast path. On WebAssembly, jax-js runs dense layers well, but convolutions are slow: a 16→32-channel 3×3 convolution on a batch of 128 takes about 240 ms in jax-js 0.1.27 in Node. Use WebGPU for CNNs.
• Keep shapes fixed. Each new input shape compiles a new program. The last, smaller batch of an epoch compiles once more.
• Batch logs cost a sync. Prefer epoch-level callbacks; verbose: 1 in Node and callbacks with wantsBatchLogs sync every batch.

MNIST measurements on an Apple M2 Pro (batch 128, Adam with cosine decay from 1e-3, 10% of the 60,000 training images held out for validation, 10 epochs):

Reproduce with examples/mnist-node.ts and the MNIST page.

Differences from Python Keras

Not yet supported: stateful RNNs, recurrent dropout, return_attention_scores, mixed precision, tf.data-style datasets, and constraints. Subclassed models are not loadable from config unless registered. The roadmap tracks these.

Parity with Python Keras

npm test runs three suites:

• Forward parity. Nine models built and saved by Python Keras 3.15: an MLP with BatchNorm and LayerNorm, a CNN that uses every conv and pooling variant, a causal Conv1D stack, an Embedding/LSTM/GRU/Bidirectional/SimpleRNN stack, a transformer block with causal self-attention and cross-attention, a multi-input, multi-output graph with shared layers, two masked models (padded token sequences through LSTM, attention, BatchNorm, Bidirectional and pooling; Masking with left, right and inner padding merged by Concatenate), and an encoder-decoder that uses return_state and initial_state. Each is loaded from its .keras file and must reproduce Keras' predictions to 2e-4.
• Training parity. Five train_on_batch steps with SGD (Nesterov), Adam, AdamW, centered RMSprop, Adagrad, and AMSGrad with global clipping, with L2 regularization and BatchNorm. Losses must match to 1e-4 and every weight to 5e-4. Also: a masked per-timestep loss with temporal sample weights and weighted metrics, class_weight in fit(), and resuming four optimizers (Adam, centered RMSprop with momentum, AMSGrad with a learning rate schedule, SGD with momentum) from Keras checkpoints.
• Python round trip. Models built and trained in keras-js are saved as .keras and loaded by Python Keras, which must reproduce keras-js' predictions; this includes a masked model, an encoder-decoder with initialState, and a checkpoint whose optimizer state Python resumes. This suite runs when KERAS_PYTHON points at a Python with Keras, JAX and safetensors.

Regenerate the fixtures with npm run parity:fixtures.