benchmarks/layer_benchmark/attention_benchmark.py/benchmark_attention
def benchmark_attention(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Attention"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[[256, 64], [256, 64]],
        flat_call_inputs=False,
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )