def test_lambda_callback_with_batches(self):
        """Test LambdaCallback's behavior with batch-level callbacks."""
        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)
        y = np.random.randn(16, 1)
        lambda_log_callback = callbacks.LambdaCallback(
            on_train_batch_begin=lambda batch, logs: logging.warning(
                "on_train_batch_begin"
            ),
            on_train_batch_end=lambda batch, logs: logging.warning(
                "on_train_batch_end"
            ),
        )
        with self.assertLogs(level="WARNING") as logs:
            model.fit(
                x,
                y,
                batch_size=batch_size,
                validation_split=0.2,
                callbacks=[lambda_log_callback],
                epochs=5,
                verbose=0,
            )
            self.assertTrue(
                any("on_train_batch_begin" in log for log in logs.output)
            )
            self.assertTrue(
                any("on_train_batch_end" in log for log in logs.output)
            )
Test LambdaCallback's behavior with batch-level callbacks.
keras/src/callbacks/lambda_callback_test.py