keras/src/backend/tensorflow/trainer.py/TensorFlowTrainer/test_on_batch
class TensorFlowTrainer:    def test_on_batch(
        self,
        x,
        y=None,
        sample_weight=None,
        return_dict=False,
    ):
        self._assert_compile_called("test_on_batch")

        def data():
            yield (x, y, sample_weight)

        # Maybe build model
        self._symbolic_build(data_batch=(x, y, sample_weight))
        self.make_test_function()

        logs = self.test_function(data())
        logs = tree.map_structure(lambda x: np.array(x), logs)
        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)