TensorFlow中tf.svd在GradientTape梯度计算时出现未实现错误求助

正在实现多张量网络的缩并操作，缩并过程中使用奇异值分解（SVD）简化计算流程。不计算梯度时代码运行完全正常，但一旦通过GradientTape监听张量启动梯度计算就会报错，暂未定位到关联原因。

相关代码 #

缩并函数与自定义SVD实现 #

import tensorflow as tf

@tf.function
def contraction_step(network, max_singular_values: int):

    bottom = network[-1]
    uppper = network[-2]

    def contract_up_down(up,dn):
        shu = tf.shape(up)
        shd = tf.shape(dn)
        c = tf.einsum("ijkxlm,nkpyqr->injpxylqmr", up, dn)
        return tf.reshape(c, (
            shu[0]*shd[0], shu[1], shd[2], shu[-3], shd[-3], shu[-2]*shd[-2], shu[-1]*shd[-1]
        ))

    new = []
    multiplier = tf.eye(tf.shape(bottom[-1])[-1]*tf.shape(uppper[-1])[-1], dtype=uppper[-1].dtype)
    for ix in reversed(range(len(bottom))):
        tensor = contract_up_down(uppper[ix], bottom[ix])
        t = tf.einsum("ludpxor,ij->ludpxoj",tensor, multiplier)
        u, s, vh = svd(t, 1, max_singular_values = max_singular_values)
        multiplier = tf.tensordot(u, s/tf.norm(s), axes=(-1,0))
        new.insert(0, vh)
    new[-1] = tf.tensordot(new[-1], multiplier, axes=(-1,0))

    return network[:-2] + [new]

def svd(tensor,pivot,max_singular_values = None,cutoff = 0.0):
    left_dims = tf.shape(tensor)[:pivot]
    right_dims = tf.shape(tensor)[pivot:]
    tensor = tf.reshape(tensor, (tf.reduce_prod(left_dims), tf.reduce_prod(right_dims)))

    s, u, v = tf.linalg.svd(tensor)

    s_shape = tf.math.count_nonzero(
        tf.cast(s >= cutoff, dtype = tf.int32), dtype = tf.int32
    )
    if max_singular_values is None:
        max_singular_values = s_shape
    else:
        max_singular_values = tf.cast(tf.constant(max_singular_values), dtype = tf.int32)
    num_sing_vals_keep = tf.maximum(
        tf.minimum(max_singular_values, s_shape), tf.constant(1, dtype = tf.int32)
    )

    s = tf.slice(s, [0], [num_sing_vals_keep])
    u = tf.slice(u, [0, 0], [tf.shape(u)[0], num_sing_vals_keep])
    v = tf.slice(v, [0, 0], [tf.shape(v)[0], num_sing_vals_keep])

    vh = tf.linalg.adjoint(v)

    dim_s = tf.shape(s)[0]  # must use tf.shape (not s.shape) to compile
    u = tf.reshape(u, tf.concat([left_dims, [dim_s]], axis = -1))
    vh = tf.reshape(vh, tf.concat([[dim_s], right_dims], axis = -1))

    return u, tf.linalg.diag(s), vh

正常运行测试用例 #

upper = [tf.random.uniform((5,3,3,2,1,5), dtype=tf.float64) for _ in range(5)]
lower = [tf.random.uniform((5,3,3,2,1,5), dtype=tf.float64) for _ in range(5)]

contracted = contraction_step([upper, lower], 2)[0]
print(f"shapes: {', '.join([str(x.shape) for x in contracted])}")

# 输出：
# shapes: (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2)

梯度测试代码与报错 #

with tf.GradientTape() as tape:
    tape.watch(upper + lower)
    contracted = contraction_step([upper, lower], 2)[0]

报错信息：

NotImplementedError: SVD gradient has not been implemented for input with unknown inner matrix shape.

已尝试的排查方案 #

• 将张量设置为tf.Variable替代手动watch，报错完全一致

系统信息 #

• 操作系统及版本：macOS v12.0.1
• TensorFlow版本：v2.6.0-rc2-32-g919f693420e 2.6.0
• Python版本：3.8.9

问题原因 #

TF 2.6版本的SVD梯度实现要求输入矩阵的两个维度在计算图编译阶段必须是静态确定的，你代码中所有形状计算用的都是tf.shape返回的动态张量，编译阶段无法确定reshape后的矩阵尺寸，因此触发未实现错误。

解决方案 #

方案一：移除@tf.function装饰器 #

以Eager模式执行运算，该模式下不需要提前编译计算图，SVD梯度可以正常处理动态形状，代码改动最小，适合网络规模不大、对性能要求不高的场景。

方案二：升级TensorFlow版本 #

升级TensorFlow到2.8及以上版本，更高版本的TensorFlow优化了SVD梯度的实现，已经支持动态形状输入的梯度计算，不需要修改业务代码即可解决问题。

方案三：静态指定SVD输入维度 #

如果必须保留TF2.6版本和@tf.function装饰器，需要修改svd函数中的形状计算逻辑，优先用张量的静态形状属性.shape计算维度乘积，确保编译阶段可以拿到确定的矩阵尺寸，修改示例：

def svd(tensor,pivot,max_singular_values = None,cutoff = 0.0):
    # 优先取静态形状
    left_dims_static = tensor.shape[:pivot]
    right_dims_static = tensor.shape[pivot:]
    if left_dims_static.is_fully_defined() and right_dims_static.is_fully_defined():
        left_prod = tf.reduce_prod(left_dims_static).numpy()
        right_prod = tf.reduce_prod(right_dims_static).numpy()
    else:
        left_prod = tf.reduce_prod(tf.shape(tensor)[:pivot])
        right_prod = tf.reduce_prod(tf.shape(tensor)[pivot:])
    tensor = tf.reshape(tensor, (left_prod, right_prod))
    # 剩余逻辑和原代码一致
    ...

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