def METHOD_NAME(self, num_elements, batch_size):
    """Test the KPL adapt implementation."""
    input_t = keras.Input(shape=(1,), dtype=tf.float32)
    layer = normalization.Normalization()
    _ = layer(input_t)
    num_repeats = 5
    starts = []
    ends = []
    for _ in range(num_repeats):
        ds = tf.data.Dataset.range(num_elements)
        ds = ds.map(lambda x: tf.expand_dims(tf.cast(x, tf.float32), -1))
        ds = ds.batch(batch_size)
        starts.append(time.time())
        # Benchmarked code begins here.
        layer.adapt(ds)
        # Benchmarked code ends here.
        ends.append(time.time())
    avg_time = np.mean(np.array(ends) - np.array(starts))
    name = f"normalization_adapt|{num_elements}_elements|batch_{batch_size}"
    baseline = self.run_dataset_implementation(num_elements, batch_size)
    extras = {
        "tf.data implementation baseline": baseline,
        "delta seconds": (baseline - avg_time),
        "delta percent": ((baseline - avg_time) / baseline) * 100,
    }
    self.report_benchmark(
        iters=num_repeats, wall_time=avg_time, extras=extras, name=name
    )