Update aten_xla_type.cpp

pytorch · Jan 17, 2024 · f9618ea · f9618ea
1 parent 66660de
commit f9618ea
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/torch_xla/csrc/aten_xla_type.cpp b/torch_xla/csrc/aten_xla_type.cpp
@@ -2648,6 +2648,14 @@ at::Tensor scatter_reduce_helper(const at::Tensor& self, int64_t dim,
   }
 }
 
+at::Tensor XLANativeFunctions::scan(const at::Tensor& self,
+                                    const Callable f,
+                                    const at::Tensor& init,
+                                    const at::Tensor& xs) {
+  TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
+  return bridge::AtenFromXlaTensor(tensor_methods::scan(f, init, xs));
+}
+
 at::Tensor XLANativeFunctions::scatter(const at::Tensor& self, int64_t dim,
                                        const at::Tensor& index,
                                        const at::Tensor& src) {