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