benchmarks/layer_benchmark/core_benchmark.py/benchmark_embedding
def benchmark_embedding(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Embedding"
    init_args = {
        "input_dim": 128,
        "output_dim": 256,
    }
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[
            256,
        ],
        jit_compile=jit_compile,
    )

    data = [np.random.randint(30, size=(num_samples, 256))]
    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
        data=data,
    )

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