I have 2 arrays (Tensors) with the same shape such as:
array1 = [[[0, 1, 0], [1, 0, 0]], [[0, 1, 0], [0, 0, 1]]]
array2 = [[[0.2, 0.7, 0.1], [0.2, 0.4, 0.4]], [[0.2, 0.8, 0.0], [0.5, 0.2, 0.3]]]
I want to create a new array (Tensor) such that values from array2 are chosen based on array1 values.
array3 = tf.magic_function(array1, array2)
# array3 = [[[0.7], [0.2]], [[0.4], [0.3]]]
I also have array1 not one-hot encoded so a solution assuming array1 = [[1, 0], [1, 2]] is also acceptable. The solution must use Tensorflow functions, NumPy is not acceptable since I can't convert my Tensors to NumPy arrays during training due to eager execution. This is for a weighted (Sparse)CategoricalCrossEntrophy loss function.
I am using Tensorflow 2.10 and Python 3.9.
Looked up tf.gather() and tf.nd_gather() but those didn't seem to do what I want.
The following code seems to do the trick, assuming array2 only has positive values:
import tensorflow as tf
array1 = [[[0, 1, 0], [1, 0, 0]], [[0, 1, 0], [0, 0, 1]]]
array2 = [[[0.2, 0.7, 0.1], [0.2, 0.4, 0.4]], [[0.2, 0.8, 0.0], [0.5, 0.2, 0.3]]]
tensor1 = tf.convert_to_tensor(array1, dtype=tf.float32)
tensor2 = tf.convert_to_tensor(array2)
array3 = tf.reduce_max(tf.math.multiply(tensor1, tensor2), axis=2, keepdims=True)
print(array3)
# tf.Tensor(
# [[[0.7]
# [0.2]]
# [[0.8]
# [0.3]]], shape=(2, 2, 1), dtype=float32)
If we can't make that assumption, we can also define array3 like so:
array3 = tf.reshape(
tf.boolean_mask(tensor2, tf.cast(tensor1, tf.bool)), (len(tensor1), -1, 1)
)