Skip to content

Commit 2de02d3

Browse files
GregoryComerGregoryComer
committed
[Backend Tester] Add reduction op tests
ghstack-source-id: ef74349 ghstack-comment-id: 3116316841 Pull-Request: #12853
1 parent 29721df commit 2de02d3

File tree

7 files changed

+2088
-2
lines changed

7 files changed

+2088
-2
lines changed

‎backends/test/suite/operators/test_amax.py‎

Lines changed: 388 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,388 @@
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-unsafe
8+
9+
from typing import List, Optional, Tuple, Union
10+
11+
import torch
12+
from executorch.backends.test.suite.flow import TestFlow
13+
14+
from executorch.backends.test.suite.operators import (
15+
dtype_test,
16+
operator_test,
17+
OperatorTest,
18+
)
19+
20+
21+
class AmaxModel(torch.nn.Module):
22+
def __init__(
23+
self,
24+
dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None,
25+
keepdim: bool = False,
26+
):
27+
super().__init__()
28+
self.dim = dim
29+
self.keepdim = keepdim
30+
31+
def forward(self, x):
32+
return torch.amax(x, dim=self.dim, keepdim=self.keepdim)
33+
34+
35+
@operator_test
36+
class Amax(OperatorTest):
37+
@dtype_test
38+
def test_amax_dtype(self, flow: TestFlow, dtype) -> None:
39+
self._test_op(
40+
AmaxModel().to(dtype),
41+
(torch.rand(10, 10).to(dtype),),
42+
flow,
43+
)
44+
45+
def test_amax_basic(self, flow: TestFlow) -> None:
46+
self._test_op(
47+
AmaxModel(),
48+
(torch.randn(10, 10),),
49+
flow,
50+
)
51+
52+
def test_amax_dim(self, flow: TestFlow) -> None:
53+
self._test_op(
54+
AmaxModel(dim=0),
55+
(torch.randn(5, 10),),
56+
flow,
57+
)
58+
59+
self._test_op(
60+
AmaxModel(dim=1),
61+
(torch.randn(5, 10),),
62+
flow,
63+
)
64+
65+
self._test_op(
66+
AmaxModel(dim=0),
67+
(torch.randn(3, 4, 5),),
68+
flow,
69+
)
70+
71+
self._test_op(
72+
AmaxModel(dim=1),
73+
(torch.randn(3, 4, 5),),
74+
flow,
75+
)
76+
77+
self._test_op(
78+
AmaxModel(dim=2),
79+
(torch.randn(3, 4, 5),),
80+
flow,
81+
)
82+
83+
self._test_op(
84+
AmaxModel(dim=1),
85+
(torch.randn(2, 3, 4, 5),),
86+
flow,
87+
)
88+
89+
self._test_op(
90+
AmaxModel(dim=-1),
91+
(torch.randn(3, 4, 5),),
92+
flow,
93+
)
94+
95+
self._test_op(
96+
AmaxModel(dim=-2),
97+
(torch.randn(3, 4, 5),),
98+
flow,
99+
)
100+
101+
def test_amax_multi_dim(self, flow: TestFlow) -> None:
102+
self._test_op(
103+
AmaxModel(dim=(0, 1)),
104+
(torch.randn(3, 4, 5),),
105+
flow,
106+
)
107+
108+
self._test_op(
109+
AmaxModel(dim=(0, 2)),
110+
(torch.randn(3, 4, 5),),
111+
flow,
112+
)
113+
114+
self._test_op(
115+
AmaxModel(dim=(1, 2)),
116+
(torch.randn(3, 4, 5),),
117+
flow,
118+
)
119+
120+
self._test_op(
121+
AmaxModel(dim=(1, 3)),
122+
(torch.randn(2, 3, 4, 5),),
123+
flow,
124+
)
125+
126+
self._test_op(
127+
AmaxModel(dim=(0, 2)),
128+
(torch.randn(2, 3, 4, 5),),
129+
flow,
130+
)
131+
132+
self._test_op(
133+
AmaxModel(dim=(-1, -3)),
134+
(torch.randn(2, 3, 4, 5),),
135+
flow,
136+
)
137+
138+
self._test_op(
139+
AmaxModel(dim=(0, 1, 2, 3)),
140+
(torch.randn(2, 3, 4, 5),),
141+
flow,
142+
)
143+
144+
def test_amax_keepdim(self, flow: TestFlow) -> None:
145+
self._test_op(
146+
AmaxModel(dim=0, keepdim=True),
147+
(torch.randn(5, 10),),
148+
flow,
149+
)
150+
151+
self._test_op(
152+
AmaxModel(dim=1, keepdim=True),
153+
(torch.randn(5, 10),),
154+
flow,
155+
)
156+
157+
self._test_op(
158+
AmaxModel(dim=1, keepdim=True),
159+
(torch.randn(3, 4, 5),),
160+
flow,
161+
)
162+
163+
self._test_op(
164+
AmaxModel(dim=2, keepdim=True),
165+
(torch.randn(2, 3, 4, 5),),
166+
flow,
167+
)
168+
169+
self._test_op(
170+
AmaxModel(dim=(1, 2), keepdim=True),
171+
(torch.randn(3, 4, 5),),
172+
flow,
173+
)
174+
175+
def test_amax_shapes(self, flow: TestFlow) -> None:
176+
self._test_op(
177+
AmaxModel(),
178+
(torch.randn(20),),
179+
flow,
180+
)
181+
self._test_op(
182+
AmaxModel(dim=0),
183+
(torch.randn(20),),
184+
flow,
185+
)
186+
187+
self._test_op(
188+
AmaxModel(),
189+
(torch.randn(5, 10),),
190+
flow,
191+
)
192+
193+
self._test_op(
194+
AmaxModel(),
195+
(torch.randn(3, 4, 5),),
196+
flow,
197+
)
198+
199+
self._test_op(
200+
AmaxModel(),
201+
(torch.randn(2, 3, 4, 5),),
202+
flow,
203+
)
204+
205+
self._test_op(
206+
AmaxModel(),
207+
(torch.randn(2, 2, 3, 4, 5),),
208+
flow,
209+
)
210+
211+
def test_amax_values(self, flow: TestFlow) -> None:
212+
x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
213+
self._test_op(
214+
AmaxModel(),
215+
(x,),
216+
flow,
217+
)
218+
self._test_op(
219+
AmaxModel(dim=0),
220+
(x,),
221+
flow,
222+
)
223+
self._test_op(
224+
AmaxModel(dim=1),
225+
(x,),
226+
flow,
227+
)
228+
229+
x = torch.tensor([[3.0, 2.0, 3.0], [6.0, 6.0, 5.0]])
230+
self._test_op(
231+
AmaxModel(),
232+
(x,),
233+
flow,
234+
)
235+
self._test_op(
236+
AmaxModel(dim=0),
237+
(x,),
238+
flow,
239+
)
240+
self._test_op(
241+
AmaxModel(dim=1),
242+
(x,),
243+
flow,
244+
)
245+
246+
x = torch.tensor([[-3.0, -2.0, -1.0], [-6.0, -5.0, -4.0]])
247+
self._test_op(
248+
AmaxModel(),
249+
(x,),
250+
flow,
251+
)
252+
self._test_op(
253+
AmaxModel(dim=0),
254+
(x,),
255+
flow,
256+
)
257+
self._test_op(
258+
AmaxModel(dim=1),
259+
(x,),
260+
flow,
261+
)
262+
263+
x = torch.tensor([[-3.0, 2.0, -1.0], [6.0, -5.0, 4.0]])
264+
self._test_op(
265+
AmaxModel(),
266+
(x,),
267+
flow,
268+
)
269+
self._test_op(
270+
AmaxModel(dim=0),
271+
(x,),
272+
flow,
273+
)
274+
self._test_op(
275+
AmaxModel(dim=1),
276+
(x,),
277+
flow,
278+
)
279+
280+
def test_amax_edge_cases(self, flow: TestFlow) -> None:
281+
x = torch.ones(3, 4)
282+
self._test_op(
283+
AmaxModel(),
284+
(x,),
285+
flow,
286+
)
287+
self._test_op(
288+
AmaxModel(dim=0),
289+
(x,),
290+
flow,
291+
)
292+
self._test_op(
293+
AmaxModel(dim=1),
294+
(x,),
295+
flow,
296+
)
297+
298+
x = torch.zeros(3, 4)
299+
self._test_op(
300+
AmaxModel(),
301+
(x,),
302+
flow,
303+
)
304+
self._test_op(
305+
AmaxModel(dim=0),
306+
(x,),
307+
flow,
308+
)
309+
self._test_op(
310+
AmaxModel(dim=1),
311+
(x,),
312+
flow,
313+
)
314+
315+
x = torch.tensor([[1.0, float("inf"), 3.0], [4.0, 5.0, float("inf")]])
316+
self._test_op(
317+
AmaxModel(),
318+
(x,),
319+
flow,
320+
)
321+
self._test_op(
322+
AmaxModel(dim=0),
323+
(x,),
324+
flow,
325+
)
326+
self._test_op(
327+
AmaxModel(dim=1),
328+
(x,),
329+
flow,
330+
)
331+
332+
x = torch.tensor([[1.0, float("-inf"), 3.0], [4.0, 5.0, float("-inf")]])
333+
self._test_op(
334+
AmaxModel(),
335+
(x,),
336+
flow,
337+
)
338+
self._test_op(
339+
AmaxModel(dim=0),
340+
(x,),
341+
flow,
342+
)
343+
self._test_op(
344+
AmaxModel(dim=1),
345+
(x,),
346+
flow,
347+
)
348+
349+
x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]])
350+
self._test_op(
351+
AmaxModel(),
352+
(x,),
353+
flow,
354+
)
355+
self._test_op(
356+
AmaxModel(dim=0),
357+
(x,),
358+
flow,
359+
)
360+
self._test_op(
361+
AmaxModel(dim=1),
362+
(x,),
363+
flow,
364+
)
365+
366+
x = torch.tensor([5.0])
367+
self._test_op(
368+
AmaxModel(),
369+
(x,),
370+
flow,
371+
)
372+
self._test_op(
373+
AmaxModel(dim=0),
374+
(x,),
375+
flow,
376+
)
377+
378+
def test_amax_scalar(self, flow: TestFlow) -> None:
379+
self._test_op(
380+
AmaxModel(),
381+
(torch.tensor([5.0]),),
382+
flow,
383+
)
384+
self._test_op(
385+
AmaxModel(dim=0),
386+
(torch.tensor([5.0]),),
387+
flow,
388+
)

0 commit comments

Comments
 (0)