xref: /aosp_15_r20/external/executorch/exir/dialects/edge/op/sample_input.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"""
8Sample inputs for Core ATen ops in Portable Kernel Library
9"""
10import torch
11from executorch.exir.dialects.edge.arg.model import InArg, InKwarg, Return
12from executorch.exir.dialects.edge.arg.type import ArgType
13
14
15SAMPLE_INPUT = {
16    "_log_softmax.default": {  # (Tensor self, int dim, bool half_to_float) -> Tensor
17        "args": [
18            InArg(ArgType.Tensor),
19            InArg(ArgType.Param, value=0),
20            InArg(ArgType.Param, value=False),
21        ],
22        "returns": [
23            Return(ArgType.Tensor),
24        ],
25    },
26    "_native_batch_norm_legit_no_training.default": {  # (Tensor input, Tensor? weight, Tensor? bias, Tensor running_mean, Tensor running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor)
27        "args": [
28            InArg(ArgType.Tensor, size=[2, 3, 4, 5]),
29            InArg(ArgType.TensorOpt, size=[3]),
30            InArg(ArgType.TensorOpt, size=[3]),
31            InArg(ArgType.Tensor, size=[3]),
32            InArg(ArgType.Tensor, size=[3]),
33            InArg(ArgType.Param, value=0.1),
34            InArg(ArgType.Param, value=1e-8),
35        ],
36        "returns": [
37            Return(ArgType.Tensor, argname="__ret0", size=[2, 3, 4, 5]),
38            Return(ArgType.Tensor, argname="__ret1", size=[0]),
39            Return(ArgType.Tensor, argname="__ret2", size=[0]),
40        ],
41    },
42    "_softmax.default": {  # (Tensor self, int dim, bool half_to_float) -> Tensor
43        "args": [
44            InArg(ArgType.Tensor),
45            InArg(ArgType.Param, value=0),
46            InArg(ArgType.Param, value=False),
47        ],
48        "returns": [
49            Return(ArgType.Tensor),
50        ],
51    },
52    "_to_copy.default": {  # (Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor
53        "args": [
54            InArg(ArgType.Tensor),
55            InKwarg(ArgType.Param, "non_blocking", value=False),
56            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
57            InKwarg(ArgType.Param, "memory_format", value=None),
58        ],
59        "returns": [
60            Return(ArgType.Tensor),
61        ],
62    },
63    "abs.default": {  # (Tensor self) -> Tensor
64        "args": [
65            InArg(ArgType.Tensor),
66        ],
67        "returns": [Return(ArgType.Tensor)],
68    },
69    "acos.default": {  # (Tensor self) -> Tensor
70        "args": [
71            InArg(ArgType.Tensor),
72        ],
73        "returns": [Return(ArgType.Tensor)],
74    },
75    "acosh.default": {  # (Tensor self) -> Tensor
76        "args": [
77            InArg(ArgType.Tensor),
78        ],
79        "returns": [Return(ArgType.Tensor)],
80    },
81    "add.Tensor": {  # (Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
82        "args": [
83            InArg(ArgType.Tensor),
84            InArg(ArgType.Tensor),
85            InKwarg(ArgType.Scalar, "alpha"),
86        ],
87        "returns": [
88            Return(ArgType.Tensor),
89        ],
90    },
91    "add.Scalar": {  # (Tensor self, Scalar other, Scalar alpha=1) -> Tensor
92        "args": [
93            InArg(ArgType.Tensor),
94            InArg(ArgType.Scalar),
95            InArg(ArgType.Scalar),
96        ],
97        "returns": [
98            Return(ArgType.Tensor),
99        ],
100    },
101    "addmm.default": {  # (Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor
102        "args": [
103            InArg(ArgType.Tensor),
104            InArg(ArgType.Tensor),
105            InArg(ArgType.Tensor),
106            InKwarg(ArgType.Scalar, "beta"),
107            InKwarg(ArgType.Scalar, "alpha"),
108        ],
109        "returns": [
110            Return(ArgType.Tensor),
111        ],
112    },
113    "alias_copy.default": {  # (Tensor self) -> Tensor
114        "args": [
115            InArg(ArgType.Tensor),
116        ],
117        "returns": [Return(ArgType.Tensor)],
118    },
119    "amax.default": {  # (Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
120        "args": [
121            InArg(ArgType.Tensor),
122            InArg(ArgType.Param, value=[0]),
123            InArg(ArgType.Param, value=True),
124        ],
125        "returns": [
126            Return(ArgType.Tensor, size=(1, 2)),
127        ],
128    },
129    "amin.default": {  # (Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
130        "args": [
131            InArg(ArgType.Tensor),
132            InArg(ArgType.Param, value=[0]),
133            InArg(ArgType.Param, value=True),
134        ],
135        "returns": [
136            Return(ArgType.Tensor, size=(1, 2)),
137        ],
138    },
139    "any.default": {  # (Tensor self) -> Tensor
140        "args": [
141            InArg(ArgType.Tensor, dtype=torch.bool),
142        ],
143        "returns": [
144            Return(ArgType.Tensor, size=[], dtype=torch.bool),
145        ],
146    },
147    "arange.default": {  # (Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
148        "args": [
149            InArg(ArgType.Scalar, value=1),
150            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
151        ],
152        "returns": [Return(ArgType.Tensor, size=[1])],
153    },
154    "arange.start_step": {  # (Scalar start, Scalar end, Scalar step=1, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
155        "args": [
156            InArg(ArgType.Scalar, value=0),
157            InArg(ArgType.Scalar, value=1),
158            InArg(ArgType.Scalar, value=1),
159            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
160        ],
161        "returns": [
162            Return(ArgType.Tensor, size=[1]),
163        ],
164    },
165    "argmax.default": {  # (Tensor self, int? dim=None, bool keepdim=False) -> Tensor
166        "args": [
167            InArg(ArgType.Tensor),
168            InArg(ArgType.Param, value=1),
169            InArg(ArgType.Param, value=False),
170        ],
171        "returns": [
172            Return(ArgType.Tensor, size=[2], dtype=torch.long),
173        ],
174    },
175    "argmin.default": {  # (Tensor self, int? dim=None, bool keepdim=False) -> Tensor
176        "args": [
177            InArg(ArgType.Tensor),
178            InArg(ArgType.Param, value=1),
179            InArg(ArgType.Param, value=False),
180        ],
181        "returns": [
182            Return(ArgType.Tensor, size=[2], dtype=torch.long),
183        ],
184    },
185    "as_strided_copy.default": {  # (Tensor self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor
186        "args": [
187            InArg(ArgType.Tensor),
188            InArg(ArgType.Param, value=[4]),
189            InArg(ArgType.Param, value=[1]),
190            InArg(ArgType.Param, value=None),
191        ],
192        "returns": [
193            Return(ArgType.Tensor, size=[4]),
194        ],
195    },
196    "asin.default": {  # (Tensor self) -> Tensor
197        "args": [
198            InArg(ArgType.Tensor),
199        ],
200        "returns": [Return(ArgType.Tensor)],
201    },
202    "asinh.default": {  # (Tensor self) -> Tensor
203        "args": [
204            InArg(ArgType.Tensor),
205        ],
206        "returns": [Return(ArgType.Tensor)],
207    },
208    "atan.default": {  # (Tensor self) -> Tensor
209        "args": [
210            InArg(ArgType.Tensor),
211        ],
212        "returns": [Return(ArgType.Tensor)],
213    },
214    "atanh.default": {  # (Tensor self) -> Tensor
215        "args": [
216            InArg(ArgType.Tensor),
217        ],
218        "returns": [Return(ArgType.Tensor)],
219    },
220    "avg_pool2d.default": {  # (Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor
221        "args": [
222            InArg(ArgType.Tensor, size=[2, 3, 14, 12]),
223            InArg(ArgType.Param, value=[4, 2]),
224            InArg(ArgType.Param, value=[1, 2]),
225            InArg(ArgType.Param, value=[1, 1]),
226            InArg(ArgType.Param, value=True),
227            InArg(ArgType.Param, value=False),
228            InArg(ArgType.Param, value=None),
229        ],
230        "returns": [
231            Return(ArgType.Tensor, size=[2, 3, 13, 7]),
232        ],
233    },
234    "bitwise_and.Scalar": {  # (Tensor self, Scalar other) -> Tensor
235        "args": [
236            InArg(ArgType.Tensor, dtype=torch.bool),
237            InArg(ArgType.Scalar, dtype=torch.bool),
238        ],
239        "returns": [
240            Return(ArgType.Tensor, dtype=torch.bool),
241        ],
242    },
243    "bitwise_and.Tensor": {  # (Tensor self, Tensor other) -> Tensor
244        "args": [
245            InArg(ArgType.Tensor, dtype=torch.bool),
246            InArg(ArgType.Tensor, dtype=torch.bool),
247        ],
248        "returns": [
249            Return(ArgType.Tensor, dtype=torch.bool),
250        ],
251    },
252    "bitwise_not.default": {  # (Tensor self) -> Tensor
253        "args": [
254            InArg(ArgType.Tensor, dtype=torch.bool),
255        ],
256        "returns": [
257            Return(ArgType.Tensor, dtype=torch.bool),
258        ],
259    },
260    "bitwise_or.Scalar": {  # (Tensor self, Scalar other) -> Tensor
261        "args": [
262            InArg(ArgType.Tensor, dtype=torch.bool),
263            InArg(ArgType.Scalar, dtype=torch.bool),
264        ],
265        "returns": [
266            Return(ArgType.Tensor, dtype=torch.bool),
267        ],
268    },
269    "bitwise_or.Tensor": {  # (Tensor self, Tensor other) -> Tensor
270        "args": [
271            InArg(ArgType.Tensor, dtype=torch.bool),
272            InArg(ArgType.Tensor, dtype=torch.bool),
273        ],
274        "returns": [
275            Return(ArgType.Tensor, dtype=torch.bool),
276        ],
277    },
278    "bitwise_xor.Scalar": {  # (Tensor self, Scalar other) -> Tensor
279        "args": [
280            InArg(ArgType.Tensor, dtype=torch.bool),
281            InArg(ArgType.Scalar, dtype=torch.bool),
282        ],
283        "returns": [
284            Return(ArgType.Tensor, dtype=torch.bool),
285        ],
286    },
287    "bitwise_xor.Tensor": {  # (Tensor self, Tensor other) -> Tensor
288        "args": [
289            InArg(ArgType.Tensor, dtype=torch.bool),
290            InArg(ArgType.Tensor, dtype=torch.bool),
291        ],
292        "returns": [
293            Return(ArgType.Tensor, dtype=torch.bool),
294        ],
295    },
296    "bmm.default": {  # (Tensor self, Tensor mat2) -> Tensor
297        "args": [
298            InArg(ArgType.Tensor, size=[1, 2, 2]),
299            InArg(ArgType.Tensor, size=[1, 2, 2]),
300        ],
301        "returns": [
302            Return(ArgType.Tensor, size=[1, 2, 2]),
303        ],
304    },
305    "cat.default": {  # (Tensor[] tensors, int dim=0) -> Tensor
306        "args": [
307            InArg(
308                ArgType.TensorList, value=[InArg(ArgType.Tensor), InArg(ArgType.Tensor)]
309            ),
310            InArg(ArgType.Param, value=0),
311        ],
312        "returns": [
313            Return(ArgType.Tensor, size=[4, 2]),
314        ],
315    },
316    "ceil.default": {  # (Tensor self) -> Tensor
317        "args": [
318            InArg(ArgType.Tensor),
319        ],
320        "returns": [Return(ArgType.Tensor)],
321    },
322    "clamp.default": {  # (Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor
323        "args": [
324            InArg(ArgType.Tensor),
325            InArg(ArgType.ScalarOpt, bounded=True),
326            InArg(ArgType.ScalarOpt, bounded=True),
327        ],
328        "returns": [
329            Return(ArgType.Tensor),
330        ],
331    },
332    "clone.default": {  # (Tensor self, *, MemoryFormat? memory_format=None) -> Tensor
333        "args": [
334            InArg(ArgType.Tensor),
335            InKwarg(ArgType.Param, "memory_format", value=None),
336        ],
337        "returns": [
338            Return(ArgType.Tensor),
339        ],
340    },
341    "constant_pad_nd.default": {  # (Tensor self, SymInt[] pad, Scalar value=0) -> Tensor
342        "args": [
343            InArg(ArgType.Tensor),
344            InArg(ArgType.Param, value=[1, 0, 0, 1]),
345            InArg(ArgType.Scalar, bounded=True),
346        ],
347        "returns": [
348            Return(ArgType.Tensor, size=[3, 3]),
349        ],
350    },
351    "convolution.default": {  # (Tensor input, Tensor weight, Tensor? bias, int[] stride, SymInt[] padding, int[] dilation, bool transposed, SymInt[] output_padding, int groups) -> Tensor
352        "args": [
353            InArg(ArgType.Tensor, size=[1, 2, 5]),
354            InArg(ArgType.Tensor, size=[4, 2, 3]),
355            InArg(ArgType.TensorOpt, size=[4]),
356            InArg(ArgType.Param, value=[2]),
357            InArg(ArgType.Param, value=[2]),
358            InArg(ArgType.Param, value=[1]),
359            InArg(ArgType.Param, value=False),
360            InArg(ArgType.Param, value=[0]),
361            InArg(ArgType.Param, value=1),
362        ],
363        "returns": [
364            Return(ArgType.Tensor, size=[1, 4, 4]),
365        ],
366    },
367    "copy.default": {  # (Tensor self, Tensor src, bool non_blocking=False) -> Tensor
368        "args": [
369            InArg(ArgType.Tensor),
370            InArg(ArgType.Tensor),
371            InArg(ArgType.Param, value=False),
372        ],
373        "returns": [
374            Return(ArgType.Tensor),
375        ],
376    },
377    "cos.default": {  # (Tensor self) -> Tensor
378        "args": [
379            InArg(ArgType.Tensor),
380        ],
381        "returns": [Return(ArgType.Tensor)],
382    },
383    "cosh.default": {  # (Tensor self) -> Tensor
384        "args": [
385            InArg(ArgType.Tensor),
386        ],
387        "returns": [Return(ArgType.Tensor)],
388    },
389    "cumsum.default": {  # (Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor
390        "args": [
391            InArg(ArgType.Tensor),
392            InArg(ArgType.Param, value=0),
393            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
394        ],
395        "returns": [
396            Return(ArgType.Tensor),
397        ],
398    },
399    "detach_copy.default": {  # (Tensor self) -> Tensor
400        "args": [
401            InArg(ArgType.Tensor),
402        ],
403        "returns": [Return(ArgType.Tensor)],
404    },
405    "div.Tensor": {  # (Tensor self, Tensor other) -> Tensor
406        "args": [
407            InArg(ArgType.Tensor),
408            InArg(ArgType.Tensor),
409        ],
410        "returns": [Return(ArgType.Tensor)],
411    },
412    "div.Tensor_mode": {  # (Tensor self, Tensor other, *, str? rounding_mode) -> Tensor
413        "args": [
414            InArg(ArgType.Tensor),
415            InArg(ArgType.Tensor),
416            InKwarg(ArgType.Param, "rounding_mode", value="floor"),
417        ],
418        "returns": [Return(ArgType.Tensor)],
419    },
420    "div.Scalar": {  # (Tensor self, Scalar other) -> Tensor
421        "args": [
422            InArg(ArgType.Tensor),
423            InArg(ArgType.Scalar),
424        ],
425        "returns": [Return(ArgType.Tensor)],
426    },
427    "embedding.default": {  # (Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
428        "args": [
429            InArg(ArgType.Tensor),
430            InArg(ArgType.Tensor, value=[[0, 1], [1, 0]], dtype=torch.long),
431            InArg(ArgType.Param, value=-1),
432            InArg(ArgType.Param, value=False),
433            InArg(ArgType.Param, value=False),
434        ],
435        "returns": [
436            Return(ArgType.Tensor, size=[2, 2, 2]),
437        ],
438    },
439    "empty.memory_format": {  # (SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
440        "args": [
441            InArg(ArgType.Param, value=[2, 2]),
442            InKwarg(ArgType.Param, "memory_format", value=None),
443            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
444        ],
445        "returns": [
446            Return(ArgType.Tensor, fill=3),
447        ],
448    },
449    "eq.Scalar": {  # (Tensor self, Scalar other) -> Tensor
450        "args": [
451            InArg(ArgType.Tensor),
452            InArg(ArgType.Scalar),
453        ],
454        "returns": [Return(ArgType.Tensor)],
455    },
456    "erf.default": {  # (Tensor self) -> Tensor
457        "args": [
458            InArg(ArgType.Tensor),
459        ],
460        "returns": [Return(ArgType.Tensor)],
461    },
462    "exp.default": {  # (Tensor self) -> Tensor
463        "args": [
464            InArg(ArgType.Tensor),
465        ],
466        "returns": [Return(ArgType.Tensor)],
467    },
468    "expand_copy.default": {  # (Tensor self, SymInt[] size, *, bool implicit=False) -> Tensor
469        "args": [
470            InArg(ArgType.Tensor, size=[2, 1]),
471            InArg(ArgType.Param, value=[-1, 3]),
472            InKwarg(ArgType.Param, "implicit", value=False),
473        ],
474        "returns": [
475            Return(ArgType.Tensor, size=[2, 3]),
476        ],
477    },
478    "fill.Scalar": {  # (Tensor self, Scalar value) -> Tensor
479        "args": [
480            InArg(ArgType.Tensor),
481            InArg(ArgType.Scalar, bounded=True),
482        ],
483        "returns": [
484            Return(ArgType.Tensor),
485        ],
486    },
487    "fill.Tensor": {  # (Tensor self, Tensor value) -> Tensor
488        "args": [
489            InArg(ArgType.Tensor),
490            InArg(
491                ArgType.Tensor,
492                size=[],
493                bounded=True,
494            ),
495        ],
496        "returns": [
497            Return(ArgType.Tensor),
498        ],
499    },
500    "floor.default": {  # (Tensor self) -> Tensor
501        "args": [
502            InArg(ArgType.Tensor),
503        ],
504        "returns": [Return(ArgType.Tensor)],
505    },
506    "floor_divide.default": {  # (Tensor self, Tensor other) -> Tensor
507        "args": [
508            InArg(ArgType.Tensor),
509            InArg(ArgType.Tensor, nonzero=True),
510        ],
511        "returns": [
512            Return(ArgType.Tensor),
513        ],
514    },
515    "fmod.Tensor": {  # (Tensor self, Tensor other) -> Tensor
516        "args": [
517            InArg(ArgType.Tensor),
518            InArg(ArgType.Tensor, nonzero=True),
519        ],
520        "returns": [
521            Return(ArgType.Tensor),
522        ],
523    },
524    "fmod.Scalar": {  # (Tensor self, Scalar other) -> Tensor
525        "args": [
526            InArg(ArgType.Tensor),
527            InArg(ArgType.Scalar, value=1),
528        ],
529        "returns": [
530            Return(ArgType.Tensor),
531        ],
532    },
533    "full.default": {  # (SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
534        "args": [
535            InArg(ArgType.Param, value=[2, 2]),
536            InArg(ArgType.Scalar, bounded=True),
537            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
538        ],
539        "returns": [
540            Return(ArgType.Tensor),
541        ],
542    },
543    "full_like.default": {  # (Tensor self, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
544        "args": [
545            InArg(ArgType.Tensor),
546            InArg(ArgType.Scalar, bounded=True),
547            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
548            InKwarg(ArgType.Param, "memory_format", value=None),
549        ],
550        "returns": [
551            Return(ArgType.Tensor),
552        ],
553    },
554    "ge.Scalar": {  # (Tensor self, Scalar other) -> Tensor
555        "args": [
556            InArg(ArgType.Tensor),
557            InArg(ArgType.Scalar),
558        ],
559        "returns": [Return(ArgType.Tensor)],
560    },
561    "ge.Tensor": {  # (Tensor self, Tensor other) -> Tensor
562        "args": [
563            InArg(ArgType.Tensor),
564            InArg(ArgType.Tensor),
565        ],
566        "returns": [Return(ArgType.Tensor)],
567    },
568    "gelu.default": {  # (Tensor self, *, str approximate="none") -> Tensor
569        "args": [
570            InArg(ArgType.Tensor),
571            InKwarg(
572                ArgType.Param,
573                "approximate",
574                value="none",
575            ),
576        ],
577        "returns": [
578            Return(ArgType.Tensor),
579        ],
580    },
581    "glu.default": {  # (Tensor self, int dim=-1) -> Tensor
582        "args": [
583            InArg(ArgType.Tensor),
584            InArg(ArgType.Param, value=0),
585        ],
586        "returns": [
587            Return(ArgType.Tensor, size=[1, 2]),
588        ],
589    },
590    "gt.Scalar": {  # (Tensor self, Scalar other) -> Tensor
591        "args": [
592            InArg(ArgType.Tensor),
593            InArg(ArgType.Scalar),
594        ],
595        "returns": [Return(ArgType.Tensor)],
596    },
597    "gt.Tensor": {  # (Tensor self, Tensor other) -> Tensor
598        "args": [
599            InArg(ArgType.Tensor),
600            InArg(ArgType.Tensor),
601        ],
602        "returns": [Return(ArgType.Tensor)],
603    },
604    "hardtanh.default": {  # (Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor
605        "args": [
606            InArg(ArgType.Tensor),
607            InArg(ArgType.Scalar, bounded=True),
608            InArg(ArgType.Scalar, bounded=True),
609        ],
610        "returns": [
611            Return(ArgType.Tensor),
612        ],
613    },
614    "index.Tensor": {  # (Tensor self, Tensor?[] indices) -> Tensor
615        "args": [
616            InArg(ArgType.Tensor),
617            InArg(
618                ArgType.TensorOptList,
619                value=[
620                    InArg(ArgType.Tensor, value=[0, 1]),
621                    InArg(ArgType.Tensor, value=[1, 1]),
622                ],
623                dtype=torch.long,
624            ),
625        ],
626        "returns": [
627            Return(ArgType.Tensor, size=[2]),
628        ],
629    },
630    "index_put.default": {  # (Tensor self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor
631        "args": [
632            InArg(ArgType.Tensor),
633            InArg(
634                ArgType.TensorOptList,
635                value=[
636                    InArg(ArgType.Tensor, value=[0, 1]),
637                    InArg(ArgType.Tensor, value=[1, 1]),
638                ],
639                dtype=torch.long,
640            ),
641            InArg(ArgType.Tensor, size=[2]),
642            InArg(ArgType.Param, value=False),
643        ],
644        "returns": [
645            Return(ArgType.Tensor),
646        ],
647    },
648    "index_select.default": {  # (Tensor self, int dim, Tensor index) -> Tensor
649        "args": [
650            InArg(ArgType.Tensor),
651            InArg(ArgType.Param, value=0),
652            InArg(ArgType.Tensor, value=[1], dtype=torch.long),
653        ],
654        "returns": [
655            Return(ArgType.Tensor, size=[1, 2]),
656        ],
657    },
658    "isinf.default": {  # (Tensor self) -> Tensor
659        "args": [
660            InArg(ArgType.Tensor),
661        ],
662        "returns": [Return(ArgType.Tensor, dtype=torch.bool)],
663    },
664    "isnan.default": {  # (Tensor self) -> Tensor
665        "args": [
666            InArg(ArgType.Tensor),
667        ],
668        "returns": [Return(ArgType.Tensor, dtype=torch.bool)],
669    },
670    "le.Scalar": {  # (Tensor self, Scalar other) -> Tensor
671        "args": [
672            InArg(ArgType.Tensor),
673            InArg(ArgType.Scalar),
674        ],
675        "returns": [Return(ArgType.Tensor)],
676    },
677    "le.Tensor": {  # (Tensor self, Tensor other) -> Tensor
678        "args": [
679            InArg(ArgType.Tensor),
680            InArg(ArgType.Tensor),
681        ],
682        "returns": [Return(ArgType.Tensor)],
683    },
684    "leaky_relu.default": {  # (Tensor self, Scalar negative_slope=0.01) -> Tensor
685        "args": [
686            InArg(ArgType.Tensor),
687            InArg(ArgType.Scalar),
688        ],
689        "returns": [Return(ArgType.Tensor)],
690    },
691    "lift_fresh_copy.default": {  # (Tensor self) -> Tensor
692        "args": [
693            InArg(ArgType.Tensor),
694        ],
695        "returns": [Return(ArgType.Tensor)],
696    },
697    "log.default": {  # (Tensor self) -> Tensor
698        "args": [
699            InArg(ArgType.Tensor),
700        ],
701        "returns": [Return(ArgType.Tensor)],
702    },
703    "logical_and.default": {  # (Tensor self, Tensor other) -> Tensor
704        "args": [
705            InArg(ArgType.Tensor),
706            InArg(ArgType.Tensor),
707        ],
708        "returns": [Return(ArgType.Tensor)],
709    },
710    "logical_not.default": {  # (Tensor self) -> Tensor
711        "args": [
712            InArg(ArgType.Tensor),
713        ],
714        "returns": [Return(ArgType.Tensor)],
715    },
716    "logical_or.default": {  # (Tensor self, Tensor other) -> Tensor
717        "args": [
718            InArg(ArgType.Tensor),
719            InArg(ArgType.Tensor),
720        ],
721        "returns": [Return(ArgType.Tensor)],
722    },
723    "logical_xor.default": {  # (Tensor self, Tensor other) -> Tensor
724        "args": [
725            InArg(ArgType.Tensor),
726            InArg(ArgType.Tensor),
727        ],
728        "returns": [Return(ArgType.Tensor)],
729    },
730    "logit.default": {  # (Tensor self, float? eps=None) -> Tensor
731        "args": [
732            InArg(ArgType.Tensor),
733            InArg(ArgType.Param, value=0.1),
734        ],
735        "returns": [
736            Return(ArgType.Tensor),
737        ],
738    },
739    "lt.Scalar": {  # (Tensor self, Scalar other) -> Tensor
740        "args": [
741            InArg(ArgType.Tensor),
742            InArg(ArgType.Scalar),
743        ],
744        "returns": [Return(ArgType.Tensor)],
745    },
746    "lt.Tensor": {  # (Tensor self, Tensor other) -> Tensor
747        "args": [
748            InArg(ArgType.Tensor),
749            InArg(ArgType.Tensor),
750        ],
751        "returns": [Return(ArgType.Tensor)],
752    },
753    "masked_fill.Scalar": {  # (Tensor self, Tensor mask, Scalar value) -> Tensor
754        "args": [
755            InArg(ArgType.Tensor),
756            InArg(
757                ArgType.Tensor, value=[[True, False], [False, True]], dtype=torch.bool
758            ),
759            InArg(ArgType.Scalar, bounded=True),
760        ],
761        "returns": [
762            Return(ArgType.Tensor),
763        ],
764    },
765    "max.dim": {  # (Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
766        "args": [
767            InArg(ArgType.Tensor),
768            InArg(ArgType.Param, value=0),
769            InArg(ArgType.Param, value=False),
770        ],
771        "returns": [
772            Return(ArgType.Tensor, argname="values", size=[2]),
773            Return(ArgType.Tensor, argname="indices", size=[2], dtype=torch.long),
774        ],
775    },
776    "max_pool2d_with_indices.default": {  # (Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)
777        "args": [
778            InArg(ArgType.Tensor, size=[2, 12, 12]),
779            InArg(ArgType.Param, value=[4, 3]),
780            InArg(ArgType.Param, value=[3, 2]),
781            InArg(ArgType.Param, value=[2, 1]),
782            InArg(ArgType.Param, value=[1, 2]),
783            InArg(ArgType.Param, value=False),
784        ],
785        "returns": [
786            Return(ArgType.Tensor, size=[2, 5, 5]),
787            Return(ArgType.Tensor, size=[2, 5, 5], dtype=torch.long),
788        ],
789    },
790    "mean.dim": {  # (Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
791        "args": [
792            InArg(ArgType.Tensor),
793            InArg(ArgType.Param, value=[0]),
794            InArg(ArgType.Param, value=False),
795            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
796        ],
797        "returns": [
798            Return(ArgType.Tensor, size=[2]),
799        ],
800    },
801    "min.dim": {  # (Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
802        "args": [
803            InArg(ArgType.Tensor),
804            InArg(ArgType.Param, value=0),
805            InArg(ArgType.Param, value=False),
806        ],
807        "returns": [
808            Return(ArgType.Tensor, argname="values", size=[2]),
809            Return(ArgType.Tensor, argname="indices", size=[2], dtype=torch.long),
810        ],
811    },
812    "minimum.default": {  # (Tensor self, Tensor other) -> Tensor
813        "args": [
814            InArg(ArgType.Tensor),
815            InArg(ArgType.Tensor),
816        ],
817        "returns": [Return(ArgType.Tensor)],
818    },
819    "mm.default": {  # (Tensor self, Tensor mat2) -> Tensor
820        "args": [
821            InArg(ArgType.Tensor),
822            InArg(ArgType.Tensor),
823        ],
824        "returns": [Return(ArgType.Tensor)],
825    },
826    "mul.Tensor": {  # (Tensor self, Tensor other) -> Tensor
827        "args": [
828            InArg(ArgType.Tensor),
829            InArg(ArgType.Tensor),
830        ],
831        "returns": [Return(ArgType.Tensor)],
832    },
833    "mul.Scalar": {  # (Tensor self, Scalar other) -> Tensor
834        "args": [
835            InArg(ArgType.Tensor),
836            InArg(ArgType.Scalar),
837        ],
838        "returns": [Return(ArgType.Tensor)],
839    },
840    "native_layer_norm.default": {  # (Tensor input, SymInt[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)
841        "args": [
842            InArg(ArgType.Tensor),
843            InArg(ArgType.Param, value=[2]),
844            InArg(ArgType.TensorOpt, size=[2]),
845            InArg(ArgType.TensorOpt, size=[2]),
846            InArg(ArgType.Param, value=1e-5),
847        ],
848        "returns": [
849            Return(ArgType.Tensor, argname="__ret0"),
850            Return(ArgType.Tensor, argname="__ret1", size=[2, 1]),
851            Return(ArgType.Tensor, argname="__ret2", size=[2, 1]),
852        ],
853    },
854    "ne.Scalar": {  # (Tensor self, Scalar other) -> Tensor
855        "args": [
856            InArg(ArgType.Tensor),
857            InArg(ArgType.Scalar),
858        ],
859        "returns": [Return(ArgType.Tensor)],
860    },
861    "ne.Tensor": {  # (Tensor self, Tensor other) -> Tensor
862        "args": [
863            InArg(ArgType.Tensor),
864            InArg(ArgType.Tensor),
865        ],
866        "returns": [Return(ArgType.Tensor)],
867    },
868    "neg.default": {  # (Tensor self) -> Tensor
869        "args": [
870            InArg(ArgType.Tensor),
871        ],
872        "returns": [Return(ArgType.Tensor)],
873    },
874    "nonzero.default": {  # (Tensor self) -> Tensor
875        "args": [
876            InArg(ArgType.Tensor, value=[[1, 0], [0, 1]]),
877        ],
878        "returns": [
879            Return(ArgType.Tensor, dtype=torch.long),
880        ],
881    },
882    "ones.default": {  # (SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
883        "args": [
884            InArg(ArgType.Param, value=[2, 2]),
885            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
886        ],
887        "returns": [Return(ArgType.Tensor, fill=3)],
888    },
889    "permute_copy.default": {  # (Tensor self, int[] dims) -> Tensor
890        "args": [
891            InArg(ArgType.Tensor),
892            InArg(ArgType.Param, value=[1, 0]),
893        ],
894        "returns": [
895            Return(ArgType.Tensor),
896        ],
897    },
898    "pixel_shuffle.default": {  # (Tensor self, int upscale_factor) -> Tensor
899        "args": [
900            InArg(ArgType.Tensor, size=[2, 4, 1, 3]),
901            InArg(ArgType.Param, value=2),
902        ],
903        "returns": [
904            Return(ArgType.Tensor, size=[2, 1, 2, 6]),
905        ],
906    },
907    "pow.Tensor_Scalar": {  # (Tensor self, Scalar exponent) -> Tensor
908        "args": [
909            InArg(ArgType.Tensor),
910            InArg(ArgType.Scalar, value=2),
911        ],
912        "returns": [
913            Return(ArgType.Tensor),
914        ],
915    },
916    "pow.Tensor_Tensor": {  # (Tensor self, Tensor exponent) -> Tensor
917        "args": [
918            InArg(ArgType.Tensor),
919            InArg(ArgType.Tensor),
920        ],
921        "returns": [Return(ArgType.Tensor)],
922    },
923    "reciprocal.default": {  # (Tensor self) -> Tensor
924        "args": [
925            InArg(ArgType.Tensor),
926        ],
927        "returns": [Return(ArgType.Tensor)],
928    },
929    "relu.default": {  # (Tensor self) -> Tensor
930        "args": [
931            InArg(ArgType.Tensor),
932        ],
933        "returns": [Return(ArgType.Tensor)],
934    },
935    "remainder.Tensor": {  # (Tensor self, Tensor other) -> Tensor
936        "args": [
937            InArg(ArgType.Tensor),
938            InArg(ArgType.Tensor, nonzero=True),
939        ],
940        "returns": [
941            Return(ArgType.Tensor),
942        ],
943    },
944    "remainder.Scalar": {  # (Tensor self, Scalar other) -> Tensor
945        "args": [
946            InArg(ArgType.Tensor),
947            InArg(ArgType.Scalar, value=1),
948        ],
949        "returns": [
950            Return(ArgType.Tensor),
951        ],
952    },
953    "repeat.default": {  # (Tensor self, SymInt[] repeats) -> Tensor
954        "args": [
955            InArg(ArgType.Tensor),
956            InArg(ArgType.Param, value=[1, 2]),
957        ],
958        "returns": [
959            Return(ArgType.Tensor, size=[2, 4]),
960        ],
961    },
962    "round.default": {  # (Tensor self) -> Tensor
963        "args": [
964            InArg(ArgType.Tensor),
965        ],
966        "returns": [Return(ArgType.Tensor)],
967    },
968    "rsqrt.default": {  # (Tensor self) -> Tensor
969        "args": [
970            InArg(ArgType.Tensor),
971        ],
972        "returns": [Return(ArgType.Tensor)],
973    },
974    "rsub.Scalar": {  # (Tensor self, Scalar other, Scalar alpha=1) -> Tensor
975        "args": [
976            InArg(ArgType.Tensor),
977            InArg(ArgType.Scalar),
978            InArg(ArgType.Scalar),
979        ],
980        "returns": [
981            Return(ArgType.Tensor),
982        ],
983    },
984    "scalar_tensor.default": {  # (Scalar s, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
985        "args": [
986            InArg(ArgType.Scalar, bounded=True),
987            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
988        ],
989        "returns": [Return(ArgType.Tensor, size=[])],
990    },
991    "scatter_add.default": {  # (Tensor self, int dim, Tensor index, Tensor src) -> Tensor
992        "args": [
993            InArg(ArgType.Tensor),
994            InArg(ArgType.Param, value=0),
995            InArg(ArgType.Tensor, value=[[0, 1]], dtype=torch.long),
996            InArg(ArgType.Tensor),
997        ],
998        "returns": [
999            Return(ArgType.Tensor),
1000        ],
1001    },
1002    "select_copy.int": {  # (Tensor self, int dim, SymInt index) -> Tensor
1003        "args": [
1004            InArg(ArgType.Tensor),
1005            InArg(ArgType.Param, value=0),
1006            InArg(ArgType.Param, value=1),
1007        ],
1008        "returns": [
1009            Return(ArgType.Tensor, size=[2]),
1010        ],
1011    },
1012    "select_scatter.default": {  # (Tensor self, Tensor src, int dim, SymInt index) -> Tensor
1013        "args": [
1014            InArg(ArgType.Tensor),
1015            InArg(ArgType.Tensor, size=[2]),
1016            InArg(ArgType.Param, value=0),
1017            InArg(ArgType.Param, value=1),
1018        ],
1019        "returns": [
1020            Return(ArgType.Tensor),
1021        ],
1022    },
1023    "sigmoid.default": {  # (Tensor self) -> Tensor
1024        "args": [
1025            InArg(ArgType.Tensor),
1026        ],
1027        "returns": [Return(ArgType.Tensor)],
1028    },
1029    "sign.default": {  # (Tensor self) -> Tensor
1030        "args": [
1031            InArg(ArgType.Tensor),
1032        ],
1033        "returns": [Return(ArgType.Tensor)],
1034    },
1035    "sin.default": {  # (Tensor self) -> Tensor
1036        "args": [
1037            InArg(ArgType.Tensor),
1038        ],
1039        "returns": [Return(ArgType.Tensor)],
1040    },
1041    "sinh.default": {  # (Tensor self) -> Tensor
1042        "args": [
1043            InArg(ArgType.Tensor),
1044        ],
1045        "returns": [Return(ArgType.Tensor)],
1046    },
1047    "slice_copy.Tensor": {  # (Tensor self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor
1048        "args": [
1049            InArg(ArgType.Tensor),
1050            InArg(ArgType.Param, value=0),
1051            InArg(ArgType.Param, value=None),
1052            InArg(ArgType.Param, value=None),
1053            InArg(ArgType.Param, value=1),
1054        ],
1055        "returns": [
1056            Return(ArgType.Tensor),
1057        ],
1058    },
1059    "slice_scatter.default": {  # (Tensor self, Tensor src, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor
1060        "args": [
1061            InArg(ArgType.Tensor),
1062            InArg(ArgType.Tensor),
1063            InArg(ArgType.Param, value=0),
1064            InArg(ArgType.Param, value=None),
1065            InArg(ArgType.Param, value=None),
1066            InArg(ArgType.Param, value=1),
1067        ],
1068        "returns": [
1069            Return(ArgType.Tensor),
1070        ],
1071    },
1072    "split_copy.Tensor": {  # (Tensor self, SymInt split_size, int dim=0) -> Tensor[]
1073        "args": [
1074            InArg(ArgType.Tensor),
1075            InArg(ArgType.Param, value=1),
1076            InArg(ArgType.Param, value=0),
1077        ],
1078        "returns": [
1079            Return(
1080                ArgType.TensorList,
1081                value=[
1082                    Return(ArgType.Tensor, size=[1, 2]),
1083                    Return(ArgType.Tensor, size=[1, 2]),
1084                ],
1085            ),
1086        ],
1087    },
1088    "split_with_sizes_copy.default": {  # (Tensor self, SymInt[] split_sizes, int dim=0) -> Tensor[]
1089        "args": [
1090            InArg(ArgType.Tensor, size=[2, 6, 3]),
1091            InArg(ArgType.Param, value=[3, 1, 2]),
1092            InArg(ArgType.Param, value=1),
1093        ],
1094        "returns": [
1095            Return(
1096                ArgType.TensorList,
1097                value=[
1098                    Return(ArgType.Tensor, size=[2, 3, 3]),
1099                    Return(ArgType.Tensor, size=[2, 1, 3]),
1100                    Return(ArgType.Tensor, size=[2, 2, 3]),
1101                ],
1102            ),
1103        ],
1104    },
1105    "sqrt.default": {  # (Tensor self) -> Tensor
1106        "args": [
1107            InArg(ArgType.Tensor),
1108        ],
1109        "returns": [Return(ArgType.Tensor)],
1110    },
1111    "squeeze_copy.dim": {  # (Tensor self, int dim) -> Tensor
1112        "args": [
1113            InArg(ArgType.Tensor, size=[1, 2]),
1114            InArg(ArgType.Param, value=0),
1115        ],
1116        "returns": [
1117            Return(ArgType.Tensor, size=[2]),
1118        ],
1119    },
1120    "squeeze_copy.dims": {  # (Tensor self, int[] dims) -> Tensor
1121        "args": [
1122            InArg(ArgType.Tensor, size=[1, 2, 1, 5]),
1123            InArg(ArgType.Param, value=[0, 2]),
1124        ],
1125        "returns": [
1126            Return(ArgType.Tensor, size=[2, 5]),
1127        ],
1128    },
1129    "stack.default": {  # (Tensor[] tensors, int dim=0) -> Tensor
1130        "args": [
1131            InArg(
1132                ArgType.TensorList,
1133                value=[
1134                    InArg(ArgType.Tensor),
1135                    InArg(ArgType.Tensor),
1136                    InArg(ArgType.Tensor),
1137                ],
1138            ),
1139            InArg(ArgType.Param, value=0),
1140        ],
1141        "returns": [
1142            Return(ArgType.Tensor, size=[3, 2, 2]),
1143        ],
1144    },
1145    "sub.Tensor": {  # (Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
1146        "args": [
1147            InArg(ArgType.Tensor),
1148            InArg(ArgType.Tensor),
1149            InKwarg(ArgType.Scalar, "alpha"),
1150        ],
1151        "returns": [
1152            Return(ArgType.Tensor),
1153        ],
1154    },
1155    "sub.Scalar": {  # (Tensor self, Scalar other, Scalar alpha=1) -> Tensor
1156        "args": [
1157            InArg(ArgType.Tensor),
1158            InArg(ArgType.Scalar),
1159            InArg(ArgType.Scalar),
1160        ],
1161        "returns": [
1162            Return(ArgType.Tensor),
1163        ],
1164    },
1165    "sum.dim_IntList": {  # (Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
1166        "args": [
1167            InArg(ArgType.Tensor),
1168            InArg(ArgType.Param, value=[0]),
1169            InArg(ArgType.Param, value=False),
1170            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
1171        ],
1172        "returns": [
1173            Return(ArgType.Tensor, size=[2]),
1174        ],
1175    },
1176    "t_copy.default": {  # (Tensor self) -> Tensor
1177        "args": [
1178            InArg(ArgType.Tensor),
1179        ],
1180        "returns": [Return(ArgType.Tensor)],
1181    },
1182    "tan.default": {  # (Tensor self) -> Tensor
1183        "args": [
1184            InArg(ArgType.Tensor),
1185        ],
1186        "returns": [Return(ArgType.Tensor)],
1187    },
1188    "tanh.default": {  # (Tensor self) -> Tensor
1189        "args": [
1190            InArg(ArgType.Tensor),
1191        ],
1192        "returns": [Return(ArgType.Tensor)],
1193    },
1194    "transpose_copy.int": {  # (Tensor self, int dim0, int dim1) -> Tensor
1195        "args": [
1196            InArg(ArgType.Tensor),
1197            InArg(ArgType.Param, value=0),
1198            InArg(ArgType.Param, value=1),
1199        ],
1200        "returns": [
1201            Return(ArgType.Tensor),
1202        ],
1203    },
1204    "tril.default": {  # (Tensor self, int diagonal=0) -> Tensor
1205        "args": [
1206            InArg(ArgType.Tensor),
1207            InArg(ArgType.Param, value=0),
1208        ],
1209        "returns": [
1210            Return(ArgType.Tensor),
1211        ],
1212    },
1213    "unbind_copy.int": {  # (Tensor self, int dim=0) -> Tensor[]
1214        "args": [
1215            InArg(ArgType.Tensor),
1216            InArg(ArgType.Param, value=0),
1217        ],
1218        "returns": [
1219            Return(
1220                ArgType.TensorList,
1221                value=[
1222                    Return(ArgType.Tensor, size=[2]),
1223                    Return(ArgType.Tensor, size=[2]),
1224                ],
1225            ),
1226        ],
1227    },
1228    "unsqueeze_copy.default": {  # (Tensor self, int dim) -> Tensor
1229        "args": [
1230            InArg(ArgType.Tensor),
1231            InArg(ArgType.Param, value=1),
1232        ],
1233        "returns": [
1234            Return(ArgType.Tensor, size=[2, 1, 2]),
1235        ],
1236    },
1237    "var.dim": {  # (Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> Tensor
1238        "args": [
1239            InArg(ArgType.Tensor),
1240            InArg(ArgType.Param, value=[0]),
1241            InArg(ArgType.Param, value=True),
1242            InArg(ArgType.Param, value=False),
1243        ],
1244        "returns": [
1245            Return(ArgType.Tensor, size=[2]),
1246        ],
1247    },
1248    "view_copy.default": {  # (Tensor self, SymInt[] size) -> Tensor
1249        "args": [
1250            InArg(ArgType.Tensor),
1251            InArg(ArgType.Param, value=[4]),
1252        ],
1253        "returns": [
1254            Return(ArgType.Tensor, size=[4]),
1255        ],
1256    },
1257    "where.self": {  # (Tensor condition, Tensor self, Tensor other) -> Tensor
1258        "args": [
1259            InArg(ArgType.Tensor, dtype=torch.bool),
1260            InArg(ArgType.Tensor),
1261            InArg(ArgType.Tensor),
1262        ],
1263        "returns": [
1264            Return(ArgType.Tensor),
1265        ],
1266    },
1267    "zeros.default": {  # (SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
1268        "args": [
1269            InArg(ArgType.Param, value=[2, 2]),
1270            InKwarg(ArgType.ScalarTypeOpt, "dtype"),
1271        ],
1272        "returns": [Return(ArgType.Tensor, fill=3)],
1273    },
1274}
1275