diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 1380431bc..7a88c6fce 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -44435,94 +44435,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 4, - "endColumn": 15, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 16, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 16, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 17, - "endColumn": 68, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 57, - "endColumn": 68, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 8, - "endColumn": 76, - "lineCount": 2 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 32, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 32, - "endColumn": 51, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 37, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 37, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 45, - "endColumn": 51, - "lineCount": 1 - } - }, { "code": "reportUnknownParameterType", "range": { @@ -93540,132 +93452,132 @@ } } ], - "./loopy/transform/privatize.py": [ + "./loopy/transform/realize_reduction.py": [ { "code": "reportUnknownParameterType", "range": { - "startColumn": 23, - "endColumn": 40, + "startColumn": 36, + "endColumn": 49, "lineCount": 1 } }, { "code": "reportMissingParameterType", "range": { - "startColumn": 23, - "endColumn": 40, + "startColumn": 36, + "endColumn": 49, "lineCount": 1 } }, { "code": "reportUnknownParameterType", "range": { - "startColumn": 42, - "endColumn": 57, + "startColumn": 51, + "endColumn": 61, "lineCount": 1 } }, { "code": "reportMissingParameterType", "range": { - "startColumn": 42, - "endColumn": 57, + "startColumn": 51, + "endColumn": 61, "lineCount": 1 } }, { - "code": "reportUnannotatedClassAttribute", + "code": "reportUnknownParameterType", "range": { - "startColumn": 13, - "endColumn": 30, + "startColumn": 12, + "endColumn": 24, "lineCount": 1 } }, { - "code": "reportUnannotatedClassAttribute", + "code": "reportMissingParameterType", "range": { - "startColumn": 13, - "endColumn": 28, + "startColumn": 12, + "endColumn": 24, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 8, - "endColumn": 34, + "startColumn": 31, + "endColumn": 41, "lineCount": 1 } }, { - "code": "reportUnannotatedClassAttribute", + "code": "reportMissingParameterType", "range": { - "startColumn": 13, - "endColumn": 34, + "startColumn": 31, + "endColumn": 41, "lineCount": 1 } }, { - "code": "reportImplicitOverride", + "code": "reportUnknownParameterType", "range": { - "startColumn": 8, - "endColumn": 21, + "startColumn": 31, + "endColumn": 46, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 12, - "endColumn": 38, + "startColumn": 31, + "endColumn": 46, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 12, - "endColumn": 45, + "startColumn": 48, + "endColumn": 54, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 46, - "endColumn": 52, + "startColumn": 48, + "endColumn": 54, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 36, - "endColumn": 46, - "lineCount": 2 + "startColumn": 25, + "endColumn": 62, + "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 44, - "endColumn": 76, + "startColumn": 25, + "endColumn": 47, "lineCount": 1 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 69, - "endColumn": 75, + "startColumn": 25, + "endColumn": 58, "lineCount": 1 } }, { - "code": "reportImplicitOverride", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 8, - "endColumn": 20, + "startColumn": 37, + "endColumn": 47, "lineCount": 1 } }, @@ -93673,191 +93585,191 @@ "code": "reportUnknownMemberType", "range": { "startColumn": 12, - "endColumn": 38, + "endColumn": 29, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 12, - "endColumn": 45, + "startColumn": 30, + "endColumn": 35, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 46, - "endColumn": 52, + "startColumn": 39, + "endColumn": 49, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 28, - "endColumn": 44, - "lineCount": 2 + "startColumn": 12, + "endColumn": 28, + "lineCount": 1 } }, { "code": "reportUnknownArgumentType", "range": { - "startColumn": 36, - "endColumn": 68, + "startColumn": 29, + "endColumn": 34, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 61, - "endColumn": 67, + "startColumn": 39, + "endColumn": 49, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 19, - "endColumn": 26, + "startColumn": 12, + "endColumn": 31, "lineCount": 1 } }, { - "code": "reportUnknownParameterType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 4, + "startColumn": 32, "endColumn": 37, "lineCount": 1 } }, { - "code": "reportUnknownParameterType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 8, - "endColumn": 14, + "startColumn": 12, + "endColumn": 29, "lineCount": 1 } }, { - "code": "reportMissingParameterType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 8, - "endColumn": 14, + "startColumn": 30, + "endColumn": 35, "lineCount": 1 } }, { - "code": "reportUnknownParameterType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 16, - "endColumn": 34, + "startColumn": 12, + "endColumn": 29, "lineCount": 1 } }, { - "code": "reportMissingParameterType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 16, - "endColumn": 34, + "startColumn": 18, + "endColumn": 28, "lineCount": 1 } }, { - "code": "reportUnknownParameterType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 36, - "endColumn": 50, + "startColumn": 31, + "endColumn": 47, "lineCount": 1 } }, { - "code": "reportMissingParameterType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 36, - "endColumn": 50, + "startColumn": 37, + "endColumn": 46, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 38, - "endColumn": 55, + "startColumn": 49, + "endColumn": 68, "lineCount": 1 } }, { "code": "reportUnknownArgumentType", "range": { - "startColumn": 31, - "endColumn": 45, + "startColumn": 55, + "endColumn": 67, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 50, - "endColumn": 75, + "startColumn": 38, + "endColumn": 44, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 49, - "endColumn": 60, + "startColumn": 38, + "endColumn": 44, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 11, - "endColumn": 28, + "startColumn": 46, + "endColumn": 59, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 14, - "endColumn": 40, + "startColumn": 46, + "endColumn": 59, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 14, - "endColumn": 47, + "startColumn": 30, + "endColumn": 43, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 42, - "endColumn": 49, - "lineCount": 1 + "startColumn": 12, + "endColumn": 50, + "lineCount": 2 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 30, - "endColumn": 38, + "startColumn": 32, + "endColumn": 63, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 39, - "endColumn": 46, + "startColumn": 32, + "endColumn": 65, "lineCount": 1 } }, @@ -93865,1377 +93777,287 @@ "code": "reportUnknownMemberType", "range": { "startColumn": 26, - "endColumn": 43, + "endColumn": 48, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 31, - "endColumn": 56, + "startColumn": 39, + "endColumn": 45, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 22, - "endColumn": 56, + "startColumn": 39, + "endColumn": 45, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 15, - "endColumn": 22, + "startColumn": 47, + "endColumn": 54, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportMissingParameterType", "range": { "startColumn": 47, - "endColumn": 82, + "endColumn": 54, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 74, - "endColumn": 81, + "startColumn": 56, + "endColumn": 65, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 47, - "endColumn": 54, + "startColumn": 56, + "endColumn": 65, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 42, - "endColumn": 62, + "startColumn": 22, + "endColumn": 45, "lineCount": 1 } }, { "code": "reportUnknownArgumentType", "range": { - "startColumn": 42, - "endColumn": 77, + "startColumn": 22, + "endColumn": 65, "lineCount": 1 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 69, - "endColumn": 76, + "startColumn": 14, + "endColumn": 31, "lineCount": 1 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 39, - "endColumn": 46, + "startColumn": 11, + "endColumn": 41, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 54, - "endColumn": 74, + "startColumn": 4, + "endColumn": 18, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 21, - "endColumn": 44, + "startColumn": 4, + "endColumn": 26, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 45, - "endColumn": 56, + "startColumn": 27, + "endColumn": 33, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportMissingParameterType", "range": { - "startColumn": 45, - "endColumn": 56, + "startColumn": 27, + "endColumn": 33, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 52, - "endColumn": 77, + "startColumn": 35, + "endColumn": 45, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportMissingParameterType", "range": { - "startColumn": 52, - "endColumn": 77, + "startColumn": 35, + "endColumn": 45, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 20, - "endColumn": 46, + "startColumn": 47, + "endColumn": 63, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 20, - "endColumn": 51, + "startColumn": 47, + "endColumn": 63, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 28, - "endColumn": 80, + "startColumn": 14, + "endColumn": 36, "lineCount": 1 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 16, - "endColumn": 24, + "startColumn": 31, + "endColumn": 50, "lineCount": 1 } }, { "code": "reportUnknownArgumentType", "range": { - "startColumn": 32, - "endColumn": 37, + "startColumn": 15, + "endColumn": 29, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownMemberType", "range": { - "startColumn": 45, - "endColumn": 56, + "startColumn": 8, + "endColumn": 29, "lineCount": 1 } }, { "code": "reportUnknownArgumentType", "range": { - "startColumn": 34, + "startColumn": 30, "endColumn": 40, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { "startColumn": 15, - "endColumn": 40, + "endColumn": 29, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportPossiblyUnboundVariable", "range": { - "startColumn": 29, - "endColumn": 34, + "startColumn": 37, + "endColumn": 56, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 23, - "endColumn": 38, - "lineCount": 1 + "startColumn": 16, + "endColumn": 46, + "lineCount": 2 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 56, - "endColumn": 67, + "startColumn": 27, + "endColumn": 33, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 22, - "endColumn": 29, + "startColumn": 27, + "endColumn": 33, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 33, - "endColumn": 40, + "startColumn": 35, + "endColumn": 45, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportMissingParameterType", "range": { - "startColumn": 32, - "endColumn": 37, + "startColumn": 35, + "endColumn": 45, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 16, - "endColumn": 35, + "startColumn": 47, + "endColumn": 58, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportMissingParameterType", "range": { - "startColumn": 39, - "endColumn": 57, + "startColumn": 47, + "endColumn": 58, "lineCount": 1 } }, { - "code": "reportUnknownArgumentType", + "code": "reportUnknownParameterType", "range": { - "startColumn": 39, - "endColumn": 54, + "startColumn": 60, + "endColumn": 77, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportMissingParameterType", "range": { - "startColumn": 19, - "endColumn": 52, + "startColumn": 60, + "endColumn": 77, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportUnusedParameter", "range": { - "startColumn": 15, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 45, - "endColumn": 63, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 23, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 36, - "endColumn": 43, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 24, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 24, - "endColumn": 71, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 53, - "endColumn": 71, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 8, - "endColumn": 24, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 25, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 11, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 23, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 23, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 41, - "endColumn": 55, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 41, - "endColumn": 55, - "lineCount": 1 - } - }, - { - "code": "reportUnannotatedClassAttribute", - "range": { - "startColumn": 13, - "endColumn": 27, - "lineCount": 1 - } - }, - { - "code": "reportUnannotatedClassAttribute", - "range": { - "startColumn": 13, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnannotatedClassAttribute", - "range": { - "startColumn": 13, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 20, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 23, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 37, - "endColumn": 68, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 16, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 48, - "endColumn": 57, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 75, - "endColumn": 84, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 4, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 8, - "endColumn": 14, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 8, - "endColumn": 14, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 16, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 16, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 36, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 36, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 38, - "endColumn": 55, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 31, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 50, - "endColumn": 75, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 49, - "endColumn": 60, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 23, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 43, - "endColumn": 57, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 8, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportUnknownLambdaType", - "range": { - "startColumn": 49, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 55, - "endColumn": 56, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 20, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 33, - "endColumn": 62, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 20, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 20, - "endColumn": 51, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 23, - "endColumn": 42, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 25, - "endColumn": 55, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 56, - "endColumn": 63, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 20, - "endColumn": 28, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 30, - "endColumn": 45, - "lineCount": 2 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 61, - "endColumn": 70, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 23, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportUnusedVariable", - "range": { - "startColumn": 32, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 59, - "endColumn": 71, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 24, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 34, - "endColumn": 45, - "lineCount": 2 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 65, - "endColumn": 78, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 27, - "endColumn": 42, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 37, - "endColumn": 45, - "lineCount": 2 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 68, - "endColumn": 84, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 33, - "endColumn": 40, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 11, - "endColumn": 22, - "lineCount": 1 - } - } - ], - "./loopy/transform/realize_reduction.py": [ - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 36, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 36, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 51, - "endColumn": 61, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 51, - "endColumn": 61, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 12, - "endColumn": 24, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 12, - "endColumn": 24, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 31, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 31, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 31, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 31, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 48, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 48, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 25, - "endColumn": 62, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 25, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 25, - "endColumn": 58, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 37, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 12, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 30, - "endColumn": 35, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 39, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 12, - "endColumn": 28, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 29, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 39, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 12, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 32, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 12, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 30, - "endColumn": 35, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 12, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 18, - "endColumn": 28, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 31, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 37, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 49, - "endColumn": 68, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 55, - "endColumn": 67, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 38, - "endColumn": 44, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 38, - "endColumn": 44, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 46, - "endColumn": 59, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 46, - "endColumn": 59, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 30, - "endColumn": 43, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 12, - "endColumn": 50, - "lineCount": 2 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 32, - "endColumn": 63, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 32, - "endColumn": 65, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 26, - "endColumn": 48, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 39, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 39, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 47, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 47, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 56, - "endColumn": 65, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 56, - "endColumn": 65, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 22, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 22, - "endColumn": 65, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 14, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 11, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 18, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 4, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 27, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 27, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 35, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 35, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 47, - "endColumn": 63, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 47, - "endColumn": 63, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 14, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 31, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 15, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 8, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 30, - "endColumn": 40, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 15, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportPossiblyUnboundVariable", - "range": { - "startColumn": 37, - "endColumn": 56, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 16, - "endColumn": 46, - "lineCount": 2 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 27, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 27, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 35, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 35, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 47, - "endColumn": 58, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 47, - "endColumn": 58, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 60, - "endColumn": 77, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 60, - "endColumn": 77, - "lineCount": 1 - } - }, - { - "code": "reportUnusedParameter", - "range": { - "startColumn": 60, - "endColumn": 77, + "startColumn": 60, + "endColumn": 77, "lineCount": 1 } }, @@ -114327,14 +113149,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 10, - "endColumn": 46, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -116995,14 +115809,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 10, - "endColumn": 46, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -117011,46 +115817,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 10, - "endColumn": 48, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 14, - "endColumn": 52, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 14, - "endColumn": 52, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 19, - "endColumn": 57, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 19, - "endColumn": 55, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 76c49ab5c..4e157be6f 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -304,7 +304,7 @@ def __init__(self, within_inames_is_final: bool | None, within_inames: frozenset[str] | None, priority: int | None, - predicates: frozenset[str] | None, + predicates: frozenset[str] | frozenset[Expression] | None, tags: frozenset[Tag] | None, *, depends_on: frozenset[str] | str | None = None, @@ -946,7 +946,7 @@ def __init__(self, within_inames_is_final: bool | None = None, within_inames: frozenset[str] | None = None, priority: int | None = None, - predicates: frozenset[str] | None = None, + predicates: frozenset[str] | frozenset[Expression] | None = None, tags: frozenset[Tag] | None = None, temp_var_type: type[_not_provided] diff --git a/loopy/preprocess.py b/loopy/preprocess.py index 6df059d98..463ed7883 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -454,7 +454,7 @@ def find_temporary_address_space(kernel): # {{{ realize_ilp -def realize_ilp(kernel): +def realize_ilp(kernel: LoopKernel) -> LoopKernel: logger.debug("%s: add axes to temporaries for ilp" % kernel.name) from loopy.kernel.data import IlpBaseTag, VectorizeTag, filter_iname_tags_by_type diff --git a/loopy/transform/privatize.py b/loopy/transform/privatize.py index e790a591f..bf4f9f45f 100644 --- a/loopy/transform/privatize.py +++ b/loopy/transform/privatize.py @@ -23,16 +23,26 @@ THE SOFTWARE. """ - import logging +from typing import TYPE_CHECKING + +from typing_extensions import override -import pymbolic import pymbolic.primitives as p from loopy.diagnostic import LoopyError from loopy.translation_unit import for_each_kernel +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from pymbolic.typing import ArithmeticExpression, Expression + + from loopy.kernel import LoopKernel + from loopy.kernel.instruction import InstructionBase + from loopy.typing import InameStr, InameStrSet + logger = logging.getLogger(__name__) @@ -51,13 +61,20 @@ class ExtraInameIndexInserter(IdentityMapper[[]]): - def __init__(self, var_to_new_inames, iname_to_lbound): + var_to_new_inames: Mapping[str, Sequence[p.Variable]] + iname_to_lbound: Mapping[str, ArithmeticExpression] + seen_priv_axis_inames: set[str] + + def __init__(self, + var_to_new_inames: Mapping[str, Sequence[p.Variable]], + iname_to_lbound: Mapping[str, ArithmeticExpression]) -> None: self.var_to_new_inames = var_to_new_inames self.iname_to_lbound = iname_to_lbound self.seen_priv_axis_inames = set() super().__init__() - def map_subscript(self, expr: p.Subscript): + @override + def map_subscript(self, expr: p.Subscript, /) -> Expression: assert isinstance(expr.aggregate, p.Variable) try: extra_idx = self.var_to_new_inames[expr.aggregate.name] @@ -71,14 +88,16 @@ def map_subscript(self, expr: p.Subscript): self.seen_priv_axis_inames.update(v.name for v in extra_idx) - new_idx = index + tuple(flatten(v - self.iname_to_lbound[v.name]) - for v in extra_idx) + new_idx = index + tuple( + flatten(v - self.iname_to_lbound[v.name]) for v in extra_idx + ) if len(new_idx) == 1: new_idx = new_idx[0] return expr.aggregate[new_idx] - def map_variable(self, expr: p.Variable): + @override + def map_variable(self, expr: p.Variable, /) -> Expression: try: new_idx = self.var_to_new_inames[expr.name] except KeyError: @@ -86,17 +105,19 @@ def map_variable(self, expr: p.Variable): else: self.seen_priv_axis_inames.update(v.name for v in new_idx) - new_idx = tuple(flatten(v - self.iname_to_lbound[v.name]) - for v in new_idx) - + new_idx = tuple(flatten(v - self.iname_to_lbound[v.name]) for v in new_idx) if len(new_idx) == 1: new_idx = new_idx[0] + return expr[new_idx] @for_each_kernel def privatize_temporaries_with_inames( - kernel, privatizing_inames, only_var_names=None): + kernel: LoopKernel, + privatizing_inames: InameStr | InameStrSet, + only_var_names: InameStr | InameStrSet | None = None, + ) -> LoopKernel: """This function provides each loop iteration of the *privatizing_inames* with its own private entry in the temporaries it accesses (possibly restricted to *only_var_names*). @@ -124,32 +145,32 @@ def privatize_temporaries_with_inames( end facilitating loop interchange of the *imatrix* loop. + .. versionadded:: 2018.1 """ if isinstance(privatizing_inames, str): privatizing_inames = frozenset( - s.strip() - for s in privatizing_inames.split(",")) + s.strip() for s in privatizing_inames.split(",") + ) if isinstance(only_var_names, str): only_var_names = frozenset( - s.strip() - for s in only_var_names.split(",")) + s.strip() for s in only_var_names.split(",") + ) # {{{ sanity checks if (only_var_names is not None and privatizing_inames <= kernel.all_inames() and not (frozenset(only_var_names) <= kernel.all_variable_names())): - raise LoopyError(f"Some variables in '{only_var_names}'" - f" not used in kernel '{kernel.name}'.") + raise LoopyError(f"some variables in '{only_var_names}'" + f" not used in kernel '{kernel.name}'") # }}} wmap = kernel.writer_map() - - var_to_new_priv_axis_iname = {} + var_to_new_priv_axis_iname: dict[str, frozenset[str]] = {} # {{{ find variables that need extra indices @@ -162,8 +183,8 @@ def privatize_temporaries_with_inames( priv_axis_inames = writer_insn.within_inames & privatizing_inames - referenced_priv_axis_inames = (priv_axis_inames - & writer_insn.write_dependency_names()) + referenced_priv_axis_inames = ( + priv_axis_inames & writer_insn.write_dependency_names()) new_priv_axis_inames = priv_axis_inames - referenced_priv_axis_inames @@ -171,18 +192,18 @@ def privatize_temporaries_with_inames( break if tv.name in var_to_new_priv_axis_iname: - if new_priv_axis_inames != set(var_to_new_priv_axis_iname[tv.name]): - raise LoopyError("instruction '%s' requires adding " - "indices for privatizing var '%s' on iname(s) '%s', " - "but previous instructions required different " - "inames '%s'" - % (writer_insn_id, tv.name, - ", ".join(new_priv_axis_inames), - ", ".join(var_to_new_priv_axis_iname[tv.name]))) + if new_priv_axis_inames != var_to_new_priv_axis_iname[tv.name]: + new_inames_str = ", ".join(new_priv_axis_inames) + prev_inames_str = ", ".join(var_to_new_priv_axis_iname[tv.name]) + raise LoopyError( + f"instruction '{writer_insn_id}' requires adding indices " + "for privatizing var '{tv.name}' on iname(s) " + f"'{new_inames_str}', but previous instructions required " + f"different inames '{prev_inames_str}'") continue - var_to_new_priv_axis_iname[tv.name] = set(new_priv_axis_inames) + var_to_new_priv_axis_iname[tv.name] = frozenset(new_priv_axis_inames) # }}} @@ -191,8 +212,8 @@ def privatize_temporaries_with_inames( from loopy.isl_helpers import static_max_of_pw_aff from loopy.symbolic import pw_aff_to_expr - priv_axis_iname_to_length = {} - iname_to_lbound = {} + priv_axis_iname_to_length: dict[str, ArithmeticExpression] = {} + iname_to_lbound: dict[str, ArithmeticExpression] = {} for priv_axis_inames in var_to_new_priv_axis_iname.values(): for iname in priv_axis_inames: if iname in priv_axis_iname_to_length: @@ -209,7 +230,7 @@ def privatize_temporaries_with_inames( from loopy.kernel.data import VectorizeTag - new_temp_vars = kernel.temporary_variables.copy() + new_temp_vars = dict(kernel.temporary_variables) for tv_name, inames in var_to_new_priv_axis_iname.items(): tv = new_temp_vars[tv_name] extra_shape = tuple(priv_axis_iname_to_length[iname] for iname in inames) @@ -218,10 +239,14 @@ def privatize_temporaries_with_inames( if shape is None: shape = () - dim_tags = ["c"] * (len(shape) + len(extra_shape)) + # NOTE: could be auto? + assert isinstance(shape, tuple) + ndim = len(shape) + + dim_tags = ["c"] * (ndim + len(extra_shape)) for i, iname in enumerate(inames): if kernel.iname_tags_of_type(iname, VectorizeTag): - dim_tags[len(shape) + i] = "vec" + dim_tags[ndim + i] = "vec" base_indices = tv.base_indices if base_indices is not None: @@ -229,20 +254,17 @@ def privatize_temporaries_with_inames( new_temp_vars[tv.name] = tv.copy(shape=shape + extra_shape, base_indices=base_indices, - # Forget what you knew about data layout, - # create from scratch. + # Forget what you knew about data layout, create from scratch. dim_tags=dim_tags, dim_names=None) # }}} - from pymbolic import var var_to_extra_iname = { - var_name: tuple(var(iname) for iname in inames) + var_name: tuple(p.Variable(iname) for iname in inames) for var_name, inames in var_to_new_priv_axis_iname.items()} - new_insns = [] - + new_insns: list[InstructionBase] = [] for insn in kernel.instructions: eiii = ExtraInameIndexInserter(var_to_extra_iname, iname_to_lbound) @@ -269,25 +291,34 @@ def privatize_temporaries_with_inames( # {{{ unprivatize temporaries with iname class _InameRemover(IdentityMapper[[bool]]): - def __init__(self, inames_to_remove, only_var_names): + only_var_names: frozenset[str] | None + inames_to_remove: frozenset[str] + var_name_to_remove_indices: dict[str, dict[int, str]] + + def __init__(self, + inames_to_remove: frozenset[str], + only_var_names: frozenset[str] | None) -> None: self.only_var_names = only_var_names self.inames_to_remove = inames_to_remove self.var_name_to_remove_indices = {} super().__init__() - def map_subscript(self, expr: p.Subscript, in_subscript: bool = False): + @override + def map_subscript(self, expr: p.Subscript, /, + in_subscript: bool = False) -> Expression: assert isinstance(expr.aggregate, p.Variable) name = expr.aggregate.name + if not self.only_var_names or name in self.only_var_names: index = expr.index if not isinstance(index, tuple): index = (index,) - remove_indices = {} - new_index = [] + remove_indices: dict[int, str] = {} + new_index: list[Expression] = [] for i, index_expr in enumerate(index): - if isinstance(index_expr, pymbolic.primitives.Variable) and \ - index_expr.name in self.inames_to_remove: + if (isinstance(index_expr, p.Variable) + and index_expr.name in self.inames_to_remove): remove_indices[i] = index_expr.name else: new_index.append(index_expr) @@ -303,8 +334,9 @@ def map_subscript(self, expr: p.Subscript, in_subscript: bool = False): self.var_name_to_remove_indices[name] = remove_indices if new_index: - new_index = new_index[0] if len(new_index) == 1 else tuple(new_index) - return expr.aggregate[new_index] + return expr.aggregate[ + new_index[0] if len(new_index) == 1 else tuple(new_index) + ] else: return expr.aggregate else: @@ -313,7 +345,9 @@ def map_subscript(self, expr: p.Subscript, in_subscript: bool = False): @for_each_kernel def unprivatize_temporaries_with_inames( - kernel, privatizing_inames, only_var_names=None): + kernel: LoopKernel, + privatizing_inames: InameStr | InameStrSet, + only_var_names: InameStr | InameStrSet | None = None) -> LoopKernel: """This function reverses the effects of :func:`privatize_temporaries_with_inames` and removes the private entries in the temporaries each loop iteration of the *privatizing_inames* @@ -342,13 +376,13 @@ def unprivatize_temporaries_with_inames( if isinstance(privatizing_inames, str): privatizing_inames = frozenset( - s.strip() - for s in privatizing_inames.split(",")) + s.strip() for s in privatizing_inames.split(",") + ) if isinstance(only_var_names, str): only_var_names = frozenset( - s.strip() - for s in only_var_names.split(",")) + s.strip() for s in only_var_names.split(",") + ) # {{{ sanity checks @@ -372,18 +406,20 @@ def unprivatize_temporaries_with_inames( from loopy.kernel.array import VectorArrayDimTag - new_temp_vars = kernel.temporary_variables.copy() + new_temp_vars = dict(kernel.temporary_variables) for tv_name, tv in new_temp_vars.items(): remove_indices = var_name_to_remove_indices.get(tv_name, {}) new_shape = tv.shape if new_shape is not None: - new_shape = tuple(dim for idim, dim in enumerate(new_shape) + assert isinstance(new_shape, tuple) + new_shape = tuple( + dim for idim, dim in enumerate(new_shape) if idim not in remove_indices) new_dim_tags = tv.dim_tags if new_dim_tags is not None: new_dim_tags = ["vec" if isinstance(dim_tag, VectorArrayDimTag) else "c" - for idim, dim_tag in enumerate(new_dim_tags)] + for _idim, dim_tag in enumerate(new_dim_tags)] new_dim_tags = tuple(dim for idim, dim in enumerate(new_dim_tags) if idim not in remove_indices)