FilesExpand file tree

master

/

nested_tensor.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

97 lines (73 loc) · 2.52 KB

master

/

nested_tensor.py

Copy path

Top

File metadata and controls

• Code
• Blame

97 lines (73 loc) · 2.52 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

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 __itruediv__(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 cpu(self):

return self.to(device="cpu")

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 __repr__(self):

return f"{type(self).__name__}({self.tensors!r})"

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)