def masked_apply(tensor, op, mask):
"""Applies `op` to tensor only at locations indicated by `mask` and sets the rest to zero.
Similar to doing `tensor = tf.where(mask, op(tensor), tf.zeros_like(tensor))` but it behaves correctly
when `op(tensor)` is NaN or inf while tf.where does not.
:param tensor: tf.Tensor
:param op: tf.Op
:param mask: tf.Tensor with dtype == bool
:return: tf.Tensor
"""
chosen = tf.boolean_mask(tensor, mask)
applied = op(chosen)
idx = tf.to_int32(tf.where(mask))
result = tf.scatter_nd(idx, applied, tf.shape(tensor))
return result
评论列表
文章目录