Stacks and flattens first dim of inputs at indices `input_indices`.
(pfor_input, input_indices)
| 1474 | |
| 1475 | |
| 1476 | def _inputs_with_flattening(pfor_input, input_indices): |
| 1477 | """Stacks and flattens first dim of inputs at indices `input_indices`.""" |
| 1478 | if input_indices is None: |
| 1479 | input_indices = [] |
| 1480 | pfor_input.stack_inputs(stack_indices=input_indices) |
| 1481 | inputs = [] |
| 1482 | for i in range(pfor_input.num_inputs): |
| 1483 | if i in input_indices: |
| 1484 | inp = pfor_input.stacked_input(i) |
| 1485 | inp = _flatten_first_two_dims(inp) |
| 1486 | else: |
| 1487 | inp = pfor_input.unstacked_input(i) |
| 1488 | inputs.append(inp) |
| 1489 | return inputs |
| 1490 | |
| 1491 | |
| 1492 | @RegisterPForWithArgs("Conv2D", dims=[0]) |
no test coverage detected