Given a source shape, a source layout and a destination layout, infer the destination shape. Parameter --------- src_shape : tuple of int or IntImm Source shape src_layout : str or Layout Source layout dst_layout : str or Layout Destination layout
(src_shape, src_layout, dst_layout)
| 404 | |
| 405 | |
| 406 | def get_shape(src_shape, src_layout, dst_layout): |
| 407 | """Given a source shape, a source layout and a destination layout, infer |
| 408 | the destination shape. |
| 409 | |
| 410 | Parameter |
| 411 | --------- |
| 412 | src_shape : tuple of int or IntImm |
| 413 | Source shape |
| 414 | |
| 415 | src_layout : str or Layout |
| 416 | Source layout |
| 417 | |
| 418 | dst_layout : str or Layout |
| 419 | Destination layout |
| 420 | |
| 421 | Returns |
| 422 | ------- |
| 423 | dst_shape : tuple of int |
| 424 | Destination shape |
| 425 | """ |
| 426 | if src_layout == dst_layout: |
| 427 | return get_const_tuple(src_shape) |
| 428 | |
| 429 | if isinstance(src_layout, str): |
| 430 | src_layout = slayout(src_layout) |
| 431 | if isinstance(dst_layout, str): |
| 432 | dst_layout = slayout(dst_layout) |
| 433 | |
| 434 | assert len(src_layout) == len(dst_layout), f"Incompatible layout {src_layout} vs {dst_layout}" |
| 435 | |
| 436 | layout_mapping = sbijective_layout(src_layout, dst_layout) |
| 437 | dst_indices = layout_mapping.forward_index(tvm.runtime.convert(list(range(len(src_layout))))) |
| 438 | |
| 439 | return get_const_tuple(tuple([src_shape[i.value] for i in dst_indices])) |
| 440 | |
| 441 | |
| 442 | def within_index(b, e, s, i): |
nothing calls this directly
no test coverage detected
searching dependent graphs…