Как реализовать логическое маскирование на ТПУ

Question

Мне нужно реализовать операцию логического маскирования с использованием TPUEstimator.tf.boolean_mask не реализован.Есть ли обходной путь?

Следующий код отлично подходит для моей цели на CPU и GPU:

  all_out = model.get_sequence_output()
  P = tf.boolean_mask(all_out, P_mask)

all_out - тензор формы [?, 128, 768]

P_mask имеет форму [?, 128], а второе измерение закодировано в одну горячую точку для представления желаемого тензора для извлечения.

Желаемая форма P равна [?, 768]

Когда я запускаю это на TPU с помощью TPUEstimator, я получаю следующее сообщение об ошибке:

Compilation failure: Detected unsupported operations when trying to
compile graph _functionalize_body_1[] on XLA_TPU_JIT: Where (No 
registered 'Where' OpKernel for XLA_TPU_JIT devices compatible with node
node boolean_mask/Where

Answer 1

Это связано с ограничением на tf.where (которое вызывается tf.boolean_mask на TPU, см. Также здесь .

    tf.where    Both x and y must be non-None. If both x and y are None, the operator would not have a static shape.

коренная причина в том, что вы не получите статическую форму, выполняя это, поэтому на сегодняшний день ТПУ не довольны этим.

Если ваша единственная цель - окончательно рассчитать убыток или сумму, возможно, выполнимопереписать свой код.

Rewrite this:
   reduce_sum(gather_nd(tf.where(cond),Y))
to this:
   reduce_sum(Y * tf.cast(cond))

Однако, если вам нужна динамическая форма маскированного вывода [?, 768], я понятия не имею.