Create N x M Tensor matrix with two Tensor vectors and apply a function for each (n,m) pair

user_5833067 0
Viewed 401

I have two Tensors X and Z of shape (100,) and I would like to create a Tensor X x Z which would result in a shape of (100, 100).

And for each pair element within this matrix I'd like to apply some function that I have already defined, i.e. fn(x,z) for each possible combination in the matrix.

I am new to TensorFlow and I'm used to thinking sequentially, with two for loops over each vector when dealing with numpy arrays.

How can I do this in TensorFlow? Many thanks.

tensorflow python
1 Answers

For the general question

You could create a custom layer to solve this.

import tensorflow as tf
from tensorflow.keras.layers import Layer


def fn(x, y):
    return x + y


class PairMatrix(Layer):
    def __init__(self, func=None):
        super(PairMatrix, self).__init__()
        self.func = func

    def call(self, inputs, **kwargs):
        X, Y = tf.meshgrid(inputs[0], inputs[1])
        return self.func(X, Y)


x = tf.constant([1, 2, 3])
y = tf.constant([1, 2, 3])

z = PairMatrix(func=fn).apply([x, y])

This should create a layer that takes a function as argument. When called, the function will be applied on the constructed mesh of input x and y.

In the simple test case above, the output produced is:

z = tf.Tensor([[2 3 4]
               [3 4 5]
               [4 5 6]], shape=(3, 3), dtype=int32)

Regarding your commment

If it's only the covariance that you are interested in, you could use tfp.stats.covariance

user_81504840
Related