1

In my problem I have a vector containing n elements. Given a window size k I want to efficiently create a matrix size n x 2k+1 which contains the banded diagonal. For example:

a = [a_1, a_2, a_3, a_4]
k = 1
b = [[0, a_1, a_2],
    [a_1, a_2, a_3],
    [a_2, a_3, a_4],
    [a_3, a_4, a_5],
    [a_4, a_5, 0]]

The naive way to implement this would be using for loops

out_data = mx.ndarray.zeros((n, 2k+1))
for i in range(0, n):
    for j in range(0, 2k+1):
        index = i - k + j
        if not (index < 0 or index >= seq_len):
            out_data[i][j] = in_data[index]

This is very slow.

Creating the full matrix would be easy by just using tile and reshape, however the masking part is not clear.

Update I found a faster, yet still very slow, implementation:

window = 2*self.windowSize + 1
in_data_reshaped = in_data.reshape((batch_size, seq_len))
out_data = mx.ndarray.zeros((seq_len * window))
for i in range(0, seq_len):
    copy_from_start = max(i - self.windowSize, 0)
    copy_from_end = min(seq_len -1, i+1+self.windowSize)
    copy_length = copy_from_end - copy_from_start
    copy_to_start = i*window + (2*self.windowSize + 1 - copy_length)
    copy_to_end = copy_to_start + copy_length
    out_data[copy_to_start:copy_to_end] = in_data_reshaped[copy_from_start:copy_from_end]
out_data = out_data.reshape((seq_len, window))
kanedo
  • 75
  • 6

1 Answers1

0

If in your operation, k and n are constant and you can do what you want using a combination of mxnet.nd.gather_nd() and mx.nd.scatter_nd. Even though generating the indices tensor is just as inefficient, because you need to do it only once, that wouldn't be a problem. You would want to use gather_nd to effectively "duplicate" your data from original array and then use scatter_nd to scatter them to the final matrix shape. Alternatively, you can simply concatenate a 0 element to your input array (for example [a_1, a_2, a_3] would turn into [0, a_1, a_2, a_3]) and then use only mxnet.nd.gather_nd() to duplicate elements into your final matrix.

Sina Afrooze
  • 960
  • 6
  • 11
  • Thanks this tip was helpful. I did an example implementation see this [gist](https://gist.github.com/kanedo/00f96bd3f9ac52746d7d7014f561e8ce). There are two caveats. (1) padding in mxnet is only working for 4- or 5-dimensional arrays [see this feature request](https://github.com/apache/incubator-mxnet/issues/9966) (2) gather_nd is not yet "compatible" with batches, meaning the indices have to be calculated explicitly [see this feature request](https://github.com/apache/incubator-mxnet/issues/9998) – kanedo Mar 07 '18 at 15:09
  • For 1st caveat, you can call `NDArray.reshape()`, to add extra dimensions to your array and reshape it back after you've done the padding: attention_padded = mx.nd.pad(attention.reshape((1, 2, 4, 1)), mode='constant', constant_value=0, pad_width=(0, 0, 0, 0, k, k, 0, 0)).reshape((2, 4, 1)) – Sina Afrooze Mar 08 '18 at 03:02