From 432a9912a35ad9d34e12e73407952edc02c047a1 Mon Sep 17 00:00:00 2001
From: Wei-Sheng Chin <wschin@outlook.com>
Date: Thu, 19 Jan 2023 16:02:40 -0800
Subject: [PATCH] Fix LORT CI failure due to PyTorch change (#14367)

As title. The fuser in LORT doesn't like "scalar". With a recent PyTorch
change, scalar is intorduced somewhere it was there before. Now, a
simple fix is to check if all inputs are tensors or some specially
allowed cases before sending ops to ORT.
---
 .../orttraining/lazy_tensor/accelerator.cc    | 35 ++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/orttraining/orttraining/lazy_tensor/accelerator.cc b/orttraining/orttraining/lazy_tensor/accelerator.cc
index f7ecbd000d..05102b74c1 100644
--- a/orttraining/orttraining/lazy_tensor/accelerator.cc
+++ b/orttraining/orttraining/lazy_tensor/accelerator.cc
@@ -36,6 +36,39 @@ namespace py = pybind11;
 namespace aten = torch::jit::aten;
 namespace prim = torch::jit::prim;
 
+bool IsFusable(const torch::jit::Node* node) {
+  // This function checks the fusion restriction inside
+  // mergeNodeIntoGroup(...). When selecting onnx-supported aten ops,
+  // we need to call this function to make sure they are fusable
+  // in mergeNodeIntoGroup(...).
+
+  // Not all inputs are fusable. For example, the fuser
+  // expects all inputs are tensors with a few exceptions.
+  // This flag denotes if we find an unsupported input.
+  bool found_not_fusable = false;
+  for (auto input : node->inputs()) {
+    if (input->type()->isSubtypeOf(*c10::TensorType::get())) {
+      continue;
+    } else if (
+        (input->type()->isSubtypeOf(*c10::FloatType::get()) &&
+         input->node()->kind() != torch::jit::prim::Constant) ||
+        (node->kind() == torch::jit::aten::_grad_sum_to_size &&
+         input->type()->isSubtypeOf(*c10::ListType::ofInts()))) {
+      continue;
+    } else if (
+        input->type()->isSubtypeOf(*c10::IntType::get()) &&
+        input->node()->kind() != torch::jit::prim::Constant) {
+      continue;
+    } else{
+      if (input->node()->kind() == torch::jit::prim::Constant) {
+        continue;
+      }
+      found_not_fusable = true;
+    }
+  }
+  return !found_not_fusable;
+}
+
 bool Accelerator::Supported(const torch::jit::Node* node) {
   if (!node) {
     return false;
@@ -74,7 +107,7 @@ bool Accelerator::Supported(const torch::jit::Node* node) {
         std::cout << "Supported op: "
                   << ToString(*node) << std::endl;
       }
-      return true;
+      return IsFusable(node);
     }
     default: {
       if (DumpAtenOpHistory()) {