diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp --- a/mlir/lib/Bindings/Python/IRAttributes.cpp +++ b/mlir/lib/Bindings/Python/IRAttributes.cpp @@ -710,6 +710,9 @@ // f16 return bufferInfo(shapedType, "e"); } + if (mlirTypeIsAIndex(elementType)) { + return bufferInfo(shapedType); + } if (mlirTypeIsAInteger(elementType) && mlirIntegerTypeGetWidth(elementType) == 32) { if (mlirIntegerTypeIsSignless(elementType) || diff --git a/mlir/test/python/ir/array_attributes.py b/mlir/test/python/ir/array_attributes.py --- a/mlir/test/python/ir/array_attributes.py +++ b/mlir/test/python/ir/array_attributes.py @@ -365,3 +365,20 @@ # CHECK: {{\[}}4 5 6]] print(np.array(attr)) + +# CHECK-LABEL: TEST: testGetDenseElementsIndex +@run +def testGetDenseElementsIndex(): + with Context(), Location.unknown(): + idx_type = IndexType.get() + array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uintp) + attr = DenseElementsAttr.get(array, type=idx_type) + # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xindex> + print(attr) + arr = np.array(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(arr) + # CHECK: True + print(arr.dtype == np.int64) +