vinx13 commented on code in PR #12720: URL: https://github.com/apache/tvm/pull/12720#discussion_r972234799
########## python/tvm/tir/function.py: ########## @@ -389,17 +389,27 @@ def from_func_with_separators(mapping_function: Callable, ndim: Optional[int] = final_indices = [] axis_separators = [] - for val in mapping: - if isinstance(val, tvm.ir.PrimExpr): - final_indices.append(val) - elif val is IndexMap.AXIS_SEPARATOR: - axis_separators.append(len(final_indices)) - else: - raise TypeError( - "Expected mapping function to return list of " - "either tvm.ir.PrimExpr or IndexMap.AXIS_SEPARATOR. " - f"Instead received {val} of type {type(val)}." - ) + + try: + iter(mapping) Review Comment: What's the use case for this? According to the doc the mapping function should return a List, it might also need update ########## python/tvm/tir/schedule/schedule.py: ########## @@ -2479,6 +2480,31 @@ def transform_layout( primitive will be called in addition to the TransformLayout primitive. + pad_value: Optional[Union[int, float, PrimExpr, IndexMap, Callable]] Review Comment: Document the assumption when pad_value is IndexMap. I remember in the RFC we assume it should contain no BufferLoad from buffers except the current buffer -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org