| 93 | |
| 94 | @register_lower_rule(mops.Dimshuffle) |
| 95 | def dim_shuffle_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 96 | assert len(args) == 1 and len(ctx.vars_in) == 1 and len(ctx.vars_out) == 1 |
| 97 | # mge dimshuffle can do transpose and broadcast simutaneously |
| 98 | # for example: |
| 99 | # case1: (16, 32, 64) with pattern [0, 2, 1] -> (16, 64, 32) |
| 100 | # case2: (16, 32, 64) with pattern [0, -1, 2, -1, 1] -> (16, 1, 64, 1, 32) |
| 101 | # case3: (16, 1, 64, 1, 32) with pattern [0, 4, 2] -> (16, 32, 64) |
| 102 | |
| 103 | pattern = ctx.op.pattern |
| 104 | inp = args[0] |
| 105 | if len(pattern) == inp.ndim: |
| 106 | permutation = pattern |
| 107 | return transpose(inp, permutation) |
| 108 | elif len(pattern) > inp.ndim: |
| 109 | permutation = [item for item in pattern if item != -1] |
| 110 | return transpose(inp, permutation).reshape(ctx.vars_out[0].shape) |
| 111 | else: |
| 112 | permutation = [i for i in range(inp.ndim) if i not in pattern] + list(pattern) |
| 113 | return transpose(inp, permutation).reshape(ctx.vars_out[0].shape) |
| 114 | |
| 115 | |
| 116 | def concat(inps, axis): |