/
transpose_sequence.py
119 lines (94 loc) · 3.49 KB
/
transpose_sequence.py
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import numpy
from chainer.backends import cuda
from chainer import function_node
from chainer.utils import type_check
def _transpose(xs, length):
if length == 0:
return ()
xp = cuda.get_array_module(*xs)
lengths = numpy.empty(length, dtype='i')
end = length
for i, x in enumerate(xs):
len_x = len(x)
if len_x == end:
continue
lengths[len_x:end] = i
end = len_x
lengths[0:end] = len(xs)
if xp is numpy:
dtype = xs[0].dtype
unit = xs[0].shape[1:]
outs = tuple([xp.empty((l,) + unit, dtype=dtype) for l in lengths])
for i, x in enumerate(xs):
for p, xi in enumerate(x):
outs[p][i] = xi
else:
offsets1 = numpy.empty(len(xs) + 1, dtype='i')
offsets1[0] = 0
numpy.cumsum([len(x) for x in xs], out=offsets1[1:])
offsets2 = numpy.empty(length + 1, dtype='i')
offsets2[0] = 0
numpy.cumsum(lengths, dtype='i', out=offsets2[1:])
x = xp.concatenate(xs, axis=0)
o = xp.empty_like(x)
unit = xs[0].size // len(xs[0])
size = length * len(xs) * unit
cuda.elementwise(
'int32 len, int32 unit, raw int32 off1, raw int32 off2, raw T vs',
'raw T hs',
'''
int ind = i / unit;
int off = i - ind * unit;
int y = ind / len;
int x = ind - y * len;
if (off2[x] + y < off2[x + 1]) {
hs[(off2[x] + y) * unit + off] = vs[(off1[y] + x) * unit + off];
}
''',
'transpose_sequence'
)(length, unit, cuda.to_gpu(offsets1), cuda.to_gpu(offsets2), x, o,
size=size)
outs = tuple(xp.split(o, offsets2[1:-1]))
return outs
class TransposeSequence(function_node.FunctionNode):
"""Function that transposes a list of Variables."""
def __init__(self, length):
self._length = length
def check_type_forward(self, xs_type):
for p, n in zip(xs_type, xs_type[1:]):
type_check.expect(
p.shape[0] >= n.shape[0],
p.shape[1:] == n.shape[1:],
)
def forward(self, xs):
if len(xs) == 0:
return ()
return _transpose(xs, self._length)
def backward(self, indexes, grad_outputs):
return TransposeSequence(len(self.inputs)).apply(grad_outputs)
def transpose_sequence(xs):
"""Transpose a list of Variables.
This function transposes a list of :class:`~chainer.Variable`\\ s and
returns a list of :class:`Variable`\\ s.
For example a user gives ``[(0, 1, 2, 3), (4, 5), (6)]``, the function
returns ``[(0, 4, 6), (1, 5), (2), (3)]``.
Note that a given list needs to be sorted by each length of
:class:`~chainer.Variable`.
Args:
xs (list of :class:`~chainer.Variable` or :class:`numpy.ndarray` or \
:class:`cupy.ndarray`): Variables to transpose.
Returns:
tuple of :class:`~chainer.Variable`: Transposed list.
.. admonition:: Example
>>> lst = [chainer.Variable(np.array([1, 1, 1])),
... chainer.Variable(np.array([2, 2])),
... chainer.Variable(np.array([3]))]
>>> lst
[variable([1, 1, 1]), variable([2, 2]), variable([3])]
>>> transposed = F.transpose_sequence(lst)
>>> transposed
(variable([1, 2, 3]), variable([1, 2]), variable([1]))
"""
if len(xs) == 0:
return ()
return TransposeSequence(len(xs[0])).apply(xs)