def test_lambda_callback_during_prediction(self):
        """Test LambdaCallback's functionality during model prediction."""
        batch_size = 4
        model = Sequential(
            [layers.Input(shape=(2,), batch_size=batch_size), layers.Dense(1)]
        )
        model.compile(
            optimizer=optimizers.SGD(), loss=losses.MeanSquaredError()
        )
        x = np.random.randn(16, 2)

        def custom_on_predict_begin(logs):
            logging.warning("on_predict_begin_executed")

        lambda_callback = callbacks.LambdaCallback(
            on_predict_begin=custom_on_predict_begin
        )
        with self.assertLogs(level="WARNING") as logs:
            model.predict(
                x, batch_size=batch_size, callbacks=[lambda_callback], verbose=0
            )
            self.assertTrue(
                any("on_predict_begin_executed" in log for log in logs.output)
            )
Test LambdaCallback's functionality during model prediction.
keras/src/callbacks/lambda_callback_test.py