---
title: "5.6.1 Keras 3 multi-backend: PyTorch, JAX and TensorFlow"
description: "Keras 3 on PyTorch and JAX: backend, keras.ops, Sequential, Functional and subclassing APIs, and a HAR trained on both with OLS as the reference."
url: "https://pythonparatrading.com/en/python-trading-course/level-5/keras-3-multi-backend/"
lang: "en"
image: "https://pythonparatrading.com/og/en/python-trading-course/level-5/keras-3-multi-backend.png"
---
# 5.6.1 · Keras 3: one model, several compute engines
> Keras 3 separates the definition of the model from the engine that does the math: the same code runs on PyTorch, JAX or TensorFlow. This lesson configures the backend, writes a model with the three Keras APIs and trains a HAR volatility model with PyTorch and with JAX to see what matches between engines, what does not and why.
- Course: [Level 5: Machine learning and deep learning (Keras/PyTorch)](https://pythonparatrading.com/en/python-trading-course/level-5/)
- Module: 5.6 · Deep learning with Keras 3
- Duration: 60 min
- Prerequisites: [5.1.2 · Supervised, unsupervised and reinforcement learning; scikit-learn pipelines](https://pythonparatrading.com/en/python-trading-course/level-5/supervised-learning-scikit-learn-pipelines/), [5.5.1 · Regularized regressions: Ridge, Lasso, Elastic Net and logistic](https://pythonparatrading.com/en/python-trading-course/level-5/ridge-lasso-elastic-net-logistic-regression/)
## By the end of this lesson
- Choose and check the Keras 3 backend (`KERAS_BACKEND`, `keras.json`) and know why it has to be set before the first `import keras`.
- Write portable computations with `keras.ops` and build the same model with the Sequential, Functional and subclassing APIs.
- Train a HAR volatility model as a dense layer with PyTorch and with JAX, check it against the least-squares solution and move the weights from one backend to another.
- Make each backend deterministic and use a Keras model with PyTorch's native tools.
## 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.
```mermaid
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:
```python
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))
```
```output
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):
```python
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)
```
```output
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:
$$
\log RV_{t} = c + \beta_d \log RV_{t-1} + \beta_s \overline{\log RV}_{t-5:t-1} + \beta_m \overline{\log RV}_{t-22:t-1} + \varepsilon_t
$$
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):
```python
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}")
```
```output
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`:
| API | How it is written | What it is for |
|---|---|---|
| Sequential | list of layers in `keras.Sequential` | simple stacks: one input, one output, no branches |
| Functional | each layer is called on symbolic tensors and `keras.Model(inputs, outputs)` | graphs: several inputs or outputs, branches, shared layers |
| Subclassing | child class of `keras.Model` with the layers in `__init__` and the computation in `call` | custom logic in the forward pass (loops, conditions, noise) |
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:
```python
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))
```
```output
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):
```python
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()}")
```
```output
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:
```python
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))
```
```output
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.
```python
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$):
```python
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)
```
```output
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
```
*Fig. — Index HAR slopes (1, 5 and 22-day means of log RV) by least squares and trained with Keras 3 on PyTorch and on JAX.* Three rows with three marks each. Daily: OLS 0.373 (highlighted), PyTorch 0.374, JAX 0.371. Weekly: 0.278, 0.278 and 0.288. Monthly: 0.273, 0.277 and 0.263. Both Keras versions stay within 0.011 of the OLS.
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:
```python
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)
```
```output
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:
```python
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))
```
```output
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.
| Backend | Seeds | Deterministic algorithms | Notes |
|---|---|---|---|
| All | `keras.utils.set_random_seed(s)`: Python, NumPy and the backend generator | — | set before creating the model |
| PyTorch | the above (calls `torch.manual_seed`) | `torch.use_deterministic_algorithms(True)` | an operation without a deterministic version raises `RuntimeError` |
| JAX | the above (Keras carries its own `SeedGenerator`) | on CPU, runs are identical with no extra setting | in pure JAX code, explicit keys (`jax.random.key`) |
| TensorFlow | the above | `tf.config.experimental.enable_op_determinism()` | also makes `tf.data` operations deterministic |
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):
```python
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))
```
```output
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)$:
```python
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))
```
```output
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:
```python skip
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:
```python skip
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:
```mermaid
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?
**Solution.**
```python
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())
```
```output
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.
**Solution.**
```python
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)
```
```output
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?
**Solution.**
It is not a bug. The Keras seed sets each backend's generator, but each backend has its own generator and its own
way of consuming it: the weight initialization and the batch shuffling come out different, and the optimizer travels
another path toward the same minimum. What must match is the computation with the same weights: load the weights of
one backend into the other (`save_weights` / `load_weights`) and compare the predictions with a floating-point
tolerance, as in the lesson (equal to $10^{-6}$), and compare the gradients of a fixed batch. If that checks out, the
differences in the trained weights are optimizer noise and should be the size of those between two seeds on the same
backend.
## Self-assessment
1. A script runs `import keras` and on the next line `os.environ["KERAS_BACKEND"] = "jax"`. Which backend does it train with?
- A. With whichever was configured at import time (the previous environment variable, `keras.json` or TensorFlow by default): the backend is set once, at the first `import keras`
- B. With JAX: the variable is read when the model is compiled
- C. With both at once
- D. It raises an error when the variable changes
*Answer A: With whichever was configured at import time (the previous environment variable, `keras.json` or TensorFlow by default): the backend is set once, at the first `import keras`.* Keras reads `KERAS_BACKEND` and `keras.json` when it is imported and does not allow switching backends afterward. That is why the variable is set in the process environment or in the first lines, before any `import keras` (indirect ones included).
2. The same script, with `set_random_seed(0)`, gives a daily slope of 0.3738 in PyTorch and 0.3708 in JAX. What does it indicate?
- A. That one of the two backends computes the dense layer incorrectly
- B. That each backend has its own random generator: the initialization and the shuffling of the batches differ and the optimizer follows another path; with the same weights, the predictions match to the millionth
- C. That the seed does not work in JAX
- D. That training has to be done in float64
*Answer B: That each backend has its own random generator: the initialization and the shuffling of the batches differ and the optimizer follows another path; with the same weights, the predictions match to the millionth.* The lesson loads in JAX the weights saved with PyTorch and gets the same predictions within floating-point tolerance. The training differences come from each backend's randomness, not from the arithmetic.
3. When is the Functional API worth it over `Sequential`?
- A. Never: they are equivalent
- B. Only with the JAX backend
- C. When the model has more than ten layers
- D. When the model has several inputs or outputs, branches or shared layers: the Functional API describes a graph, the Sequential one only a stack
*Answer D: When the model has several inputs or outputs, branches or shared layers: the Functional API describes a graph, the Sequential one only a stack.* `Sequential` chains layers one after another. A model with the HAR averages through one input and the asset identifier through another (with an `Embedding`) needs a graph, and the Functional API builds it and leaves it inspectable.
4. In a subclass of `keras.Model` with `GaussianNoise` or `Dropout`, what is the `training` argument of `call` for?
- A. To choose the backend
- B. To enable gradient computation in JAX
- C. So that those layers act only when training: at inference the output has to be deterministic, and `predict` and `fit` pass the right value
- D. It is useless in Keras 3
*Answer C: So that those layers act only when training: at inference the output has to be deterministic, and `predict` and `fit` pass the right value.* With `training=True` the noise layer changes the output (in the lesson, from 0.4735 to -0.104 in the first row); with `training=False`, it does not. A subclass that does not propagate `training` to its layers puts noise into the predictions.
5. The dense layer trained with Adam reaches the same test MSE as OLS (0.1738) with slopes that differ in the third decimal place. Why?
- A. Because Adam is a better estimator than OLS
- B. Because the test is too short
- C. Because Keras has no intercept term
- D. Because the three HAR averages are highly correlated: there are many combinations of slopes with almost the same error, and minibatch noise leaves the result somewhere in that flat valley
*Answer D: Because the three HAR averages are highly correlated: there are many combinations of slopes with almost the same error, and minibatch noise leaves the result somewhere in that flat valley.* OLS gives the exact minimum; minibatch gradient descent stays close, within a region where the loss barely changes. For predicting it does not matter; for interpreting coefficients, it does.
## Key points
- The Keras 3 backend is chosen before the first `import keras`: `KERAS_BACKEND` overrides `keras.json`, and with neither of them Keras tries to use TensorFlow.
- `keras.ops` is a NumPy-style API that returns backend tensors (here `torch.Tensor`); a function written with it works in PyTorch, JAX and TensorFlow.
- Sequential for simple stacks, Functional for graphs with several inputs or outputs, subclassing for custom logic in `call` (respecting `training`). With the same weights, all three give the same output.
- A dense layer trained with Adam recovers the HAR OLS (test MSE 0.1738 on both backends); each backend has its own randomness, but weights saved in one predict the same in the other.
- Determinism: `keras.utils.set_random_seed` on every backend, plus `torch.use_deterministic_algorithms(True)` in PyTorch and `enable_op_determinism()` in TensorFlow; fixed threads and rounded outputs.
- With PyTorch, a Keras model is a `torch.nn.Module`: PyTorch's autograd, optimizers and loops work with it (lesson 5.7.1).
## Resources
### Documentation
- [Keras: Getting started (installation and backend configuration)](https://keras.io/getting_started/) · keras.io
- [Keras: The Functional API](https://keras.io/guides/functional_api/) · keras.io
- [Keras: Making new layers and models via subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) · keras.io
- [Keras: Reproducibility in Keras Models](https://keras.io/examples/keras_recipes/reproducibility_recipes/) · keras.io
- [PyTorch: Reproducibility (notes on randomness and deterministic algorithms)](https://docs.pytorch.org/docs/stable/notes/randomness.html) · docs.pytorch.org
- [PyPI: tensorflow 2.22.0rc0 (the first version with wheels for Python 3.14)](https://pypi.org/project/tensorflow/2.22.0rc0/) · pypi.org
### Data
- [Level 5 data: returns and realized variance of six assets (Parquet)](https://pythonparatrading.com/curso/datos/nivel-5/vol5_diario.parquet) · Download · vol5_diario.parquet
Previous lesson: [5.5.6 · Learning to rank for stock selection](https://pythonparatrading.com/en/python-trading-course/level-5/learning-to-rank-stock-selection/)
Next lesson: [5.6.2 · An MLP on a stock panel: regularize and compare](https://pythonparatrading.com/en/python-trading-course/level-5/mlp-for-tabular-financial-data/)