Added tensorResultTypeAtIndex to Patterns.h

Need this for LayerNorm
2023-12-29 11:21:55 -08:00 · 2023-12-29 11:21:55 -08:00 · 80bd093d56
parent 9adad9bc40
commit 80bd093d56
1 changed files with 10 additions and 0 deletions
--- a/include/torch-mlir/Conversion/TorchOnnxToTorch/Patterns.h
+++ b/include/torch-mlir/Conversion/TorchOnnxToTorch/Patterns.h
@ -95,6 +95,16 @@ struct OpBinder {
    return success();
  }

+  ParseResult tensorResultTypeAtIndex(Torch::ValueTensorType &typeIdx, int64_t idx) {
+    if (idx >= op->getNumResults())
+      return failure();
+    auto t = toValidTensorType(op->getResult(idx).getType());
+    if (!t)
+      return failure();
+    typeIdx = t;
+    return success();
+  }
+
  // Attribute accessors.
  ParseResult s64BoolAttr(bool &value, StringRef nameSuffix,
                          bool defaultValue = false) {