如何在TPU上实现布尔掩码

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]

当我使用TPUEstimator在TPU上运行它时，我收到以下错误消息：

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

这是由于limitation上的tf.where（在TPU上由tf.boolean_mask调用，也见here）。

    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.

根本原因在于你不会通过这样做获得静态形状，因此截至今天，tpu对此并不满意。

如果您的唯一目的是最终计算损失或总和，则可能需要重写您的代码。

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

但是，如果你确实需要蒙版输出[?, 768]的动态形状，我不知道。