diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp --- a/mlir/lib/Bindings/Python/IRTypes.cpp +++ b/mlir/lib/Bindings/Python/IRTypes.cpp @@ -284,6 +284,19 @@ }, "Returns whether the given value is used as a placeholder for dynamic " "strides and offsets in shaped types."); + c.def_property_readonly( + "shape", + [](PyShapedType &self) { + self.requireHasRank(); + + std::vector shape; + int64_t rank = mlirShapedTypeGetRank(self); + shape.reserve(rank); + for (int64_t i = 0; i < rank; ++i) + shape.push_back(mlirShapedTypeGetDimSize(self, i)); + return shape; + }, + "Returns the shape of the ranked shaped type as a list of integers."); } private: diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py --- a/mlir/test/python/ir/builtin_types.py +++ b/mlir/test/python/ir/builtin_types.py @@ -315,6 +315,9 @@ # Encoding should be None. assert RankedTensorType.get(shape, f32).encoding is None + tensor = RankedTensorType.get(shape, f32) + assert tensor.shape == shape + # CHECK-LABEL: TEST: testUnrankedTensorType @run @@ -396,6 +399,8 @@ else: print("Exception not produced") + assert memref.shape == shape + # CHECK-LABEL: TEST: testUnrankedMemRefType @run