Keras 3: un modelo, varios motores de cálculo.

Una API, varios motores

Keras nació como una capa sobre otros motores de cálculo (Theano, después TensorFlow) y durante años fue, en la práctica, la API de alto nivel de TensorFlow. Keras 3, publicado en noviembre de 2023, volvió al diseño original: la definición del modelo (capas, pérdidas, optimizadores, fit) es independiente del motor que hace las cuentas, el backend. La documentación oficial enumera tres para entrenar: JAX, TensorFlow y PyTorch. El código del paquete trae además un backend de NumPy y otro de OpenVINO, pensados para inferencia.

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

Para un equipo de research esto tiene tres consecuencias prácticas:

• El modelo no ata al motor. Un modelo escrito en Keras puede entrenarse en JAX (rápido para modelos pequeños, como se verá) y servirse dentro de una aplicación de PyTorch, o al revés, con los mismos pesos.
• Los tensores son los del backend. Con PyTorch, cada tensor es un torch.Tensor y el modelo es un torch.nn.Module; con JAX, un jax.Array. Se puede bajar al motor cuando hace falta.
• El azar no es portable. Cada backend tiene su generador de números aleatorios: la misma semilla no da los mismos pesos iniciales en dos backends. Hay que saberlo para no confundir ruido con un error.

En el entorno del curso están PyTorch 2.14.1 (solo CPU) y JAX 0.12.0.1, con Keras 3.15.1. TensorFlow no: la última versión estable, la 2.21.0, no publica ruedas para Python 3.14, y la primera que las tiene, la 2.22.0rc0, es una versión candidata. Su código aparece al final, sin ejecutar.

Configurar el backend

