xref: /aosp_15_r20/external/executorch/devtools/bundled_program/schema/bundled_program_schema.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7# pyre-strict
8
9from dataclasses import dataclass
10from typing import List, Union
11
12from executorch.exir.scalar_type import ScalarType
13
14
15@dataclass
16class Tensor:
17    """All information we need to bundle for a tensor EValue input."""
18
19    # The scalar type of Tensor
20    scalar_type: ScalarType
21    # The target sizes of the tensor.
22    sizes: List[int]
23    # The contents of the corresponding tensor.
24    data: bytes
25    dim_order: List[bytes]
26
27
28@dataclass
29class Int:
30    int_val: int
31
32
33@dataclass
34class Bool:
35    bool_val: bool
36
37
38@dataclass
39class Double:
40    double_val: float
41
42
43ValueUnion = Union[
44    Tensor,
45    Int,
46    Double,
47    Bool,
48]
49
50
51@dataclass
52class Value:
53    """Abstraction for BundledIOSet values"""
54
55    val: "ValueUnion"
56
57
58@dataclass
59class BundledMethodTestCase:
60    """All inputs and referenced outputs needs for single verification."""
61
62    # All inputs required by Program for execution. Its length should be
63    # equal to the length of program inputs.
64    inputs: List[Value]
65
66    # The expected outputs generated while running the model in eager mode
67    # using the inputs provided. Its length should be equal to the length
68    # of program outputs.
69    expected_outputs: List[Value]
70
71
72@dataclass
73class BundledMethodTestSuite:
74    """Context for testing and verifying a Method."""
75
76    # The name of the method to test; e.g., "forward" for the forward() method
77    # of an nn.Module. This name match a method defined by the ExecuTorch
78    # program.
79    method_name: str
80
81    # Sets of input/outputs to test with.
82    test_cases: List[BundledMethodTestCase]
83
84
85@dataclass
86class BundledProgram:
87    """ExecuTorch program bunlded with data for verification."""
88
89    # Schema version.
90    version: int
91
92    # Test sets and other meta datas to verify the whole program.
93    # Each BundledMethodTestSuite contains the test cases for one of the Method's
94    # present inside the ExecuTorchProgram of the same BundledProgram. The method_name
95    # present inside the BundledMethodTestSuite is what is used to link to the appropriate Method.
96    method_test_suites: List[BundledMethodTestSuite]
97
98    # The binary data of a serialized ExecuTorchProgram.
99    program: bytes
100