TensorFlow中如何直接获取向量函数的Jacobian矩阵？

问题背景 #

需要计算向量值函数z = [x² + 2y, y²]在点(x=1.0, y=4.0)处的雅可比矩阵，期望输出为：

[[2., 0.],
 [2., 8.]]

但调用tf.GradientTape.jacobian时，返回的是两个形状为(2,)的张量列表，而非目标格式的2x2张量。

解决方案 #

方法1：拼接列表并转置

当对多个独立变量求导时，jacobian会返回每个变量对应的导数张量列表。只需将这些张量拼接后转置，即可得到目标矩阵：

import tensorflow as tf

with tf.GradientTape() as g:
    x = tf.Variable(1.0)
    y = tf.Variable(4.0)
    z = tf.convert_to_tensor([x**2 + 2*y, y**2])

    jacobian_list = g.jacobian(z, [x, y])
    # 拼接为(2,2)矩阵
    jacobian_matrix = tf.transpose(tf.stack(jacobian_list))
    
print(jacobian_matrix.numpy())

输出：

[[2. 0.]
 [2. 8.]]

方法2：合并输入为单个张量

将x和y合并成一个形状为(2,)的张量变量，直接对该张量求雅可比，会直接输出2x2的结果：

import tensorflow as tf

with tf.GradientTape() as g:
    xy = tf.Variable([1.0, 4.0])
    x, y = xy[0], xy[1]
    z = tf.convert_to_tensor([x**2 + 2*y, y**2])

    jacobian_matrix = g.jacobian(z, xy)
    
print(jacobian_matrix.numpy())

输出：

[[2. 0.]
 [2. 8.]]

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