onnx_diagnostic.torch_models.llms

onnx_diagnostic.torch_models.llms.get_tiny_llm(batch_size: int = 2, input_cache: bool = True, dynamic_rope: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model.

Parameters:

• batch_size – batch size

• input_cache – generate data for this iteration with or without cache

• dynamic_rope – use dynamic rope (see transformers.LlamaConfig)

• kwargs – to overwrite the configuration, example num_hidden_layers=1

Returns:

dictionary

See Steel method forward to guess the dynamic shapes for an example.