Keras 3: one model, several compute engines.

One API, several engines

Keras was born as a layer on top of other compute engines (Theano, then TensorFlow) and for years it was, in practice, TensorFlow's high-level API. Keras 3, released in November 2023, went back to the original design: the definition of the model (layers, losses, optimizers, fit) is independent of the engine that does the math, the backend. The official documentation lists three for training: JAX, TensorFlow and PyTorch. The package code also ships a NumPy backend and an OpenVINO one, intended for inference.

flowchart TB
  M[Keras model] --> O[keras.ops]
  O --> T[PyTorch]
  O --> J[JAX]
  O --> F[TensorFlow]
  T --> C[CPU or GPU]
  J --> C
  F --> C

For a research team this has three practical consequences:

• The model is not tied to the engine. A model written in Keras can be trained in JAX (fast for small models, as you will see) and served inside a PyTorch application, or the other way around, with the same weights.
• The tensors are the backend's. With PyTorch, each tensor is a torch.Tensor and the model is a torch.nn.Module; with JAX, a jax.Array. You can drop down to the engine when needed.
• Randomness is not portable. Each backend has its own random number generator: the same seed does not give the same initial weights on two backends. You need to know this so as not to confuse noise with a bug.

The course environment has PyTorch 2.14.1 (CPU only) and JAX 0.12.0.1, with Keras 3.15.1. Not TensorFlow: the latest stable version, 2.21.0, does not publish wheels for Python 3.14, and the first one that does, 2.22.0rc0, is a release candidate. Its code appears at the end, not executed.

Configuring the backend

Keras decides the backend only once, when it is imported, and does not allow changing it afterward. It looks for it in this order: the KERAS_BACKEND environment variable and, if it is not set, the backend field of keras.json (in ~/.keras/ or in the KERAS_HOME folder); with neither of them, it uses TensorFlow. The runner of the level 5 lessons sets KERAS_BACKEND=torch and points KERAS_HOME to a temporary folder:

import json
import os
import numpy as np
import pandas as pd
import torch
import keras
from keras import layers, ops

torch.use_deterministic_algorithms(True)
print("keras", keras.__version__)
print("KERAS_BACKEND:", os.environ.get("KERAS_BACKEND"))
print("active backend:", keras.config.backend())
with open(os.path.join(os.environ["KERAS_HOME"], "keras.json")) as fh:
    print("keras.json:", json.load(fh))

keras 3.15.1
KERAS_BACKEND: torch
active backend: torch
keras.json: {'floatx': 'float32', 'epsilon': 1e-07, 'backend': 'tensorflow', 'image_data_format': 'channels_last'}

Look at the last line: the keras.json that Keras wrote the first time it was imported in that folder says tensorflow (the default value), but the active backend is torch, because the environment variable takes precedence. Exercise 1 checks it with a custom keras.json. The practical rule: the backend is set in the process environment (in the shell, in the container, in the service), or in the first lines of the program with os.environ["KERAS_BACKEND"] = "jax" before any import keras, indirect ones included (a library that imports Keras internally also counts).

keras.ops: portable computation

Keras layers are written with keras.ops, an API shaped like NumPy (ops.mean, ops.std, ops.matmul, ops.where…) plus neural network operations (ops.relu, ops.softmax, convolutions). Each function calls the backend's and returns a backend tensor. A loss or a metric written with keras.ops works on all three; one written with torch or with numpy inside the graph does not. As an example, the annualized Sharpe of the six synthetic assets in vol5_diario.parquet (lesson 5.1.3 for the Sharpe as a metric):

v = pd.read_parquet("vol5_diario.parquet")
r = v.pivot(index="fecha", columns="activo", values="rendimiento")
t = ops.convert_to_tensor(r.to_numpy("float32"))
print(type(t).__module__, type(t).__name__, tuple(t.shape))

def sharpe(x, eje=0):
    """Annualized Sharpe with keras.ops: works on any backend."""
    m, s = ops.mean(x, axis=eje), ops.std(x, axis=eje)
    return m / s * ops.sqrt(252.0)

pd.DataFrame({"keras.ops": ops.convert_to_numpy(sharpe(t)),
              "pandas": r.mean() / r.std(ddof=0) * np.sqrt(252)}
             ).round(3)

