Equivalent tensorflow expression to numpy mask-CodePudding

I have a numpy array named PixelData of unknown shape, and I am using the following condition to filter values in the array greater than some value x using a mask:

PixelData[PixelData>=x] = PixelData[PixelData>=x] - x

When I convert this numpy array to a tensor, I cannot perform the same masking operation. I have tried using tf.where as follows:

PixelData = tf.where(PixelData>=x, PixelData - x, PixelData)

In the official documentation, they always seem to define the mask dimensions in advance to equal the dimensions of the tensor being masked, but then they talk about the dimensions being broadcasted automatically, so I am a bit confused. Are these two functions equivalent? Are there any situations where they may produce different outputs?

CodePudding user response：

Not sure what PixelData looks like, but here is working example with both methods:

import numpy as np
import tensorflow as tf

x = 2
np_pixel_data = np.array([[3, 4, 5, 1],
                       [6, 4, 2, 5]], dtype=np.float32)

np_pixel_data[np_pixel_data>=x] = np_pixel_data[np_pixel_data>=x] - x

tf_pixel_data = tf.constant([[3, 4, 5, 1],
                       [6, 4, 2, 5]], dtype=tf.float32)

tf_pixel_data = tf.where(tf.greater_equal(tf_pixel_data, x), tf_pixel_data - x, tf_pixel_data)

print(np_pixel_data)
print(tf_pixel_data)

[[1. 2. 3. 1.]
 [4. 2. 0. 3.]]
tf.Tensor(
[[1. 2. 3. 1.]
 [4. 2. 0. 3.]], shape=(2, 4), dtype=float32)

You might have some minor rounding differences, but nothing significant.