tf.boolean_mask

tf.boolean_mask就是true 对应位置的张量

import tensorflow as tf
import numpy as np
# 1-D example
tensor = [0, 1, 2, 3]
mask = np.array([True, False, True, False])
c=tf.boolean_mask(tensor, mask)  # [0, 2]

print(tf.keras.backend.eval(c))

[0 2]

# 2-D example
tensor = [[1, 2], [3, 4], [5, 6]]
mask = np.array([True, False, True])
c=tf.boolean_mask(tensor, mask)  # [[1, 2], [5, 6]]

print(tf.keras.backend.eval(c))

[[1 2]
 [5 6]]