What is the theano counterpart of numpy apply_alon

问题:

I want to apply some function to a tensor with Theano. Here's my numpy version script. But, I'm lost when converting this to a Theano version.

For example, I want to convert

array([[[0, 0, 1, 2, 3],
        [0, 1, 2, 3, 4]]])

to

array([[[0, 0, 3, 2, 1],
        [0, 4, 3, 2, 1]]])

My numpy script is as follows:

def reverse_and_shift(vec):
    '''
    Returns reversed array of the given vector except the leading 0's.
    vec: 1-d array.
    >>> reverse_and_shift(np.array([0,0,1,2,3]))
    array([0, 0, 3, 2, 1])
    '''
    reversed = vec[::-1]
    num_zero = len(vec) - np.count_nonzero(vec)
    shifted = np.roll(reversed, num_zero)
    return shifted

np.apply_along_axis(reverse_and_shift, -1, A)

How can I execute this in a tensor?

回答1:

Try it on arrays of different dimensions:

For a 1d array, it doesn't add any functionality:

In [36]: A=np.arange(10)
In [38]: np.apply_along_axis(reverse_and_shift,-1,A)
Out[38]: array([0, 9, 8, 7, 6, 5, 4, 3, 2, 1])
In [39]: reverse_and_shift(A)
Out[39]: array([0, 9, 8, 7, 6, 5, 4, 3, 2, 1])

For a 2d array, it's equivalent to iterating on the 1st dimension, applying the function to the last (-1):

In [41]: A=np.arange(9).reshape(3,3)
In [42]: np.apply_along_axis(reverse_and_shift,-1,A)
Out[42]: 
array([[0, 2, 1],
       [5, 4, 3],
       [8, 7, 6]])

In [43]: reverse_and_shift(A[0,:])
Out[43]: array([0, 2, 1])
In [44]: reverse_and_shift(A[1,:])
Out[44]: array([5, 4, 3])
In [45]: reverse_and_shift(A[2,:])
Out[45]: array([8, 7, 6])

In [46]: np.array([reverse_and_shift(x) for x in A])
Out[46]: 
array([[0, 2, 1],
       [5, 4, 3],
       [8, 7, 6]])

For 3d it iterates on all values of the 1st 2 dimensions, applying your function to the last.

In [48]: A=np.arange(8).reshape(2,2,2)
In [49]: np.apply_along_axis(reverse_and_shift,-1,A)
Out[49]: 
array([[[0, 1],
        [3, 2]],

       [[5, 4],
        [7, 6]]])

There are various ways of doing your own iteration on these values. One of the easiest is to reshape the array to 2d (and then back)

In [50]: np.array([reverse_and_shift(x) for x in A.reshape(-1,2)])
Out[50]: 
array([[0, 1],
       [3, 2],
       [5, 4],
       [7, 6]])
In [51]: _.reshape(2,2,2)
Out[51]: 
array([[[0, 1],
        [3, 2]],

       [[5, 4],
        [7, 6]]])

This reshape approach is easiest if the application axis is the last or first. If you want to apply along a middle axis, it maybe easier to first swap some axis, moving that axis to the end.

apply_along_axis is written in Python, so you can read it, and extract the relevant parts as needed. All the iteration is done in Python, so there's no performance loss if you do your own iteration.

=======================

In your example case, A is 3d; apply preserves that shape

In [70]: A.shape
Out[70]: (1, 2, 5)
In [71]: np.apply_along_axis(reverse_and_shift,-1,A)
Out[71]: 
array([[[0, 0, 3, 2, 1],
        [0, 4, 3, 2, 1]]])
In [72]: _.shape
Out[72]: (1, 2, 5)

the reshape iterative case removes a dimension, which would have to be added back.

In [73]: np.array([reverse_and_shift(x) for x in A.reshape(-1,A.shape[-1])])
Out[73]: 
array([[0, 0, 3, 2, 1],
       [0, 4, 3, 2, 1]])
In [74]: _.shape
Out[74]: (2, 5)

apply_along_axis does (roughly)

outarr=np.zeros_like(A)
for <all applicable i,j>:
    ind = (i,j,slice(None))
    outarr[ind] = reverse_and_shift(A[ind])