import torch


def merge_inner_list(ys, dim=0):
  """Concatenate tensors in the inner list; 
    this is mainly used for construct the entire output from batch outputs;
    this is used in pred_func in ...vin.vin.py
    for ease of description, use integers to represent tensors, 
      and tuples to represent concatenated tensors in the following toy examples
      ys = [1,2,3] --> (1,2,3)
      ys = [[1,2,3], [4,5,6]] --> [(1,4), (2,5), (3,6)]
      ys = [[[1,2,3], [4,5]], [[6,7,8], [9,10]]] --> [[(1,6),(2,7),(3,8)], [(4,9),(5,10)]]

  Args:
    ys: a list of tensors, or a nested list of tensors
    dim: default 0 (for almost all the cases), passed for torch.cat

  Examples:
    ys = [[[torch.tensor([1]), torch.tensor([2]), torch.tensor([3])], 
            [torch.tensor([4]), torch.tensor([5])]], 
          [[torch.tensor([6]), torch.tensor([7]), torch.tensor([8])], 
            [torch.tensor([9]), torch.tensor([10])]]]
    merge_inner_list(ys)
    
  """
  if isinstance(ys[0], torch.Tensor):
    return torch.cat(ys, dim=dim)
  elif isinstance(ys[0], (list, tuple)):
    merged = list(map(list, zip(*ys)))
    return [merge_inner_list(s, dim=dim) for s in merged]