1from __future__ import absolute_import, print_function, division
2import theano
3from theano.gradient import DisconnectedType
4from theano.gof import Op, Apply, TopoOptimizer
5from theano.gof.opt import copy_stack_trace
6from theano import tensor
7
8
9def get_diagonal_subtensor_view(x, i0, i1):
10    """
11    Helper function for DiagonalSubtensor and IncDiagonalSubtensor.
12
13    Notes
14    -----
15    It returns a partial view of x, not a partial copy.
16
17    """
18    # We have to cast i0 and i0 to int because python
19    # do not support indexing with 0-dim, 'int*' ndarrays.
20    i0 = int(i0)
21    i1 = int(i1)
22    if x.shape[i0] < x.shape[i1]:
23        raise NotImplementedError('is this allowed?')
24    idx = [slice(None)] * x.ndim
25    idx[i0] = slice(x.shape[i1] - 1, None, None)
26    xview = x.__getitem__(tuple(idx))
27    strides = list(xview.strides)
28    if x.shape[i1] != 1:
29        strides[i1] -= strides[i0]
30        xview.strides = strides
31    return xview
32
33
34class DiagonalSubtensor(Op):
35    """
36    Return a form a nd diagonal subtensor.
37
38    Parameters
39    ----------
40    x
41        n-d tensor
42    i0
43        Axis index in x
44    i1
45        Axis index in x
46
47    Notes
48    -----
49    Work on the GPU.
50
51    Extended summary
52    ----------------
53    ``x`` is some n-dimensional tensor, but this Op only deals with a
54    matrix-shaped slice, using axes i0 and i1. Without loss of
55    generality, suppose that ``i0`` picks out our ``row`` dimension,
56    and i1 the ``column`` dimension.
57
58    So the relevant part of ``x`` is some matrix ``u``. Suppose it has 7 rows
59    and 4 columns::
60
61        [ 0 0 0 0 ]
62        [ 0 0 0 0 ]
63        [ 0 0 0 0 ]
64        [ 0 0 0 0 ]
65        [ 0 0 0 0 ]
66        [ 0 0 0 0 ]
67
68    The view returned by this function is also a matrix. It's a thick,
69    diagonal ``stripe`` across u that discards the lower left triangle
70    and the upper right triangle:
71
72        [ x 0 0 0 ]
73        [ x x 0 0 ]
74        [ x x x 0 ]
75        [ 0 x x x ]
76        [ 0 0 x x ]
77        [ 0 0 0 x ]
78
79    In this case the return value would be this view of shape 3x4. The
80    returned view has the same number of dimensions as the input
81    ``x``, and the only difference is that the shape along dimension
82    ``i0`` has been reduced by ``shape[i1] - 1`` because of the
83    triangles that got chopped out.
84
85    The NotImplementedError is meant to catch the case where shape[i0]
86    is too small for the stripe to reach across the matrix, in which
87    case it's not clear what this function should do. Maybe always
88    raise an error. I'd look back to the call site in the Conv3D to
89    see what's necessary at that point.
90
91    """
92
93    __props__ = ("inplace",)
94
95    def __str__(self):
96        if self.inplace:
97            return "%s{inplace}" % self.__class__.__name__
98        return "%s" % self.__class__.__name__
99
100    def __init__(self, inplace=False):
101        self.inplace = inplace
102        if inplace:
103            self.view_map = {0: [0]}
104
105    def make_node(self, x, i0, i1):
106        _i0 = tensor.as_tensor_variable(i0)
107        _i1 = tensor.as_tensor_variable(i1)
108        return Apply(self, [x, _i0, _i1], [x.type()])
109
110    def perform(self, node, inputs, output_storage):
111        xview = get_diagonal_subtensor_view(*inputs)
112        if self.inplace:
113            output_storage[0][0] = xview
114        else:
115            output_storage[0][0] = xview.copy()
116
117    def grad(self, inputs, g_outputs):
118        z = tensor.zeros_like(inputs[0])
119        gx = inc_diagonal_subtensor(z, inputs[1], inputs[2], g_outputs[0])
120        return [gx, DisconnectedType()(), DisconnectedType()()]
121
122    def connection_pattern(self, node):
123        rval = [[True], [False], [False]]
124        return rval
125
126diagonal_subtensor = DiagonalSubtensor(False)
127
128
129class IncDiagonalSubtensor(Op):
130    """
131    The gradient of DiagonalSubtensor.
132
133    """
134
135    __props__ = ("inplace",)
136
137    def __str__(self):
138        if self.inplace:
139            return "%s{inplace}" % self.__class__.__name__
140        return "%s" % self.__class__.__name__
141
142    def __init__(self, inplace=False):
143        self.inplace = inplace
144        if inplace:
145            self.destroy_map = {0: [0]}
146
147    def make_node(self, x, i0, i1, amt):
148        _i0 = tensor.as_tensor_variable(i0)
149        _i1 = tensor.as_tensor_variable(i1)
150        return Apply(self, [x, _i0, _i1, amt], [x.type()])
151
152    def perform(self, node, inputs, output_storage):
153        x, i0, i1, amt = inputs
154        if not self.inplace:
155            x = x.copy()
156        xview = get_diagonal_subtensor_view(x, i0, i1)
157        xview += amt
158        output_storage[0][0] = x
159
160    def grad(self, inputs, g_outputs):
161        x, i0, i1, amt = inputs
162        gy = g_outputs[0]
163        return [gy, DisconnectedType()(), DisconnectedType()(),
164                diagonal_subtensor(gy, i0, i1)]
165
166    def connection_pattern(self, node):
167        rval = [[True], [False], [False], [True]]
168        return rval
169inc_diagonal_subtensor = IncDiagonalSubtensor(False)
170
171
172def conv3d(signals, filters,
173           signals_shape=None, filters_shape=None,
174           border_mode='valid'):
175    """
176    Convolve spatio-temporal filters with a movie.
177
178    It flips the filters.
179
180    Parameters
181    ----------
182    signals
183        Timeseries of images whose pixels have color channels.
184        Shape: [Ns, Ts, C, Hs, Ws].
185    filters
186        Spatio-temporal filters.
187        Shape: [Nf, Tf, C, Hf, Wf].
188    signals_shape
189        None or a tuple/list with the shape of signals.
190    filters_shape
191        None or a tuple/list with the shape of filters.
192    border_mode
193        One of 'valid', 'full' or 'half'.
194
195    Notes
196    -----
197    Another way to define signals: (batch,  time, in channel, row, column)
198    Another way to define filters: (out channel,time,in channel, row, column)
199
200    For the GPU, use nnet.conv3d.
201
202    See Also
203    --------
204    Someone made a script that shows how to swap the axes between
205    both 3d convolution implementations in Theano. See the last
206    `attachment <https://groups.google.com/d/msg/theano-users/1S9_bZgHxVw/0cQR9a4riFUJ>`_
207
208    """
209
210    if isinstance(border_mode, str):
211        border_mode = (border_mode, border_mode, border_mode)
212
213    if signals_shape is None:
214        _signals_shape_5d = signals.shape
215    else:
216        _signals_shape_5d = signals_shape
217
218    if filters_shape is None:
219        _filters_shape_5d = filters.shape
220    else:
221        _filters_shape_5d = filters_shape
222
223    Ns, Ts, C, Hs, Ws = _signals_shape_5d
224    Nf, Tf, C, Hf, Wf = _filters_shape_5d
225
226    _signals_shape_4d = (Ns * Ts, C, Hs, Ws)
227    _filters_shape_4d = (Nf * Tf, C, Hf, Wf)
228
229    if border_mode[1] != border_mode[2]:
230        raise NotImplementedError('height and width bordermodes must match')
231    conv2d_signal_shape = _signals_shape_4d
232    conv2d_filter_shape = _filters_shape_4d
233    if signals_shape is None:
234        conv2d_signal_shape = None
235    if filters_shape is None:
236        conv2d_filter_shape = None
237
238    out_4d = tensor.nnet.conv2d(
239        signals.reshape(_signals_shape_4d),
240        filters.reshape(_filters_shape_4d),
241        input_shape=conv2d_signal_shape,
242        filter_shape=conv2d_filter_shape,
243        border_mode=border_mode[1])  # ignoring border_mode[2]
244
245    # compute the intended output size
246    if border_mode[1] == 'valid':
247        Hout = Hs - Hf + 1
248        Wout = Ws - Wf + 1
249    elif border_mode[1] == 'full':
250        Hout = Hs + Hf - 1
251        Wout = Ws + Wf - 1
252    elif border_mode[1] == 'half':
253        Hout = Hs - (Hf % 2) + 1
254        Wout = Ws - (Wf % 2) + 1
255    elif border_mode[1] == 'same':
256        raise NotImplementedError()
257    else:
258        raise ValueError('invalid border mode', border_mode[1])
259
260    # reshape the temporary output to restore its original size
261    out_tmp = out_4d.reshape((Ns, Ts, Nf, Tf, Hout, Wout))
262
263    # now sum out along the Tf to get the output
264    # but we have to sum on a diagonal through the Tf and Ts submatrix.
265    if Tf == 1:
266        # for Tf==1, no sum along Tf, the Ts-axis of the output is unchanged!
267        out_5d = out_tmp.reshape((Ns, Ts, Nf, Hout, Wout))
268    else:
269        # for some types of convolution, pad out_tmp with zeros
270        if border_mode[0] == 'valid':
271            Tpad = 0
272        elif border_mode[0] == 'full':
273            Tpad = Tf - 1
274        elif border_mode[0] == 'half':
275            Tpad = Tf // 2
276        elif border_mode[0] == 'same':
277            raise NotImplementedError()
278        else:
279            raise ValueError('invalid border mode', border_mode[0])
280
281        if Tpad == 0:
282            out_5d = diagonal_subtensor(out_tmp, 1, 3).sum(axis=3)
283        else:
284            # pad out_tmp with zeros before summing over the diagonal
285            out_tmp_padded = tensor.zeros(dtype=out_tmp.dtype, shape=(
286                Ns, Ts + 2 * Tpad, Nf, Tf, Hout, Wout
287            ))
288            out_tmp_padded = tensor.set_subtensor(
289                out_tmp_padded[:, Tpad:(Ts + Tpad), :, :, :, :],
290                out_tmp)
291            out_5d = diagonal_subtensor(out_tmp_padded, 1, 3).sum(axis=3)
292
293    return out_5d
294
295
296@theano.gof.local_optimizer([DiagonalSubtensor, IncDiagonalSubtensor])
297def local_inplace_DiagonalSubtensor(node):
298    """Also work for IncDiagonalSubtensor."""
299    if (isinstance(node.op, (DiagonalSubtensor, IncDiagonalSubtensor)) and
300            not node.op.inplace):
301        new_op = node.op.__class__(inplace=True)
302        new_node = new_op(*node.inputs)
303        copy_stack_trace(node.outputs[0], new_node)
304        return [new_node]
305    return False
306theano.compile.optdb.register(
307    'local_inplace_DiagonalSubtensor',
308    TopoOptimizer(
309        local_inplace_DiagonalSubtensor,
310        failure_callback=TopoOptimizer.warn_inplace),
311    60, 'fast_run', 'inplace')
312