Theano / Keras: Set the K-first values of Tensor to a value

Question

I have a custom Keras layer and i use Theano as the backend and i want to do the following operation:

Suppose we have a tensor with shape (N,). I want to set the K first values to a fixed value x (3 or whatever...). How do i do that? I assume i have to use argsort but i don't know how to implement it using Theano.

For example in a simple FF layer, how can i set the first N values of the tensor a to a value x?

def call(self, x, mask=None):
    a = K.dot(x, self.W)

    if self.bias:
        a += self.b

    return a

For example using numpy i can set the 10 first values of a to 1 like this:

a[a.argsort() <= 10] = 1

I want to do the same using Keras backend functions or Theano functions.

Answer 1

For a[a.argsort() <= 10] = 1 the equivalent code will be:

import numpy as np
from keras import backend as K

a = np.asarray([[3,4,2,44,22,4,5,6,77,86,3,2,3,23,44,21],
                [3,4,22,44,2,4,54,6,77,8,3,2,36,23,4,2]], dtype=np.float)
a_t = K.variable(a)

a[a.argsort() <= 10] = 1
print a

arg_sort = K.T.argsort(a_t)
_cond = K.lesser_equal(arg_sort,10)
a_new = K.switch(_cond, 1, a_t)
print K.eval(a_new)

In a single statement it will be:

a_new = K.switch(K.lesser_equal(K.T.argsort(a_t),10), 1, a_t)
print K.eval(a_new)

I am not sure if this is exactly what you need.