Keras decide el backend una sola vez, al importarse, y no permite cambiarlo después. Lo busca en este orden: la variable de entorno KERAS_BACKEND y, si no está, el campo backend de keras.json (en ~/.keras/ o en la carpeta de KERAS_HOME); sin ninguno de los dos, usa TensorFlow. El ejecutor de las lecciones del nivel 5 define KERAS_BACKEND=torch y apunta KERAS_HOME a una carpeta temporal:

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("backend activo:", 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
backend activo: torch
keras.json: {'floatx': 'float32', 'epsilon': 1e-07, 'backend': 'tensorflow', 'image_data_format': 'channels_last'}

Fíjate en la última línea: el keras.json que Keras escribió al importarse por primera vez en esa carpeta dice tensorflow (el valor por defecto), pero el backend activo es torch, porque la variable de entorno manda. El ejercicio 1 lo comprueba con un keras.json propio. La regla práctica: el backend se fija en el entorno del proceso (en la shell, en el contenedor, en el servicio), o en las primeras líneas del programa con os.environ["KERAS_BACKEND"] = "jax" antes de cualquier import keras, incluidos los indirectos (una librería que importe Keras por dentro también cuenta).

keras.ops: cálculo portable

Las capas de Keras están escritas con keras.ops, una API con la forma de NumPy (ops.mean, ops.std, ops.matmul, ops.where…) más las operaciones de redes (ops.relu, ops.softmax, convoluciones). Cada función llama a la del backend y devuelve un tensor del backend. Una pérdida o una métrica escrita con keras.ops vale en los tres; una escrita con torch o con numpy dentro del grafo, no. Como ejemplo, el Sharpe anualizado de los seis activos sintéticos de vol5_diario.parquet (lección 5.1.3 para el Sharpe como métrica):

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):
    """Sharpe anualizado con keras.ops: vale en cualquier 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

El tensor es un torch.Tensor de 4.529 sesiones por 6 activos, y el resultado coincide con pandas en las tres cifras mostradas (la desviación típica de ops.std es la poblacional, de ahí ddof=0). La lección 5.6.4 usa esta misma función como base de una pérdida.

El banco de pruebas: un HAR con respuesta conocida

Para comparar motores hace falta un problema cuya solución se conozca por otra vía. El modelo HAR de Corsi predice el logaritmo de la varianza realizada de mañana con sus medias de 1, 5 y 22 días:

En los datos del nivel 5 el proceso verdadero es un HAR en logaritmos (README de los datos), así que la regresión lineal es casi el mejor modelo posible. Una red de una sola capa densa sin activación, Dense(1), es esa regresión: $\hat y = w^\top x + b$. Entrenada con el error cuadrático, su óptimo es el de mínimos cuadrados. Así que hay respuesta conocida: la OLS. Se usa el índice, con entrenamiento de 2008 a 2017 y prueba de 2018 a 2025 (el mismo corte que la tabla de previsiones del README):

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"filas: entreno {len(Xe)}, prueba {len(Xp)}")
print("OLS: constante", round(b_ols[0], 3),
      "pendientes", np.round(b_ols[1:], 3))
print(f"MSE de prueba de la OLS: {np.mean((p_ols - yp) ** 2):.4f}")

filas: entreno 2496, prueba 2011
OLS: constante -0.743 pendientes [0.373 0.278 0.273]
MSE de prueba de la OLS: 0.1738

Las pendientes suman 0,92, la persistencia del proceso verdadero. El MSE de prueba de la OLS, 0,1738, es la vara de medir del resto de la lección.

Tres formas de escribir un modelo

Keras ofrece tres APIs para definir un modelo, y las tres producen el mismo tipo de objeto, con fit, predict, save y summary:

Para comprobar que son el mismo cálculo, la celda construye una red de 3 entradas, 8 neuronas ocultas con ReLU y una salida de las tres formas, copia los pesos de la primera a las otras dos y compara las salidas sobre cuatro filas de prueba tipificadas:

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"))      # crea los pesos
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} parámetros",
          np.round(ops.convert_to_numpy(m(x)).ravel(), 4))

Sequential  41 parámetros [0.4735 0.8178 1.1518 1.2612]
Functional  41 parámetros [0.4735 0.8178 1.1518 1.2612]
MLP         41 parámetros [0.4735 0.8178 1.1518 1.2612]

Las tres tienen 41 parámetros (3×8 + 8 de la capa oculta, 8 + 1 de la salida) y dan las mismas cuatro predicciones. La subclase necesita una llamada previa (subclase(np.zeros((1, 3)))) para crear sus pesos, porque Keras no conoce la forma de la entrada hasta que la ve; las otras dos la declaran con keras.Input.

La API funcional se gana el sitio cuando el modelo deja de ser una pila. Un caso típico en finanzas: un modelo de panel que comparte las pendientes del HAR entre los seis activos y da a cada uno su propio nivel. Las medias entran por un lado y el identificador del activo por otro, a través de un Embedding de dimensión 1, que no es más que una constante por activo aprendida (el equivalente de los efectos fijos de una regresión de panel):

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

El grafo tiene dos entradas, una capa nivel con 6 parámetros (uno por activo) y una salida con 5 (cuatro pendientes y la constante). El listado se lee como el grafo: qué capa produce qué forma y cuántos parámetros tiene. La lección 5.6.2 usa esta API para un MLP con una rama lineal.

La subclase es la API más flexible y la más fácil de romper. Todo lo que hace el modelo está en call, y Keras no puede inspeccionarlo antes de ejecutarlo. El detalle que más errores produce es el argumento training: capas como Dropout o GaussianNoise se comportan distinto al entrenar y al predecir, y la subclase tiene que pasarlo a cada una. Un ejemplo con ruido gaussiano en las entradas, una forma de regularización que se usa con datos financieros ruidosos:

class MLPRuido(MLP):
    """Subclase: ruido gaussiano en las entradas solo al entrenar."""
    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("inferencia :", np.round(ops.convert_to_numpy(
    mr(x, training=False)).ravel(), 4))
print("entrenando :", np.round(ops.convert_to_numpy(
    mr(x, training=True)).ravel(), 4))

inferencia : [0.4735 0.8178 1.1518 1.2612]
entrenando : [-0.104   0.215   1.2085  1.0265]

Con training=False (lo que hacen predict y evaluate) la salida es la de la red sin ruido, idéntica a las de la celda anterior; con training=True (lo que hace fit), cada llamada añade ruido de desviación 0,5 a las entradas y la primera predicción pasa de 0,4735 a −0,104. Keras 3 lleva un contexto de llamada que propaga training a las capas internas aunque call no lo pase (comprobado con esta versión), pero pasarlo de forma explícita deja claro qué hace el modelo en cada modo y no depende de ese mecanismo cuando la capa se llama desde código propio.

El mismo entrenamiento en PyTorch y en JAX

Para entrenar con dos backends hacen falta dos procesos. La celda guarda los datos tipificados con la media y la desviación del entrenamiento (la tipificación con datos de prueba sería una fuga, lección 5.2.4) y escribe un guion corto que entrena Dense(1) con Adam (tasa 0,01, lotes de 128, 40 épocas), guarda los pesos y las predicciones, y escribe un JSON con las pendientes y la curva de validación. Si recibe un fichero de pesos como argumento, no entrena: los carga y predice.

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:             # pesos guardados con otro 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)

La siguiente celda lanza el guion dos veces con subprocess y el mismo intérprete, cambiando solo KERAS_BACKEND, y devuelve las pendientes a la escala original (multiplicando por $\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          # a la escala original
    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

Las dos versiones de Keras llegan al MSE de prueba de la OLS a la cuarta cifra (0,1738 y 0,1737), pero las pendientes no son idénticas: la de PyTorch queda a menos de 0,004 de la OLS en las tres; la de JAX se separa hasta 0,010 en la semanal y la mensual. No es un error. Las tres medias del HAR están muy correlacionadas (cada una contiene a la anterior), y la pérdida forma un valle casi plano en el que muchas combinaciones de pendientes dan el mismo error; el descenso por gradiente con minilotes se para en algún punto del valle que depende del azar de la inicialización y del barajado. Para predecir da igual. Para leer coeficientes, la OLS (o la lección 5.5.1) es la herramienta.

La curva de validación de cada backend, época a época:

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

La figura de la cabecera dibuja las dos curvas completas. Empiezan en sitios distintos (0,282 en PyTorch y 0,386 en JAX tras la primera época) porque parten de pesos iniciales distintos con la misma semilla, y se juntan sobre la línea de la OLS hacia la época 20. A partir de ahí, las dos oscilan en la cuarta cifra decimal.

Pesos portables entre backends

Si los entrenamientos difieren por el azar, la prueba de que los dos backends calculan lo mismo es otra: los mismos pesos tienen que dar las mismas predicciones. save_weights escribe un fichero .weights.h5 (HDF5) con los pesos de cada capa, independiente del backend. La celda carga en un proceso de JAX los pesos guardados por PyTorch:

cruzado = ejecuta("jax", "har_torch.weights.h5")
p_torch = np.load("pred_torch1.npy")
p_cruz = np.load("pred_jax2.npy")
print("mismos pesos:", np.allclose(cruzado["w"], res["torch"]["w"]))
print("predicciones iguales (tolerancia 1e-6):",
      np.allclose(p_torch, p_cruz, rtol=0, atol=1e-6))

mismos pesos: True
predicciones iguales (tolerancia 1e-6): True

Las predicciones coinciden dentro de una tolerancia de $10^{-6}$ en las 2.011 sesiones de prueba. Las diferencias que quedan, del orden de $10^{-7}$, son de redondeo en float32: cada backend suma en su propio orden. Un modelo entrenado en JAX en una estación de trabajo puede servirse con PyTorch dentro de otra aplicación, o al revés; la lección 5.6.4 lo hace con el formato .keras completo.

Determinismo en cada backend

Un resultado que cambia al volver a ejecutar no se puede revisar ni depurar. En CPU, con hilos fijos, basta con fijar todas las fuentes de azar; en GPU hacen falta además algoritmos deterministas, que a veces son más lentos.

Además, el número de hilos de BLAS y OpenMP tiene que ser fijo (el ejecutor del curso pone 4): repartir una suma entre otro número de hilos cambia el orden y, con él, los últimos decimales. Y las salidas se redondean a lo que la conclusión necesita, porque otra CPU puede dar decimales lejanos distintos con el mismo código. La comprobación en este proceso (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("semilla 0 dos veces, idénticos:", np.array_equal(a, b))
print("sin volver a fijar la semilla :", np.array_equal(a, c))

semilla 0 dos veces, idénticos: True
sin volver a fijar la semilla : False

Con la semilla fijada antes de cada entrenamiento, dos ajustes dan pesos idénticos bit a bit. Sin volver a fijarla, el generador sigue donde se quedó y el segundo modelo empieza con otros pesos y otro orden de lotes. Por eso cada experimento de las lecciones fija la semilla justo antes de construir su modelo.

Lo nativo de cada backend

La portabilidad no impide bajar al motor. Con el backend de PyTorch, un modelo de Keras es un torch.nn.Module: sus pesos son torch.nn.Parameter, autograd calcula sus gradientes y cualquier optimizador o bucle de PyTorch puede entrenarlo. La celda calcula el gradiente del error cuadrático de un lote de 256 filas con autograd y lo compara con la fórmula a mano, $\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("es 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("gradiente autograd:", np.round(g_torch, 5))
print("gradiente a mano  :", np.round(g_mano, 5))

es torch.nn.Module: True
[(1,), (3, 1)]
gradiente autograd: [0.06436 0.02329 0.05671 0.07229]
gradiente a mano  : [0.06436 0.02329 0.05671 0.07229]

parameters() devuelve la constante (forma (1,)) y la matriz de pesos (3, 1), y el gradiente de autograd coincide con el de la fórmula en las cinco cifras mostradas, en el orden pendiente diaria, semanal, mensual y constante. La lección 5.7.1 desarrolla el bucle de entrenamiento de PyTorch; un modelo de Keras se puede meter en él.

Con TensorFlow, el código es el mismo cambiando la variable; con la versión estable de TensorFlow que no se puede instalar aquí, quedaría así:

import os
os.environ["KERAS_BACKEND"] = "tensorflow"   # antes de importar keras
import keras
import tensorflow as tf

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

Con JAX, el equivalente de bajar al motor es la función sin estado. JAX trabaja con funciones puras: los pesos no viven dentro del modelo, se pasan como argumento. stateless_call evalúa el modelo con unos pesos dados y devuelve la salida y el estado nuevo, así que jax.grad y jax.jit se aplican directamente (la guía de Keras «Writing a training loop from scratch in JAX» lo desarrolla). Este bloque necesita un proceso con KERAS_BACKEND=jax, por eso no se ejecuta dentro de esta lección:

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)

Qué backend elegir

No hay una respuesta universal, pero sí criterios para un equipo de research cuantitativo:

flowchart TB
  A[Modelo pequeño en CPU] --> J[JAX]
  B[Integrar con PyTorch] --> T[PyTorch]
  C[Servir en TF Serving] --> F[TensorFlow]
  J --> K[Mismo código Keras]
  T --> K
  F --> K

• PyTorch si el resto del sistema ya es PyTorch (los modelos fundacionales del módulo 5.8 lo son) o si hace falta mezclar capas de Keras con módulos de PyTorch.
• JAX para entrenar muchos modelos pequeños en CPU, como en una búsqueda de hiperparámetros o un walk-forward con reentrenamiento (lección 5.4.3): la compilación con XLA amortiza el coste fijo de cada paso.
• TensorFlow si el despliegue usa su ecosistema (TF Serving, LiteRT). En Python 3.14 aún no hay versión estable.

En los tres casos, el modelo y los pesos se pueden mover. Lo que no se mueve es el azar del entrenamiento: un resultado que solo aparece con un backend, o con una semilla, es ruido hasta que se demuestre lo contrario (lección 5.6.2).

Ejercicios

1. Comprueba el orden de prioridad de la configuración: en un proceso hijo sin KERAS_BACKEND y con un KERAS_HOME propio cuyo keras.json diga "backend": "jax", ¿qué backend arranca? ¿Y si además se define KERAS_BACKEND=torch?

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 "solo keras.json", "->", out.stdout.strip())

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

2. Añade al HAR la parte negativa del rendimiento de ayer, $\ell_{t-1} = \min(r_{t-1}, 0)$, como segunda entrada de un modelo funcional (las tres medias por un lado y $\ell$ por otro, unidas con Concatenate). Compara sus coeficientes con los de la OLS con los cuatro regresores.

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. Un compañero entrena el mismo modelo con keras.utils.set_random_seed(0) en PyTorch y en JAX y le preocupa que los pesos finales no coincidan. ¿Es un error? ¿Qué comprobarías para saber si los dos backends calculan lo mismo?