import torch


def nested_stack(params, roll: bool = False):
    """Form a tensor from a nexted list of tensors.

    This function is a generalization of torch.stack. For proper usage,
    it's important that params is a nested list with shape consistent with
    and array. The innermost elements of that nested list should be PyTorch
    tensors, all of which have identical size.

    For an example, suppose that a, b, c, and d are all tensors of size (5,).
    Then, nested_stack([[a, b], [c, d]]) returns a tensor of size (2, 2, 5).

    If roll is set to True, then the dimensions of the tensors (like a, b, c
    and d in the example above) will be permuted to the start of the output.
    This is useful if those dimensions were supposed to be batch dimensions.
    In the example, the output with roll=True would have size (5, 2, 2).
    If instead a, b, c, and d all had size (6, 9, 8), then the output size
    would be (6, 9, 8, 2, 2) if roll=True and (2, 2, 6, 9, 8) if roll=False.
    """
    def recursive_stack(params_):
        if isinstance(params_[0], torch.Tensor):
            return torch.stack(params_)
        num_rows = len(params_)
        return torch.stack(
            [nested_stack(params_[i]) for i in range(num_rows)]
        )
    stacked = recursive_stack(params).squeeze(0)

    if roll:
        inner = params[0]
        while not isinstance(inner, torch.Tensor):
            inner = inner[0]
        inner_dim = inner.dim()
        perm = list(range(stacked.dim()-inner_dim, stacked.dim())) + list(range(stacked.dim()-inner_dim))
        return stacked.permute(perm)
    else:
        return stacked