Add scatter operation to direct bindings

python/python_direct/ops.cpp

@@ -2256,15 +2256,15 @@ list of Val
 
 void bindIndexingOps(py::module_& ops) {
   ops.def(
-         "index_select",
-         [](TensorView* arg, TensorView* index, int64_t dim) -> TensorView* {
-           return indexSelect(arg, dim, index);
-         },
-         py::arg("arg"),
-         py::arg("index"),
-         py::arg("dim"),
-         py::return_value_policy::reference,
-         R"(
+      "index_select",
+      [](TensorView* arg, TensorView* index, int64_t dim) -> TensorView* {
+        return indexSelect(arg, dim, index);
+      },
+      py::arg("arg"),
+      py::arg("index"),
+      py::arg("dim"),
+      py::return_value_policy::reference,
+      R"(
 Select elements from a tensor along a specified dimension.
 
 Parameters
@@ -2278,17 +2278,17 @@ Returns
 -------
 TensorView
     The selected tensor.
-)")
-      .def(
-          "select",
-          [](TensorView* arg, Val* index, int64_t dim) -> TensorView* {
-            return select(arg, dim, index);
-          },
-          py::arg("arg"),
-          py::arg("index"),
-          py::arg("dim"),
-          py::return_value_policy::reference,
-          R"(
+)");
+  ops.def(
+      "select",
+      [](TensorView* arg, Val* index, int64_t dim) -> TensorView* {
+        return select(arg, dim, index);
+      },
+      py::arg("arg"),
+      py::arg("index"),
+      py::arg("dim"),
+      py::return_value_policy::reference,
+      R"(
 Select elements from a tensor along a specified dimension.
 
 Parameters
@@ -2303,6 +2303,55 @@ Returns
 TensorView
     The selected tensor.
 )");
+  ops.def(
+      "scatter",
+      [](TensorView* arg1, TensorView* index, TensorView* src, int64_t dim)
+          -> TensorView* {
+        NVF_CHECK(
+            arg1->nDims() == index->nDims() && arg1->nDims() == src->nDims(),
+            "Tensor arguments have different dimensions ",
+            arg1->nDims(),
+            ", ",
+            index->nDims(),
+            " and ",
+            src->nDims());
+        auto num_dims = (int64_t)arg1->nDims();
+        NVF_CHECK(
+            dim >= -num_dims && dim < num_dims,
+            "Tensor arguments have dimension ",
+            num_dims,
+            " so dim argument must satisfy ",
+            -num_dims,
+            " <= dim < ",
+            num_dims,
+            ", but received ",
+            dim);
+        return scatter(arg1, dim, index, src);
+      },
+      py::arg("arg1"),
+      py::arg("index"),
+      py::arg("src"),
+      py::arg("dim"),
+      R"(
+Scatter a tensor.
+
+Parameters
+----------
+arg1 : TensorView
+    The tensor to scatter into.
+index : TensorView
+    The tensor containing the indices.
+src : TensorView
+    The source tensor to scatter from.
+dim : int
+    The dimension to scatter along.
+
+Returns
+-------
+TensorView
+    The scattered tensor.
+)",
+      py::return_value_policy::reference);
 }
 
 template <class ShapeType>

python/python_direct/python_translate.cpp

@@ -1100,6 +1100,20 @@ class PythonTranslator : public OptInConstDispatch {
         {out_tv});
   }
 
+  // Map ScatterOp to python frontend
+  void handle(const ScatterOp* sop) final {
+    NVF_ERROR(sop != nullptr);
+    TensorView* out_tv = sop->output(0)->as<TensorView>();
+    visited_vals_.insert(out_tv);
+    static const std::vector<std::string> argument_names = {"dim"};
+    printer_.generateKwargsOperation(
+        "fd.ops.scatter",
+        std::make_tuple(sop->selfTv(), sop->indexTv(), sop->srcTv()),
+        argument_names,
+        std::make_tuple(sop->dim()),
+        {out_tv});
+  }
+
   // Map TopKOp to python frontend
   void handle(const TopKOp* topkop) final {
     NVF_ERROR(topkop != nullptr);

tests/python/opinfo/opinfos.py

@@ -1090,6 +1090,7 @@ def scatter_wrapper(
         ArgumentType.Symbolic,
         ArgumentType.Constant,
     ),
+    supports_direct_bindings=True,
 )
 shape_ops.append(scatter_opinfo)