| import torch |
|
|
| class NestedTensor: |
| def __init__(self, tensors): |
| self.tensors = list(tensors) |
| self.is_nested = True |
|
|
| def _copy(self): |
| return NestedTensor(self.tensors) |
|
|
| def apply_operation(self, other, operation): |
| o = self._copy() |
| if isinstance(other, NestedTensor): |
| for i, t in enumerate(o.tensors): |
| o.tensors[i] = operation(t, other.tensors[i]) |
| else: |
| for i, t in enumerate(o.tensors): |
| o.tensors[i] = operation(t, other) |
| return o |
|
|
| def __add__(self, b): |
| return self.apply_operation(b, lambda x, y: x + y) |
|
|
| def __sub__(self, b): |
| return self.apply_operation(b, lambda x, y: x - y) |
|
|
| def __mul__(self, b): |
| return self.apply_operation(b, lambda x, y: x * y) |
|
|
| |
| |
|
|
| def __truediv__(self, b): |
| return self.apply_operation(b, lambda x, y: x / y) |
|
|
| def __getitem__(self, *args, **kwargs): |
| return self.apply_operation(None, lambda x, y: x.__getitem__(*args, **kwargs)) |
|
|
| def unbind(self): |
| return self.tensors |
|
|
| def to(self, *args, **kwargs): |
| o = self._copy() |
| for i, t in enumerate(o.tensors): |
| o.tensors[i] = t.to(*args, **kwargs) |
| return o |
|
|
| def new_ones(self, *args, **kwargs): |
| return self.tensors[0].new_ones(*args, **kwargs) |
|
|
| def float(self): |
| return self.to(dtype=torch.float) |
|
|
| def chunk(self, *args, **kwargs): |
| return self.apply_operation(None, lambda x, y: x.chunk(*args, **kwargs)) |
|
|
| def size(self): |
| return self.tensors[0].size() |
|
|
| @property |
| def shape(self): |
| return self.tensors[0].shape |
|
|
| @property |
| def ndim(self): |
| dims = 0 |
| for t in self.tensors: |
| dims = max(t.ndim, dims) |
| return dims |
|
|
| @property |
| def device(self): |
| return self.tensors[0].device |
|
|
| @property |
| def dtype(self): |
| return self.tensors[0].dtype |
|
|
| @property |
| def layout(self): |
| return self.tensors[0].layout |
|
|
|
|
| def cat_nested(tensors, *args, **kwargs): |
| cated_tensors = [] |
| for i in range(len(tensors[0].tensors)): |
| tens = [] |
| for j in range(len(tensors)): |
| tens.append(tensors[j].tensors[i]) |
| cated_tensors.append(torch.cat(tens, *args, **kwargs)) |
| return NestedTensor(cated_tensors) |
|
|
