xref: /aosp_15_r20/external/pytorch/android/pytorch_android/test_asset.jit (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1def forward(self, input):
2    return None
3
4def eqBool(self, input: bool) -> bool:
5    return input
6
7def eqInt(self, input: int) -> int:
8    return input
9
10def eqFloat(self, input: float) -> float:
11    return input
12
13def eqStr(self, input: str) -> str:
14    return input
15
16def eqTensor(self, input: Tensor) -> Tensor:
17    return input
18
19def eqDictStrKeyIntValue(self, input: Dict[str, int]) -> Dict[str, int]:
20    return input
21
22def eqDictIntKeyIntValue(self, input: Dict[int, int]) -> Dict[int, int]:
23    return input
24
25def eqDictFloatKeyIntValue(self, input: Dict[float, int]) -> Dict[float, int]:
26    return input
27
28def listIntSumReturnTuple(self, input: List[int]) -> Tuple[List[int], int]:
29    sum = 0
30    for x in input:
31        sum += x
32    return (input, sum)
33
34def listBoolConjunction(self, input: List[bool]) -> bool:
35    res = True
36    for x in input:
37        res = res and x
38    return res
39
40def listBoolDisjunction(self, input: List[bool]) -> bool:
41    res = False
42    for x in input:
43        res = res or x
44    return res
45
46def tupleIntSumReturnTuple(self, input: Tuple[int, int, int]) -> Tuple[Tuple[int, int, int], int]:
47    sum = 0
48    for x in input:
49        sum += x
50    return (input, sum)
51
52def optionalIntIsNone(self, input: Optional[int]) -> bool:
53    return input is None
54
55def intEq0None(self, input: int) -> Optional[int]:
56    if input == 0:
57        return None
58    return input
59
60def str3Concat(self, input: str) -> str:
61    return input + input + input
62
63def newEmptyShapeWithItem(self, input):
64    return torch.tensor([int(input.item())])[0]
65
66def testAliasWithOffset(self) -> List[Tensor]:
67    x = torch.tensor([100, 200])
68    a = [x[0], x[1]]
69    return a
70
71def testNonContiguous(self):
72    x = torch.tensor([100, 200, 300])[::2]
73    assert not x.is_contiguous()
74    assert x[0] == 100
75    assert x[1] == 300
76    return x
77
78def conv2d(self, x: Tensor, w: Tensor, toChannelsLast: bool) -> Tensor:
79    r = torch.conv2d(x, w)
80    if (toChannelsLast):
81        # memory_format=torch.channels_last
82        r = r.contiguous(memory_format=2)
83    else:
84        r = r.contiguous()
85    return r
86
87def conv3d(self, x: Tensor, w: Tensor, toChannelsLast: bool) -> Tensor:
88    r = torch.conv3d(x, w)
89    if (toChannelsLast):
90        # memory_format=torch.channels_last_3d
91        r = r.contiguous(memory_format=2)
92    else:
93        r = r.contiguous()
94    return r
95
96def contiguous(self, x: Tensor) -> Tensor:
97    return x.contiguous()
98
99def contiguousChannelsLast(self, x: Tensor) -> Tensor:
100    # memory_format=torch.channels_last
101    return x.contiguous(memory_format=2)
102
103def contiguousChannelsLast3d(self, x: Tensor) -> Tensor:
104    # memory_format=torch.channels_last_3d
105    return x.contiguous(memory_format=3)
106