In Keras 3, how do you override train_step(), and what must it return?

## The seam Keras gives you `fit()` is a loop that does a great deal you do not want to rewrite: epoch iteration, batching, shuffling, the progress bar, callback dispatch, the validation pass at each epoch end, `History` bookkeeping, and distribution. All of the per-batch *arithmetic* lives in one overridable method, `Model.train_step(self, data)`. Overriding it is the supported way to change what a training batch does — a custom gradient manipulation, an extra regulariser, a GAN's two-optimizer dance — without giving up the machinery around it. `test_step` is its counterpart for `evaluate()` and for the validation pass inside `fit()`, and `predict_step` for `predict()`. ## The contract ``` class CustomModel(keras.Model): def train_step(self, data): x, y = data # 1. forward pass in training mode # 2. scalar loss # 3. gradients -> optimizer # 4. metric.update_state(...) return {m.name: m.result() for m in self.metrics} ``` Four obligations: 1. **Unpack `data` yourself.** It is exactly what your input yields — a `(x, y)` tuple from arrays, a `(x, y, sample_weight)` triple, or a single tensor for an unsupervised model. Handle the arity you actually use. 2. **Call the model with `training=True`.** `self(x, training=True)` is what puts `Dropout` and `BatchNormalization` into training behaviour. Forget the flag and dropout is off during training — a silent quality bug. 3. **Use `self.compute_loss(x=..., y=..., y_pred=..., sample_weight=...)`** rather than calling the loss function directly. It applies the loss you passed to `compile()`, handles multi-output losses and `loss_weights`, and includes any regularisation losses layers added. Calling `keras.losses.mse(y, y_pred)` by hand bypasses all of that. 4. **Return a dict of metric names to scalar results.** The returned mapping is the logs — the progress bar prints it, callbacks receive it as `logs`, and `History.history` accumulates it. Returning the raw loss tensor instead of a dict breaks all three. Metrics themselves are stateful objects: you call `update_state()` inside the step, and `fit()` resets them at each epoch boundary, so the printed value is the epoch-so-far accumulation. ## The backend-specific part Keras 3 runs on TensorFlow, JAX and PyTorch, and *gradient computation is the one part of `train_step` that cannot be written backend-agnostically*, because each framework's autodiff has a different shape: - **TensorFlow backend** — wrap the forward pass in a `tf.GradientTape`, get `tape.gradient(loss, self.trainable_variables)`, then `self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))`. - **PyTorch backend** — clear existing gradients, call `loss.backward()`, and hand the resulting `.grad` values to the optimizer. - **JAX backend** — the step is functional and stateless: you override `compute_loss_and_updates` and a `train_step(self, state, data)` that receives and returns the state, because JAX cannot mutate variables in place inside a transformed function. So a `train_step` written against `tf.GradientTape` pins the model to the TensorFlow backend. If portability matters, that is a real cost to weigh, and everything *else* in the step — the forward call, `compute_loss`, metric updates — can be written with `keras.ops` and stays portable. ## What changed from Keras 2 Keras 2 examples call `self.compiled_loss(y, y_pred, regularization_losses=self.losses)` and `self.compiled_metrics.update_state(y, y_pred)`. **Neither attribute exists in Keras 3.** They were replaced by the public `self.compute_loss(...)` and `self.compute_metrics(...)`, plus explicit iteration over `self.metrics`. Pasting a Keras 2 custom-step snippet into a Keras 3 project fails with an attribute error, and recognising that instantly is a good signal in an interview. ## When it is the right tool — and when it is not Override `train_step` when the *batch computation* is non-standard: multiple optimizers, gradient surgery, a teacher-student distillation loss, an adversarial inner step, or a loss that needs the inputs as well as the targets. Do not reach for it when a simpler seam exists — a custom loss function, a custom `keras.metrics.Metric`, or a callback covers a large share of what people first attempt here, and each is far less code to maintain. ## The forgotten half The most common defect after a successful `train_step` override is leaving `test_step` alone. Validation inside `fit()` and every call to `evaluate()` go through `test_step`, so if your training loss is a custom composite and your evaluation is not, `loss` and `val_loss` are measuring different quantities — and the two curves will never be comparable, which usually gets misread as overfitting.