How to to 0/1 subsetting in Theano?

Question

The objective is to obtain the subset of an array of elements through values provided in another array.

import theano
import theano.tensor as T

a = T.vector('X', dtype='int64')
b = T.vector('Y', dtype='int64')
c = a[b]
g = function([a,b],c)

x = np.array([5,3,2,3,4,6], dtype=int)
y = np.array([0,0,1,0,0,1], dtype=int)
print g(x,y)

This prints

[5 5 3 5 5 3]

instead of

[2 6]

How do I get the expected result?

Answer 1

try to use nonzero() function.

Example in your case:

import theano
import theano.tensor as T

a = T.vector('X', dtype='int64')
b = T.vector('Y', dtype='int64')
c = a[b.nonzero()]
g = function([a,b],c)

x = np.array([5,3,2,3,4,6], dtype=int)
y = np.array([0,0,1,0,0,1], dtype=int)
print g(x,y)

hope it helps