python - 使用两个张量向量创建 N x M 张量矩阵并为每个 (n,m) 对应用一个函数

Question

我有两个形状为 (100,) 的张量 X 和 Z，我想创建一个 Tensor X x Z，它的形状为 (100, 100).

对于这个矩阵中的每个对元素，我想应用一些我已经定义的函数，即 fn(x,z) 用于矩阵中的每个可能组合。

我是 TensorFlow 的新手，我习惯于按顺序思考，在处理 numpy 数组时对每个向量使用两个 for 循环。

我如何在 TensorFlow 中执行此操作？非常感谢。

Answer 1

一般问题

您可以创建一个自定义层来解决这个问题。

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

这应该创建一个将函数作为参数的层。调用时，该函数将应用于输入 x 和 y 的构造网格。

在上面的简单测试用例中，产生的输出是:

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

关于您的评论

如果您只对协方差感兴趣，可以使用 tfp.stats.covariance