torch Tensor (4529, 6)
        keras.ops  pandas
activo
accion      0.956   0.956
bono       -0.281  -0.281
cripto      0.424   0.424
divisa      0.243   0.243
indice      0.784   0.784
oro         0.349   0.349

The tensor is a torch.Tensor of 4,529 sessions by 6 assets, and the result matches pandas in the three digits shown (the standard deviation of ops.std is the population one, hence ddof=0). Lesson 5.6.4 uses this same function as the basis of a loss.

The test bench: a HAR with a known answer

To compare engines you need a problem whose solution is known by another route. Corsi's HAR model predicts the logarithm of tomorrow's realized variance with its 1-, 5- and 22-day averages:

In the level 5 data the true process is a HAR in logarithms (data README), so linear regression is almost the best possible model. A network with a single dense layer without activation, Dense(1), is that regression: $\hat y = w^\top x + b$. Trained with squared error, its optimum is the least-squares one. So there is a known answer: OLS. The index (indice) is used, with training from 2008 to 2017 and testing from 2018 to 2025 (the same split as the README's forecast table):

d = v[v["activo"] == "indice"].set_index("fecha")
lrv = np.log(d["rv_5min"])
har = pd.DataFrame({"d1": lrv, "s5": lrv.rolling(5).mean(),
                    "m22": lrv.rolling(22).mean()}).shift(1)
har["y"] = lrv
har = har.dropna()
COLS = ["d1", "s5", "m22"]
ent = har.index < "2018-01-01"
Xe, ye = har.loc[ent, COLS].to_numpy(), har.loc[ent, "y"].to_numpy()
Xp, yp = har.loc[~ent, COLS].to_numpy(), har.loc[~ent, "y"].to_numpy()

b_ols = np.linalg.lstsq(np.c_[np.ones(len(Xe)), Xe], ye, rcond=None)[0]
p_ols = np.c_[np.ones(len(Xp)), Xp] @ b_ols
print(f"rows: training {len(Xe)}, test {len(Xp)}")
print("OLS: intercept", round(b_ols[0], 3),
      "slopes", np.round(b_ols[1:], 3))
print(f"OLS test MSE: {np.mean((p_ols - yp) ** 2):.4f}")

rows: training 2496, test 2011
OLS: intercept -0.743 slopes [0.373 0.278 0.273]
OLS test MSE: 0.1738

The slopes add up to 0.92, the persistence of the true process. The OLS test MSE, 0.1738, is the yardstick for the rest of the lesson.

Three ways to write a model

Keras offers three APIs to define a model, and all three produce the same type of object, with fit, predict, save and summary:

To check that they are the same computation, the cell builds a network with 3 inputs, 8 hidden ReLU neurons and one output in all three ways, copies the weights of the first to the other two and compares the outputs on four standardized test rows:

keras.utils.set_random_seed(0)
secuencial = keras.Sequential([
    keras.Input((3,)),
    layers.Dense(8, activation="relu"),
    layers.Dense(1)])

entrada = keras.Input((3,), name="har")
h = layers.Dense(8, activation="relu")(entrada)
funcional = keras.Model(entrada, layers.Dense(1)(h))

class MLP(keras.Model):
    def __init__(self, ocultas=8, **kwargs):
        super().__init__(**kwargs)
        self.oculta = layers.Dense(ocultas, activation="relu")
        self.salida = layers.Dense(1)

    def call(self, x, training=False):
        return self.salida(self.oculta(x))

subclase = MLP()
subclase(np.zeros((1, 3), "float32"))      # creates the weights
funcional.set_weights(secuencial.get_weights())
subclase.set_weights(secuencial.get_weights())

x = ((Xp[:4] - Xe.mean(0)) / Xe.std(0)).astype("float32")
for m in (secuencial, funcional, subclase):
    print(f"{type(m).__name__:<11}{m.count_params():>3} parameters",
          np.round(ops.convert_to_numpy(m(x)).ravel(), 4))

Sequential  41 parameters [0.4735 0.8178 1.1518 1.2612]
Functional  41 parameters [0.4735 0.8178 1.1518 1.2612]
MLP         41 parameters [0.4735 0.8178 1.1518 1.2612]

All three have 41 parameters (3×8 + 8 for the hidden layer, 8 + 1 for the output) and give the same four predictions. The subclass needs a prior call (subclase(np.zeros((1, 3)))) to create its weights, because Keras does not know the shape of the input until it sees it; the other two declare it with keras.Input.

The Functional API earns its place when the model stops being a stack. A typical case in finance: a panel model that shares the HAR slopes across the six assets and gives each one its own level. The averages come in on one side and the asset identifier on another, through an Embedding of dimension 1, which is nothing more than a learned constant per asset (the equivalent of the fixed effects of a panel regression):

rv_in = keras.Input((3,), name="har")
activo_in = keras.Input((1,), dtype="int32", name="activo")
e = layers.Flatten()(layers.Embedding(6, 1, name="nivel")(activo_in))
h = layers.Concatenate()([rv_in, e])
salida = layers.Dense(1, name="prevision")(h)
panel = keras.Model([rv_in, activo_in], salida)
for capa in panel.layers:
    print(f"{capa.name:<12}{type(capa).__name__:<13}"
          f"{str(capa.output.shape):<14}{capa.count_params()}")

activo      InputLayer   (None, 1)     0
nivel       Embedding    (None, 1, 1)  6
har         InputLayer   (None, 3)     0
flatten     Flatten      (None, 1)     0
concatenate Concatenate  (None, 4)     0
prevision   Dense        (None, 1)     5

The graph has two inputs, a nivel (level) layer with 6 parameters (one per asset) and an output with 5 (four slopes and the intercept). The listing reads like the graph: which layer produces which shape and how many parameters it has. Lesson 5.6.2 uses this API for an MLP with a linear branch.

Subclassing is the most flexible API and the easiest to break. Everything the model does is in call, and Keras cannot inspect it before running it. The detail that causes the most bugs is the training argument: layers like Dropout or GaussianNoise behave differently when training and when predicting, and the subclass has to pass it to each one. An example with Gaussian noise on the inputs, a form of regularization used with noisy financial data:

class MLPRuido(MLP):
    """Subclass: Gaussian noise on the inputs only when training."""
    def __init__(self, sigma=0.5, **kwargs):
        super().__init__(**kwargs)
        self.ruido = layers.GaussianNoise(sigma, seed=0)

    def call(self, x, training=False):
        return super().call(self.ruido(x, training=training))

mr = MLPRuido()
mr(x)
mr.set_weights(secuencial.get_weights())
print("inference:", np.round(ops.convert_to_numpy(
    mr(x, training=False)).ravel(), 4))
print("training :", np.round(ops.convert_to_numpy(
    mr(x, training=True)).ravel(), 4))

inference: [0.4735 0.8178 1.1518 1.2612]
training : [-0.104   0.215   1.2085  1.0265]

With training=False (what predict and evaluate do) the output is that of the network without noise, identical to those of the previous cell; with training=True (what fit does), each call adds noise with standard deviation 0.5 to the inputs and the first prediction goes from 0.4735 to −0.104. Keras 3 carries a call context that propagates training to the inner layers even if call does not pass it (checked with this version), but passing it explicitly makes it clear what the model does in each mode and does not depend on that mechanism when the layer is called from custom code.

The same training in PyTorch and in JAX

Training with two backends takes two processes. The cell saves the data standardized with the mean and standard deviation of the training set (standardizing with test data would be leakage, lesson 5.2.4) and writes a short script that trains Dense(1) with Adam (rate 0.01, batches of 128, 40 epochs), saves the weights and the predictions, and writes a JSON with the slopes and the validation curve. If it receives a weights file as an argument, it does not train: it loads them and predicts.

mu, sd = Xe.mean(0), Xe.std(0)
my, sy = ye.mean(), ye.std()
np.savez("har.npz",
         Xe=((Xe - mu) / sd).astype("float32"),
         ye=((ye - my) / sy).astype("float32"),
         Xp=((Xp - mu) / sd).astype("float32"),
         yp=((yp - my) / sy).astype("float32"))

GUION = '''
import json, sys
import numpy as np
import keras

keras.utils.set_random_seed(0)
z = np.load("har.npz")
m = keras.Sequential([keras.Input((3,)), keras.layers.Dense(1)])
b = keras.config.backend()
if len(sys.argv) > 1:             # weights saved with another backend
    m.load_weights(sys.argv[1])
    curva = []
else:
    m.compile(keras.optimizers.Adam(0.01), "mse")
    h = m.fit(z["Xe"], z["ye"], epochs=40, batch_size=128,
              verbose=0, validation_data=(z["Xp"], z["yp"]))
    m.save_weights(f"har_{b}.weights.h5")
    curva = h.history["val_loss"]
p = m.predict(z["Xp"], verbose=0).ravel()
np.save(f"pred_{b}{len(sys.argv)}.npy", p)
w, c = m.layers[-1].get_weights()
print(json.dumps({"w": w.ravel().tolist(), "c": float(c[0]),
                  "curva": curva}))
'''
with open("har_keras.py", "w") as fh:
    fh.write(GUION)

The next cell launches the script twice with subprocess and the same interpreter, changing only KERAS_BACKEND, and returns the slopes to the original scale (multiplying by $\sigma_y / \sigma_x$):

import subprocess
import sys

def ejecuta(backend, *args):
    entorno = {**os.environ, "KERAS_BACKEND": backend}
    out = subprocess.run([sys.executable, "har_keras.py", *args],
                         env=entorno, capture_output=True, text=True,
                         check=True).stdout
    return json.loads(out.strip().splitlines()[-1])

res = {b: ejecuta(b) for b in ("torch", "jax")}
filas = {"OLS": list(b_ols[1:]) + [np.mean((p_ols - yp) ** 2)]}
for b, r_ in res.items():
    pend = np.array(r_["w"]) * sy / sd          # to the original scale
    pred = np.load(f"pred_{b}1.npy") * sy + my
    filas[f"keras-{b}"] = list(pend) + [np.mean((pred - yp) ** 2)]
pd.DataFrame(filas, index=COLS + ["MSE prueba"]).T.round(4)

                 d1      s5     m22  MSE prueba
OLS          0.3734  0.2777  0.2728      0.1738
keras-torch  0.3738  0.2780  0.2767      0.1738
keras-jax    0.3708  0.2879  0.2633      0.1737

Both Keras versions reach the OLS test MSE to the fourth digit (0.1738 and 0.1737), but the slopes are not identical: PyTorch's stay within 0.004 of OLS on all three; JAX's drift up to 0.010 on the weekly and monthly ones. It is not a bug. The three HAR averages are highly correlated (each one contains the previous one), and the loss forms an almost flat valley in which many combinations of slopes give the same error; minibatch gradient descent stops at some point in the valley that depends on the randomness of the initialization and the shuffling. For predicting it does not matter. For reading coefficients, OLS (or lesson 5.5.1) is the tool.

The validation curve of each backend, epoch by epoch:

curvas = pd.DataFrame({b: np.array(r_["curva"]) * sy**2
                       for b, r_ in res.items()},
                      index=range(1, 41)).rename_axis("época")
curvas.iloc[[0, 1, 2, 4, 9, 19, 39]].round(4)

        torch     jax
época
1      0.2820  0.3861
2      0.2354  0.2905
3      0.2291  0.2682
5      0.2128  0.2330
10     0.1859  0.1899
20     0.1735  0.1741
40     0.1738  0.1737

The header figure plots the two full curves. They start at different places (0.282 in PyTorch and 0.386 in JAX after the first epoch) because they start from different initial weights with the same seed, and they converge on the OLS line around epoch 20. From there on, both oscillate in the fourth decimal place.

Portable weights across backends

If the training runs differ because of randomness, the proof that the two backends compute the same thing is another: the same weights have to give the same predictions. save_weights writes a .weights.h5 file (HDF5) with the weights of each layer, independent of the backend. The cell loads in a JAX process the weights saved by PyTorch:

cruzado = ejecuta("jax", "har_torch.weights.h5")
p_torch = np.load("pred_torch1.npy")
p_cruz = np.load("pred_jax2.npy")
print("same weights:", np.allclose(cruzado["w"], res["torch"]["w"]))
print("equal predictions (tolerance 1e-6):",
      np.allclose(p_torch, p_cruz, rtol=0, atol=1e-6))

same weights: True
equal predictions (tolerance 1e-6): True

The predictions match within a tolerance of $10^{-6}$ over the 2,011 test sessions. The remaining differences, on the order of $10^{-7}$, are float32 rounding: each backend sums in its own order. A model trained in JAX on a workstation can be served with PyTorch inside another application, or the other way around; lesson 5.6.4 does it with the full .keras format.

Determinism on each backend

A result that changes when you rerun it cannot be reviewed or debugged. On CPU, with fixed threads, it is enough to fix every source of randomness; on GPU you also need deterministic algorithms, which are sometimes slower.

In addition, the number of BLAS and OpenMP threads has to be fixed (the course runner sets 4): splitting a sum across a different number of threads changes the order and, with it, the last decimals. And outputs are rounded to what the conclusion needs, because another CPU can give different far decimals with the same code. The check in this process (PyTorch):

z = np.load("har.npz")

def entrena(semilla, epocas=5):
    if semilla is not None:
        keras.utils.set_random_seed(semilla)
    m = keras.Sequential([keras.Input((3,)), layers.Dense(1)])
    m.compile(keras.optimizers.Adam(0.01), "mse")
    m.fit(z["Xe"], z["ye"], epochs=epocas, batch_size=128, verbose=0)
    return m.layers[-1].get_weights()[0].ravel()

a, b = entrena(0), entrena(0)
c = entrena(None)
print("seed 0 twice, identical   :", np.array_equal(a, b))
print("without resetting the seed:", np.array_equal(a, c))

seed 0 twice, identical   : True
without resetting the seed: False

With the seed set before each training run, two fits give bit-for-bit identical weights. Without resetting it, the generator continues where it left off and the second model starts with other weights and another batch order. That is why every experiment in the lessons sets the seed right before building its model.

What is native to each backend

Portability does not prevent you from dropping down to the engine. With the PyTorch backend, a Keras model is a torch.nn.Module: its weights are torch.nn.Parameter, autograd computes its gradients and any PyTorch optimizer or loop can train it. The cell computes the gradient of the squared error of a batch of 256 rows with autograd and compares it with the formula by hand, $\nabla_w = \tfrac{2}{n} X^\top (Xw + b - y)$:

m = keras.Sequential([keras.Input((3,)), layers.Dense(1)])
m.load_weights("har_torch.weights.h5")
print("is torch.nn.Module:", isinstance(m, torch.nn.Module))
print([tuple(p.shape) for p in m.parameters()])

xb = torch.tensor(z["Xe"][:256])
yb = torch.tensor(z["ye"][:256])
perdida = torch.mean((m(xb).squeeze(-1) - yb) ** 2)
perdida.backward()
capa = m.layers[-1]
g_torch = np.r_[capa.kernel.value.grad.numpy().ravel(),
                capa.bias.value.grad.numpy()]

w, c = capa.get_weights()
resid = z["Xe"][:256] @ w.ravel() + c[0] - z["ye"][:256]
g_mano = 2 * np.r_[z["Xe"][:256].T @ resid, resid.sum()] / 256
print("autograd gradient:", np.round(g_torch, 5))
print("manual gradient  :", np.round(g_mano, 5))

is torch.nn.Module: True
[(1,), (3, 1)]
autograd gradient: [0.06436 0.02329 0.05671 0.07229]
manual gradient  : [0.06436 0.02329 0.05671 0.07229]

parameters() returns the intercept (shape (1,)) and the weight matrix (3, 1), and the autograd gradient matches the formula's in the five digits shown, in the order daily, weekly and monthly slope and intercept. Lesson 5.7.1 develops the PyTorch training loop; a Keras model can be put into it.

With TensorFlow, the code is the same changing the variable; with the stable version of TensorFlow that cannot be installed here, it would look like this:

import os
os.environ["KERAS_BACKEND"] = "tensorflow"   # before importing keras
import keras
import tensorflow as tf

keras.utils.set_random_seed(0)
tf.config.experimental.enable_op_determinism()
print(keras.config.backend())   # tensorflow

With JAX, the equivalent of dropping down to the engine is the stateless function. JAX works with pure functions: the weights do not live inside the model, they are passed as an argument. stateless_call evaluates the model with given weights and returns the output and the new state, so jax.grad and jax.jit apply directly (the Keras guide “Writing a training loop from scratch in JAX” develops it). This block needs a process with KERAS_BACKEND=jax, which is why it is not run inside this lesson:

import os
os.environ["KERAS_BACKEND"] = "jax"
import jax
import keras

m = keras.Sequential([keras.Input((3,)), keras.layers.Dense(1)])
entrenables = [v.value for v in m.trainable_variables]
no_entrenables = [v.value for v in m.non_trainable_variables]

def perdida(entrenables, x, y):
    pred, _ = m.stateless_call(entrenables, no_entrenables, x)
    return jax.numpy.mean((pred[:, 0] - y) ** 2)

grad = jax.jit(jax.grad(perdida))(entrenables, x, y)

Which backend to choose

There is no universal answer, but there are criteria for a quantitative research team:

flowchart TB
  A[Small model on CPU] --> J[JAX]
  B[Integrate with PyTorch] --> T[PyTorch]
  C[Serve on TF Serving] --> F[TensorFlow]
  J --> K[Same Keras code]
  T --> K
  F --> K

• PyTorch if the rest of the system is already PyTorch (the foundation models of module 5.8 are) or if you need to mix Keras layers with PyTorch modules.
• JAX to train many small models on CPU, as in a hyperparameter search or a walk-forward with retraining (lesson 5.4.3): XLA compilation amortizes the fixed cost of each step.
• TensorFlow if the deployment uses its ecosystem (TF Serving, LiteRT). On Python 3.14 there is no stable version yet.

In all three cases, the model and the weights can be moved. What does not move is the randomness of training: a result that only appears with one backend, or with one seed, is noise until proven otherwise (lesson 5.6.2).

Exercises

1. Check the configuration's order of precedence: in a child process without KERAS_BACKEND and with its own KERAS_HOME whose keras.json says "backend": "jax", which backend starts? And if KERAS_BACKEND=torch is also set?

import tempfile

casa = tempfile.mkdtemp()
with open(os.path.join(casa, "keras.json"), "w") as fh:
    json.dump({"backend": "jax", "floatx": "float32",
               "epsilon": 1e-07, "image_data_format": "channels_last"},
              fh)
orden = "import keras; print(keras.config.backend())"
base = {k: v for k, v in os.environ.items() if k != "KERAS_BACKEND"}
for extra in ({}, {"KERAS_BACKEND": "torch"}):
    entorno = {**base, "KERAS_HOME": casa, **extra}
    out = subprocess.run([sys.executable, "-c", orden], env=entorno,
                         capture_output=True, text=True, check=True)
    print(extra or "keras.json only", "->", out.stdout.strip())

keras.json only -> jax
{'KERAS_BACKEND': 'torch'} -> torch

2. Add to the HAR the negative part of yesterday's return, $\ell_{t-1} = \min(r_{t-1}, 0)$, as a second input of a functional model (the three averages on one side and $\ell$ on the other, joined with Concatenate). Compare its coefficients with those of the OLS with the four regressors.

neg = np.minimum(d["rendimiento"], 0).shift(1).reindex(har.index)
L = neg.to_numpy()
Le, Lp = L[ent], L[~ent]
b4 = np.linalg.lstsq(np.c_[np.ones(len(Xe)), Xe, Le], ye,
                     rcond=None)[0]
ml, sl = Le.mean(), Le.std()

keras.utils.set_random_seed(0)
e1 = keras.Input((3,), name="har")
e2 = keras.Input((1,), name="negativo")
s = layers.Dense(1)(layers.Concatenate()([e1, e2]))
har_l = keras.Model([e1, e2], s)
har_l.compile(keras.optimizers.Adam(0.01), "mse")
har_l.fit([z["Xe"], ((Le - ml) / sl)[:, None].astype("float32")],
          z["ye"], epochs=40, batch_size=128, verbose=0)
w = har_l.layers[-1].get_weights()[0].ravel()
pd.DataFrame({"OLS": b4[1:],
              "keras": w * sy / np.r_[sd, sl]},
             index=COLS + ["negativo"]).round(3)

             OLS   keras
d1         0.267   0.258
s5         0.344   0.341
m22        0.278   0.270
negativo -18.385 -19.045

3. A colleague trains the same model with keras.utils.set_random_seed(0) in PyTorch and in JAX and is worried that the final weights do not match. Is it a bug? What would you check to know whether the two backends compute the same thing?