Масштабирование множества строк в тензорном множителе

Question

TL; DR Как масштабировать часть тензора на 2 (индексы строк присутствуют в списке tf)

---

подробности:

indices_of_scaling_ids: список магазинов row_ids Tensor("Squeeze:0", dtype=int64, device=/device:GPU:0) [1, 4, 5, 6, 12]

emb_inputs = tf.nn.embedding_lookup(embedding, self.all_rows) #tensor with shape (batch_size=4, all_row_len, emb_size=128)

Итак, для каждого self.all_rows оценивается emb_inputs.

Вопрос / Задача, с которой столкнулся: мне нужно масштабировать emb_inputs на 2.0 для каждого row_ids, упомянутого в indices_of_scaling_ids. Я пробовал разные способы сращивания, но не могу найти хорошее решение. Может кто-то предложить? Спасибо

N.B. Новичок в Tensorflow

Answer 1

Попробуйте что-то вроде этого:

SCALE = 2
emb_inputs = ...
indices_of_scaling_ids = ...
emb_shape = tf.shape(emb_inputs)
# Select indices in boolean array
r = tf.range(emb_shape[1])
mask = tf.reduce_any(tf.equal(r[:, tf.newaxis], indices_of_scaling_ids), axis=1)
# Tile the mask
mask = tf.tile(mask[tf.newaxis, :, tf.newaxis], (emb_shape[0], 1, emb_shape[2]))
# Choose scaled or not depending on indices
result = tf.where(mask, SCALE * emb_inputs, emb_inputs)