keras/src/backend/numpy/trainer.py/NumpyTrainer/predict_on_batch
class NumpyTrainer:    def predict_on_batch(self, x):
        self.make_predict_function()
        batch_outputs = self.predict_function([(x,)])
        batch_outputs = tree.map_structure(
            backend.convert_to_numpy, batch_outputs
        )
        return batch_outputs