Add sample input for a subset of core ATen ops (#262)

Summary:
Pull Request resolved: #262

Basically move all the sample inputs to functional variants. Then we can generate sample inputs for out variants based on that. This is also aligned with our compiler flow where we lower from core ATen ops to their out variants.

Reviewed By: manuelcandales

Differential Revision: D49088856

fbshipit-source-id: 71105773bdc07e7e0ca850ef7e79b3af7f1d5fa5

Author: larryliu0820

exir/dialects/edge/arg/TARGETS

@@ -5,6 +5,7 @@ oncall("ai_infra_mobile_platform")
 python_library(
     name = "lib",
     srcs = [
+        "constraints.py",
         "model.py",
         "type.py",
     ],

exir/dialects/edge/arg/constraints.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+def scatter_add_index_size_max(x, dim, src, d):
+    max_d = src.size(d)
+    if d != dim:
+        max_d = min(max_d, x.size(d))
+    return max_d

exir/dialects/edge/arg/model.py

@@ -359,6 +359,8 @@ def to_out(self, *, name: Optional[str] = None) -> OutArg:
             fill=self.fill,
             size=self.size,
             dtype=self.dtype,
+            value=self.value,
+            deps=self.deps,
         )
 
 

exir/dialects/edge/op/TARGETS

@@ -7,10 +7,12 @@ python_library(
     name = "lib",
     srcs = [
         "api.py",
+        "sample_input.py",
     ],
     deps = [
         "//caffe2:torch",
         "//caffe2/torchgen:torchgen",
+        "//executorch/exir/dialects/edge/arg:lib",
         "//executorch/exir/operator:convert",
     ],
 )