experimental_experiment.torch_bench._bash_bench_models_helper

experimental_experiment.torch_bench._bash_bench_models_helper.get_dummy_model() → Tuple[Callable, Tuple[Any, ...]]

Returns a dummy model used to validate the command line.

experimental_experiment.torch_bench._bash_bench_models_helper.get_dummy_model_fail() → Tuple[Callable, Tuple[Any, ...]]

Returns a dummy model failing.

experimental_experiment.torch_bench._bash_bench_models_helper.get_dummy_model_fail_convert() → Tuple[Callable, Tuple[Any, ...]]

Returns a dummy model.

experimental_experiment.torch_bench._bash_bench_models_helper.get_llama_model_layer(num_hidden_layers: int = 1) → Tuple[Callable, Tuple[Any, ...]]

Returns a llama model with a specific number of layers.

experimental_experiment.torch_bench._bash_bench_models_helper.get_speech2text2_causal_ml_not_trained_model() → Tuple[Callable, Tuple[Any, ...]]

Returns a model. See Speech2Text2Config.

experimental_experiment.torch_bench._bash_bench_models_helper.ids_tensor(shape, vocab_size, rng=None, name=None)

Creates a random int32 tensor of the shape within the vocab size.