File tree Expand file tree Collapse file tree 1 file changed +2
-9
lines changed Expand file tree Collapse file tree 1 file changed +2
-9
lines changed Original file line number Diff line number Diff line change @@ -61,15 +61,8 @@ def get_problem_size(self) -> ProblemSize:
6161 contraction_dims = linalg .infer_contraction_dimensions (root_op )
6262 assert contraction_dims , "no contraction dimensions"
6363
64- # TODO(Bangtian): Expose Python bindings for getting indexing maps.
65- indexing_maps_attr = None
66- for attr in root_op .opview .attributes :
67- if attr .name == "indexing_maps" and isinstance (attr .attr , ir .ArrayAttr ):
68- indexing_maps_attr = attr .attr
69- break
70-
71- assert indexing_maps_attr , "indexing_maps attribute not found"
72- maps = [attr .value for attr in indexing_maps_attr ]
64+ res_maps = linalg .get_indexing_maps (root_op )
65+ maps = [map_attr .value for map_attr in res_maps ]
7366 lhs_dims = get_map_result_dim_positions (maps [0 ])
7467 rhs_dims = get_map_result_dim_positions (maps [1 ])
7568 res_dims = get_map_result_dim_positions (maps [2 ])
You can’t perform that action at this time.
0 commit comments