How to create a custom train step in Keras

Built-in Keras training works well until a model needs batch-level logic that compile() cannot express. A custom train_step() lets the model own the forward pass, loss calculation, gradient update, and metric logs while still calling fit().

With the JAX backend, Keras passes model, optimizer, and metric variables into train_step() as explicit state. The method returns both the log dictionary for fit() and the updated state for the next batch, so the code must use stateless layer, loss, optimizer, and metric APIs.

Use this pattern when the project needs custom training math but still wants Keras callbacks, epoch handling, metrics, and history output. The script assumes Keras and JAX are installed, and that KERAS_BACKEND is set before the first Keras import.

1. Create custom_train_step.py with the backend selection, imports, and training arrays.

   custom_train_step.py

   import os
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import jax
   import jax.numpy as jnp
   import keras
   import numpy as np
    
    
   keras.utils.set_random_seed(7)
    
   x_train = np.array(
       [
           [0.0, 0.0],
           [0.0, 1.0],
           [1.0, 0.0],
           [1.0, 1.0],
       ],
       dtype="float32",
   )
   y_train = np.array([[0.0], [1.0], [1.0], [2.0]], dtype="float32")

   Set KERAS_BACKEND before importing keras. A later environment change does not move an already imported process to another backend.
   Related: How to set the Keras backend

2. Add the custom model subclass and stateless loss helper below the training arrays.

   class CustomTrainStepModel(keras.Model):
       def compute_loss_and_updates(
           self,
           trainable_variables,
           non_trainable_variables,
           metrics_variables,
           x,
           y,
           sample_weight,
           training=False,
       ):
           y_pred, non_trainable_variables = self.stateless_call(
               trainable_variables,
               non_trainable_variables,
               x,
               training=training,
           )
           loss, (
               trainable_variables,
               non_trainable_variables,
               metrics_variables,
           ) = self.stateless_compute_loss(
               trainable_variables,
               non_trainable_variables,
               metrics_variables,
               x=x,
               y=y,
               y_pred=y_pred,
               sample_weight=sample_weight,
               training=training,
           )
           return loss, (y_pred, non_trainable_variables, metrics_variables)

   stateless_call() and stateless_compute_loss() let JAX trace the training update without mutating the model object inside the gradient calculation.

3. Add train_step(self, state, data) below the loss helper.

       def train_step(self, state, data):
           (
               trainable_variables,
               non_trainable_variables,
               optimizer_variables,
               metrics_variables,
           ) = state
           x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
    
           grad_fn = jax.value_and_grad(self.compute_loss_and_updates, has_aux=True)
           (loss, (y_pred, non_trainable_variables, metrics_variables)), grads = grad_fn(
               trainable_variables,
               non_trainable_variables,
               metrics_variables,
               x,
               y,
               sample_weight,
               training=True,
           )
           trainable_variables, optimizer_variables = self.optimizer.stateless_apply(
               optimizer_variables,
               grads,
               trainable_variables,
           )
    
           new_metrics_vars = []
           logs = {}
           for metric in self.metrics:
               this_metric_vars = metrics_variables[
                   len(new_metrics_vars) : len(new_metrics_vars) + len(metric.variables)
               ]
               if metric.name == "loss":
                   this_metric_vars = metric.stateless_update_state(
                       this_metric_vars,
                       loss,
                       sample_weight=sample_weight,
                   )
               else:
                   this_metric_vars = metric.stateless_update_state(
                       this_metric_vars,
                       y,
                       y_pred,
                       sample_weight=sample_weight,
                   )
               logs[metric.name] = metric.stateless_result(this_metric_vars)
               new_metrics_vars += this_metric_vars
    
           logs["custom_step_marker"] = jnp.array(1.0)
           state = (
               trainable_variables,
               non_trainable_variables,
               optimizer_variables,
               new_metrics_vars,
           )
           return logs, state

   custom_step_marker is a visible proof value for this short script. Replace it with a real custom metric, loss component, or training signal in project code.

4. Add the model build, compile, fit() call, and printed checks below the class.

   inputs = keras.Input(shape=(2,))
   hidden = keras.layers.Dense(8, activation="relu")(inputs)
   outputs = keras.layers.Dense(1)(hidden)
   model = CustomTrainStepModel(inputs, outputs)
    
   model.compile(
       optimizer=keras.optimizers.Adam(learning_rate=0.05),
       loss=keras.losses.MeanSquaredError(),
       metrics=[keras.metrics.MeanAbsoluteError(name="mae")],
   )
    
   history = model.fit(x_train, y_train, batch_size=4, epochs=3, verbose=0)
    
   print(f"backend: {keras.backend.backend()}")
   print("history keys:", ", ".join(sorted(history.history.keys())))
   print("epochs completed:", len(history.history["loss"]))
   print("custom marker last epoch:", f"{history.history['custom_step_marker'][-1]:.1f}")
   print("final mae:", f"{history.history['mae'][-1]:.4f}")

   history.history receives keys from the dictionary returned by train_step(). Callback logs receive the same names during fit().
   Related: How to compile a model in Keras

5. Run the script and confirm that fit() reports the custom log key.

   $ python3 custom_train_step.py
   backend: jax
   history keys: custom_step_marker, loss, mae
   epochs completed: 3
   custom marker last epoch: 1.0
   final mae: 0.2888