如何创建在 tensorflow2 中训练时仅更新特定部分的 tensorflow 变量

Question

例如，y=Ax 其中 A 是对角矩阵，其可训练权重块 (w1, w2, w3) 在对角线上。

A = [w1 ... ...
    ...  w2 ...
    ... ... w3]

我首先创建了 tf.Variables w1、w2 和 w3。然后我通过将它们连接在一起制作了 A(9x9)。我将大小为 (3x3) 的 tf.Variable 的（w1、w2 和 w3）与适当大小的常量零数组连接起来以生成 A.

当我在旧版本的 TensorFlow 中优化我的网络时，它运行良好。它能够找到梯度并仅在指定区域最小化权重。我不需要明确提及可训练变量。 但是，当我在 tf2 中编写代码并使用 tf.gradientTape 为优化器查找梯度时，发现返回的梯度是 NoneType。 这可能是因为对角矩阵 A 的可训练部分没有指向与连接创建它的原始 tf.Variable 相同的内存地址吗？