import torch
from typing import Sequence


def _flatten(x):
    """Flattens the tree of tensors to flattened sequence of tensors"""
    if isinstance(x, torch.Tensor):
        return [x]
    if isinstance(x, Sequence):
        res = []
        for xi in x:
            res += _flatten(xi)
        return res
    return []