Function save_bundled_program

backends/vulkan/test/utils.py:717–773

Export a bundled .pte file containing the model and test cases. Args: model: The PyTorch model to export sample_inputs: Sample inputs for the model output_path: Path where the bundled .pte file should be saved (should end with .bpte) method_name: Name of the

(
    model: torch.nn.Module,
    sample_inputs: Tuple[torch.Tensor],
    output_path: str,
    method_name: str = "forward",
    sample_kwargs=None,
    et_program: Optional[ExecutorchProgramManager] = None,
    dynamic_shapes=None,
)

def save_bundled_program(
    model: torch.nn.Module,
    sample_inputs: Tuple[torch.Tensor],
    output_path: str,
    method_name: str = "forward",
    sample_kwargs=None,
    et_program: Optional[ExecutorchProgramManager] = None,
    dynamic_shapes=None,
) -> str:
    """
    Export a bundled .pte file containing the model and test cases.

    Args:
        model: The PyTorch model to export
        sample_inputs: Sample inputs for the model
        output_path: Path where the bundled .pte file should be saved (should end with .bpte)
        method_name: Name of the method to test (default: "forward")
        et_program: Optional pre-exported ExecutorchProgramManager. If None, will export to Vulkan
        dynamic_shapes: Optional dynamic shapes for export

    Returns:
        str: Path to the saved bundled program file
    """
    # If no ExecutorchProgramManager provided, export to Vulkan
    if et_program is None:
        et_program = export_model_to_vulkan(
            model,
            sample_inputs,
            sample_kwargs=sample_kwargs,
            dynamic_shapes=dynamic_shapes,
        )

    if sample_kwargs is None:
        sample_kwargs = {}

    # Generate expected outputs by running the model
    expected_outputs = [getattr(model, method_name)(*sample_inputs, **sample_kwargs)]

    # Flatten sample inputs with kwargs to match expected format
    inputs_flattened, _ = tree_flatten((sample_inputs, sample_kwargs))

    # Create bundled program
    bp_buffer = create_bundled_program(
        et_program,
        tuple(inputs_flattened),
        expected_outputs,
        method_name,
    )

    # Ensure output path has correct extension
    if not output_path.endswith(".bpte"):
        output_path = output_path + ".bpte"

    # Write to file
    with open(output_path, "wb") as file:
        file.write(bp_buffer)
    return output_path