diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -947,9 +947,15 @@
 }
 
 void ShapeOfOp::build(OpBuilder &builder, OperationState &result, Value arg) {
-  Type type = arg.getType().isa<ShapedType>()
-                  ? (Type)getExtentTensorType(builder.getContext())
-                  : (Type)builder.getType<ShapeType>();
+  Type type;
+  Type indexTy = builder.getIndexType();
+  if (auto shapedTy = arg.getType().dyn_cast<ShapedType>()) {
+    type = RankedTensorType::get(
+        {shapedTy.hasRank() ? shapedTy.getRank() : ShapedType::kDynamicSize},
+        indexTy);
+  } else {
+    type = builder.getType<ShapeType>();
+  }
   return ShapeOfOp::build(builder, result, type, arg);
 }