pythonpython-3.xtensorflowtensorflow2.0

How To Choose Values From Tensor Using Another Tensor In Tensorflow


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.


Solution

  • 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)
    )