Python – how to apply a mask to a tensor and maintain its original shape

how to apply a mask to a tensor and maintain its original shape… here is a solution to the problem.

how to apply a mask to a tensor and maintain its original shape

I have two tensors: one containing the data and the other containing the bool value mask. If the bool value is False, I want to set all values in the data tensor to zero while keeping the original shape of the data tensor.
So far, I’ve only been able to implement mask when it’s a numpy array.

https://www.tensorflow.org/api_docs/python/tf/boolean_mask Affects the shape of the tensor, I can’t use it.

How?

import numpy as np
import tensorflow as tf
tf.enable_eager_execution()

# create dummy data
data_np = np.ones((4,2,3))
mask_np = np.array([[True, True],[False, True],[True, True],[False, False]])

# prepare tensors
data = tf.convert_to_tensor(data_np)
mask = tf.convert_to_tensor(mask_np)

# how to perform the same while avoiding numpy?
mask = np.expand_dims(mask, -1)
data *= mask

Solution

Use tf.cast() and tf.expand_dims():

import tensorflow as tf
import numpy as np

mask_np = np.array([[True, True],[False, True],[True, True],[False, False]])
data_np = np.ones((4,2,3))

mask = tf.convert_to_tensor(mask_np, dtype=tf.bool)
mask = tf.expand_dims(tf.cast(mask, dtype=tf.float32), axis=len(mask.shape))
data = tf.convert_to_tensor(data_np, dtype=tf.float32)

result = mask * data

print(result.numpy())
# [[[1. 1. 1.]
#   [1. 1. 1.]]
# 
#  [[0. 0. 0.]
#   [1. 1. 1.]]
# 
#  [[1. 1. 1.]
#   [1. 1. 1.]]
# 
#  [[0. 0. 0.]
#   [0. 0. 0.]]]

Related Problems and Solutions