Блог

Я пытаюсь создать слой в keras, который принимает плоский тензор x (не имеет нулевого значения в нем и shape = (batch_size, единицы измерения)), умноженный на a mask (той же формы), и он будет сортировать его таким образом, чтобы замаскированные значения были помещены первыми в вывод (порядок значений элементов не имеет значения). Для наглядности вот пример (batch_size = 1, units = 8):

Это кажется простым, но проблема в том, что я не могу найти хорошее решение. Приветствуется любой код или идея.

Мой текущий код выглядит следующим образом, если вы знаете более эффективный способ, пожалуйста, дайте мне знать.

 class Sort(keras.layers.Layer):
  
  def call(self, inputs):
    x = inputs.numpy()
    nonx, nony = x.nonzero()    # idxs of nonzero elements
    zero = [np.where(x == 0)[0][0], np.where(x == 0)[1][0]]   # idx of first zero
    x_shape = tf.shape(inputs)
    result = np.zeros((x_shape[0], x_shape[1], 2), dtype = 'int')   # mapping matrix

    result[:, :, 0]  = zero[0]
    result[:, :, 1]  = zero[1]
    p = np.zeros((x_shape[0]), dtype = 'int')
    for i, j in zip(nonx, nony):
      result[i, p[i]] = [i, j]
      p[i]  = 1

    y = tf.gather_nd(inputs, result)
    return y