mirror of https://github.com/llvm/torch-mlir
[torch] Add edgecase for aten.shape_to_tensor for rank-0 input (#2962)
Currently lowering uses `tensor.from_elements` which does not allow zero inputs. In this case we return a `tensor.empty` operation.pull/2958/head
parent
08bc013fcd
commit
dd673cfa8d
|
@ -84,6 +84,12 @@ public:
|
|||
getTypeConverter()->convertType(op.getType()).cast<RankedTensorType>();
|
||||
|
||||
int64_t rank = operandTy.getRank();
|
||||
if (rank == 0) {
|
||||
rewriter.replaceOpWithNewOp<tensor::EmptyOp>(op, resultTy.getShape(),
|
||||
resultTy.getElementType());
|
||||
return success();
|
||||
}
|
||||
|
||||
SmallVector<Value> dims;
|
||||
for (int i = 0; i < rank; ++i) {
|
||||
Value dim = rewriter.createOrFold<tensor::DimOp>(loc, operand, i);
|
||||
|
|
Loading…
Reference in New Issue