benchmarks/layer_benchmark/rnn_benchmark.py/benchmark_time_distributed
def benchmark_time_distributed(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "TimeDistributed"
    init_args = {}
    keras_layer = keras.layers.TimeDistributed(keras.layers.Conv2D(16, (3, 3)))
    tf_keras_layer = tf.keras.layers.TimeDistributed(
        tf.keras.layers.Conv2D(16, (3, 3))
    )
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[10, 32, 32, 3],
        jit_compile=jit_compile,
        keras_layer=keras_layer,
        tf_keras_layer=tf_keras_layer,
    )

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

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