TopolearnIt is how fit() works under JAX and what you need to combine Keras models with custom JAX code.
JAX functions must be pure: no hidden mutable state. Keras variables are mutable, so on the JAX backend Keras provides stateless versions of everything that touches state:
model.stateless_call(trainable_values, non_trainable_values, x, training=...)runs the forward pass using the arrays you pass in and returns(outputs, new_non_trainable_values), for example updated BatchNorm statistics.optimizer.stateless_apply(optimizer_values, grads, trainable_values)returns new trainable values and new optimizer values.- Metrics have
stateless_update_stateandstateless_result.
A training step is then a pure function from (state, batch) to (loss, new_state), which you can wrap in jax.jit, differentiate with jax.value_and_grad(..., has_aux=True), or shard with JAX's parallelism tools. The state is just lists of arrays gathered from model.trainable_variables, model.non_trainable_variables and optimizer.variables.
After the loop, write the final values back with variable.assign(value) so the Keras model (and model.save) reflect the training.
This is exactly what fit() does internally on JAX, and it is the pattern to use when integrating Keras layers into Flax-style or research JAX code. Remember to optimizer.build(model.trainable_variables) before reading its variables.
Example
import os
os.environ["KERAS_BACKEND"] = "jax"
import jax, keras, numpy as np
model = keras.Sequential([keras.Input((4,)), keras.layers.Dense(1)])
optimizer = keras.optimizers.SGD(learning_rate=0.1)
optimizer.build(model.trainable_variables)
loss_fn = keras.losses.MeanSquaredError()
def compute_loss(tv, ntv, x, y):
y_pred, ntv = model.stateless_call(tv, ntv, x, training=True)
return loss_fn(y, y_pred), ntv
@jax.jit
def train_step(state, x, y):
tv, ntv, ov = state
(loss, ntv), grads = jax.value_and_grad(compute_loss, has_aux=True)(tv, ntv, x, y)
tv, ov = optimizer.stateless_apply(ov, grads, tv)
return loss, (tv, ntv, ov)
state = ([v.value for v in model.trainable_variables],
[v.value for v in model.non_trainable_variables],
[v.value for v in optimizer.variables])
x = np.random.rand(256, 4).astype("float32"); y = x.sum(axis=1, keepdims=True)
for _ in range(300):
loss, state = train_step(state, x, y)
for var, val in zip(model.trainable_variables, state[0]):
var.assign(val) # write results back
print(float(loss) < 0.01) # TrueRemember
- JAX needs pure functions; Keras offers stateless_call / stateless_apply
- State = lists of trainable, non-trainable and optimizer values
- Assign final values back to the Keras variables
Watch out for
- Forgetting to write the trained values back, then saving a model with its initial weights.