Merge branch 'dliddell-aten-index-select' of https://github.com/davel…

…iddell/torch-mlir into dliddell-aten-index-select
daveliddell · Feb 6, 2024 · 6d14771 · 6d14771
2 parents ce873e4 + f0a3f1a
commit 6d14771
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp
@@ -2843,20 +2843,17 @@ OpFoldResult AtenIndexSelectOp::fold(FoldAdaptor adaptor) {
   auto self = getSelf();
   auto index = getIndex();
   auto selfTy = cast<ValueTensorType>(self.getType());
-  auto indexTy = cast<ValueTensorType>(index.getType());
+  assert(index.getType().isa<IntegerType>());
   auto resultTy = cast<ValueTensorType>(getType());
 
   auto selfSizes = selfTy.getSizes();
-  auto indexSizes = indexTy.getSizes();
   auto resultSizes = resultTy.getSizes();
 
 
   if (selfTy.getDtype() != resultTy.getDtype())
     return nullptr;
   if (selfSizes.size() != resultSizes.size())
     return nullptr;
-  if (indexSizes.size() != 1)
-    return nullptr;
 
   // If the selection results in a tensor of the same dimensions as the
   // input, the selection must have specified every index of the input,
@@ -2869,7 +2866,7 @@ OpFoldResult AtenIndexSelectOp::fold(FoldAdaptor adaptor) {
     fullTensor &= resultSizes[i] != Torch::kUnknownSize;
   }
 
-  if (fullTensor && indexSizes[0] == 1)
+  if (fullTensor)
     return self;
 
   // If the input tensor, index dimension, or indexes are non-constant,