(self)
| 124 | return layout if not layout else layout[num_expand:] |
| 125 | |
| 126 | def unfold_input(self) -> ArgData: |
| 127 | input_desc = self.input_data.desc |
| 128 | num_expand = len(input_desc.expandable_prefix) |
| 129 | unfolded_input = unfold_batches(self.input_data.data, num_expand) |
| 130 | if input_desc.layout: |
| 131 | unfolded_layout = input_desc.layout[num_expand:] |
| 132 | else: |
| 133 | unfolded_layout = input_desc.layout |
| 134 | self.unfolded_input = ArgData( |
| 135 | desc=ArgDesc(input_desc.name, "", input_desc.dest_device, unfolded_layout), |
| 136 | data=unfolded_input, |
| 137 | ) |
| 138 | return self.unfolded_input |
| 139 | |
| 140 | def compute_params(self) -> List[ArgData]: |
| 141 | raise NotImplementedError |
no test coverage detected