Function parse_groups

monai/networks/trt_compiler.py:255–302

Implements parsing of 'output_lists' arg of trt_compile(). Args: ret: plain list of Tensors output_lists: list of output group sizes: to form some Lists/Tuples out of 'ret' List, this will be a list of group dimensions, like [[], [5], [-1]] for returning Te

(
    ret: list[torch.Tensor], output_lists: list[list[int]]
)

def parse_groups(
    ret: list[torch.Tensor], output_lists: list[list[int]]
) -> tuple[torch.Tensor | list[torch.Tensor], ...]:
    """
    Implements parsing of 'output_lists' arg of trt_compile().

    Args:
      ret: plain list of Tensors

      output_lists: list of output group sizes: to form some Lists/Tuples out of 'ret' List, this will be a list
                    of group dimensions, like [[], [5], [-1]] for returning Tensor, list of 5 items and dynamic list.
        Format: [[group_n] | [], ...]
          [] or group_n == 0 : next output from ret is a scalar
          group_n > 0  :       next output from ret is a list of group_n length
          group_n == -1:       next output is a dynamic list. This entry can be at any
                               position in output_lists, but can appear only once.
    Returns:
       Tuple of Union[torch.Tensor, List[torch.Tensor]], according to the grouping in output_lists

    """
    groups: tuple[torch.Tensor | list[torch.Tensor], ...] = ()
    cur = 0
    for idx in range(len(output_lists)):
        gl = output_lists[idx]
        assert len(gl) == 0 or len(gl) == 1
        if len(gl) == 0 or gl[0] == 0:
            groups = (*groups, ret[cur])
            cur = cur + 1
        elif gl[0] > 0:
            groups = (*groups, ret[cur : cur + gl[0]])
            cur = cur + gl[0]
        elif gl[0] == -1:
            rev_groups: tuple[torch.Tensor | list[torch.Tensor], ...] = ()
            rcur = len(ret)
            for rl in range(len(output_lists) - 1, idx, -1):
                rgl = output_lists[rl]
                assert len(rgl) == 0 or len(rgl) == 1
                if len(rgl) == 0 or rgl[0] == 0:
                    rcur = rcur - 1
                    rev_groups = (*rev_groups, ret[rcur])
                elif rgl[0] > 0:
                    rcur = rcur - rgl[0]
                    rev_groups = (*rev_groups, ret[rcur : rcur + rgl[0]])
                else:
                    raise ValueError("Two -1 lists in output")
            groups = (*groups, ret[cur:rcur], *rev_groups[::-1])
            break
    return groups


class TrtCompiler: