diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp index 36c26f26c..ea5156a0c 100644 --- a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp +++ b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp @@ -1662,10 +1662,15 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ( auto shapeType = Torch::ValueTensorType::get( binder.op->getContext(), SmallVector{inputRank}, resultType.getOptionalDtype()); - Value shape = rewriter.create( binder.getLoc(), shapeType, operand); + if (inputRank == 0) { + rewriter.replaceOpWithNewOp( + binder.op, resultType, shape); + return success(); + } + if (start == 0 && end == -1) { rewriter.replaceOp(binder.op, shape); return success(); @@ -1673,18 +1678,13 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ( Value sv = rewriter.create( binder.getLoc(), rewriter.getI64IntegerAttr(start)); - Value ev = rewriter.create( binder.getLoc(), rewriter.getI64IntegerAttr(end)); - Value step = rewriter.create(binder.getLoc(), 1); - Value dim = rewriter.create(binder.getLoc(), 0); - shape = rewriter.create( - binder.getLoc(), resultType, shape, dim, sv, ev, step); - - rewriter.replaceOp(binder.op, shape); + rewriter.replaceOpWithNewOp( + binder.op, resultType, shape, dim, sv, ev, step); return success(); }); diff --git a/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir b/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir index af2a1e002..bd2a92874 100644 --- a/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir +++ b/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir @@ -2833,6 +2833,15 @@ func.func @test_shape_start_1_end_negative_1(%arg0: !torch.vtensor<[3,4,5],f32>) return %0 : !torch.vtensor<[1],si64> } +// ----- + +// CHECK-LABEL: func.func @test_shape_scalar +func.func @test_shape_scalar(%arg0: !torch.vtensor<[],si64> ) -> !torch.vtensor<[?],si64> attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.1.0"} { + // CHECK: %[[SHAPE:.+]] = torch.aten._shape_as_tensor %arg0 : !torch.vtensor<[],si64> -> !torch.vtensor<[0],si64> + // CHECK: %[[CAST:.+]] = torch.tensor_static_info_cast %[[SHAPE]] : !torch.vtensor<[0],si64> to !torch.vtensor<[?],si64> + %0 = torch.operator "onnx.Shape"(%arg0) : (!torch.vtensor<[],si64>) -> !torch.vtensor<[?],si64> + return %0: !torch.vtensor<[?],si64> +} // -----