(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 174 | |
| 175 | @register_lower_rule(mops.Split, "Split") |
| 176 | def split_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 177 | nr_inp, nr_oup = len(ctx.vars_in), len(ctx.vars_out) |
| 178 | assert len(args) == nr_inp and nr_inp == nr_oup + 1 and len(args) > 1 |
| 179 | |
| 180 | assert isinstance(ctx.param["axis"], int) |
| 181 | axis = ctx.param["axis"] |
| 182 | |
| 183 | sections = [] |
| 184 | for i in range(nr_oup): |
| 185 | section = ctx.vars_out[i].shape[axis] |
| 186 | if ctx.vars_in[i + 1].bound_data is not None: |
| 187 | assert section == _parse_var_as_value(ctx.vars_in[i + 1]) |
| 188 | sections.append(section) |
| 189 | |
| 190 | return split(args[0], sections, axis) |
| 191 | |
| 192 | |
| 193 | def fill(value, shape, dtype=np.float32): |
nothing calls this directly
no test coverage detected