TensorFlow创建含变量元素的对角矩阵遇形状问题求解决方案

我来帮你搞定这个TensorFlow矩阵创建的问题～

先说说你的代码为什么会出形状问题 #

你现在的A0和A1都是形状为(1,2)的行向量，当用tf.stack([A0,A1])时，默认是在**最外层的维度（axis=0）**堆叠，相当于把两个(1,2)的张量摞在一起，就变成了(2,1,2)的3维张量，这自然不符合你想要的2x2矩阵。

要修复这个问题，其实把最后一行的tf.stack换成tf.concat就行，指定在axis=0维度合并：

A = tf.concat([A0, A1], axis=0)

这样就能得到形状为(2,2)的矩阵了。

更简便的创建对角矩阵的方法 #

其实TensorFlow有专门的工具，比手动拼接更省事，给你推荐几种：

1. 用tf.linalg.diag（最推荐）

这是TensorFlow专门生成对角矩阵的函数，直接把你的参数a和b拼成一个向量传进去就行：

a = tf.Variable(1.0, name="a", dtype=tf.float64)
b = tf.Variable(1.0, name="b", dtype=tf.float64)
# 先把a和b拼成一维向量，再生成对角矩阵
A = tf.linalg.diag(tf.stack([a, b]))

最终得到的A就是标准的(2,2)对角矩阵，完全符合你的需求。

2. 直接构造小矩阵（直观易懂）

因为你的矩阵是固定2x2的，也可以直接像写普通矩阵一样构造，可读性拉满：

a = tf.Variable(1.0, name="a", dtype=tf.float64)
b = tf.Variable(1.0, name="b", dtype=tf.float64)
A = tf.convert_to_tensor([[a, 0.0], [0.0, b]], dtype=tf.float64)

这种方式适合维度很小的矩阵，一眼就能看出来矩阵的结构。

3. 调整原有代码的堆叠方式（不推荐，但可以试试）

如果不想大改原有代码，也可以在tf.stack时指定axis=1，再转置一下：

A = tf.transpose(tf.stack([A0, A1], axis=1))

不过这种写法绕了个弯，不如前面两种直接。

内容的提问来源于stack exchange，提问作者measure_theory