experimental_experiment.torch_test_helper

More complex helpers used in unit tests.

experimental_experiment.torch_test_helper.check_model_ort(onx: ModelProto, providers: str | List[str] | None = None, dump_file: str | None = None) → onnxruntime.InferenceSession

Loads a model with onnxruntime.

Parameters:

• onx – ModelProto

• providers – list of providers, None fur CPU, cpu for CPU, cuda for CUDA

• dump_file – if not empty, dumps the model into this file if an error happened

Returns:

InferenceSession

experimental_experiment.torch_test_helper.dummy_llm(cls_name: str | None = None, dynamic_shapes: bool = False) → Tuple[torch.nn.Module, Tuple[torch.Tensor, ...]]

Creates a dummy LLM for test purposes.

Parameters:

• cls_name – None for whole model or a piece of it

• dynamic_shapes – returns dynamic shapes as well

<<<

from experimental_experiment.torch_test_helper import dummy_llm

print(dummy_llm())

>>>

    (LLM(
      (embedding): Embedding(
        (embedding): Embedding(1024, 16)
        (pe): Embedding(1024, 16)
      )
      (decoder): DecoderLayer(
        (attention): MultiAttentionBlock(
          (attention): ModuleList(
            (0-1): 2 x AttentionBlock(
              (query): Linear(in_features=16, out_features=16, bias=False)
              (key): Linear(in_features=16, out_features=16, bias=False)
              (value): Linear(in_features=16, out_features=16, bias=False)
            )
          )
          (linear): Linear(in_features=32, out_features=16, bias=True)
        )
        (feed_forward): FeedForward(
          (linear_1): Linear(in_features=16, out_features=128, bias=True)
          (relu): ReLU()
          (linear_2): Linear(in_features=128, out_features=16, bias=True)
        )
        (norm_1): LayerNorm((16,), eps=1e-05, elementwise_affine=True)
        (norm_2): LayerNorm((16,), eps=1e-05, elementwise_affine=True)
      )
    ), (tensor([[923, 801, 883, 749,  61, 342, 462, 402, 576, 442, 967, 544, 619, 376,
             575, 392, 167, 828, 740, 165, 492, 729, 335, 905,  83,  65, 551, 100,
             495, 412]]),))

experimental_experiment.torch_test_helper.export_to_onnx(model: Any, *args: List[Any], verbose: int = 0, return_builder: bool = False, torch_script: bool = True, target_opset: int = 18, prefix: str | None = None, rename_inputs: bool = False, optimize: str | bool = True, folder: str | None = 'dump_test', export_options: ExportOptions | None = None) → Dict[str, str | ModelProto | GraphBuilder]

Exports a model to ONNX.

Parameters:

• model – model to export

• args – arguments

• verbose – verbosity

• return_builder – returns the builder

• torch_script – export with torch.script as well

• target_opset – opset to export into

• prefix – prefix to choose to export into

• rename_inputs – rename the inputs into input_{i}

• optimize – enable, disable optimizations of pattern to test

• folder – where to dump the model, creates it if it does not exist

• export_options – see ExportOptions

Returns:

dictionary with ModelProto, builder, filenames