TensorFlow量子PINN求解器二阶导数计算始终返回None问题排查

二阶导数返回None的核心原因是外层GradientTape没有追踪输入张量x的依赖关系。你的代码中，tape2仅包裹了tape1的上下文和一阶导数计算，但未显式watch x，导致TensorFlow无法建立从x到dpsi_dx的梯度追踪链路，因此无法计算二阶导数。

调整GradientTape的嵌套逻辑，确保二阶导数的依赖链能被正确追踪：

• 在外层tape2的上下文内显式调用tape2.watch(x)，让TensorFlow记录x的变化
• 保证psi和dpsi_dx的计算全程处于tape2的追踪范围内

import tensorflow as tf

class QuantumPINNSolver:
    def __init__(self, potential_func):
        self.potential_func = potential_func
        self.model = self.build_model()

    def build_model(self):
        return tf.keras.Sequential([
            tf.keras.layers.Dense(50, activation='relu', input_shape=(1,), dtype=tf.float64),
            tf.keras.layers.Dense(50, activation='relu', dtype=tf.float64),
            tf.keras.layers.Dense(50, activation='relu', dtype=tf.float64),
            tf.keras.layers.Dense(1, activation=None, dtype=tf.float64)
        ])

    @tf.function
    def schrodinger_loss(self, x, E):
        # 外层tape用于计算二阶导数，显式watch x
        with tf.GradientTape(persistent=True) as tape2:
            tape2.watch(x)
            # 内层tape用于计算一阶导数
            with tf.GradientTape(persistent=True) as tape1:
                tape1.watch(x)
                psi = self.model(x)
            dpsi_dx = tape1.gradient(psi, x)
        # 此时tape2可追踪x -> psi -> dpsi_dx的完整链路，能正确计算二阶导
        d2psi_dx2 = tape2.gradient(dpsi_dx, x)
        
        print("psi:", psi)
        print("dpsi_dx:", dpsi_dx)
        print("d2psi_dx2:", d2psi_dx2)
        
        V_x = self.potential_func(x)
        schrodinger_eq = -d2psi_dx2 + (V_x - E) * psi
        return tf.reduce_mean(tf.square(schrodinger_eq))

# 使用示例
def harmonic_oscillator_potential(x):
    return 0.5 * tf.square(x)

solver = QuantumPINNSolver(potential_func=harmonic_oscillator_potential)
x = tf.random.uniform((100, 1), minval=-5, maxval=5, dtype=tf.float64)
E = tf.constant(0.5, dtype=tf.float64)
loss = solver.schrodinger_loss(x, E)

• 由于使用了persistent=True，函数结束后TensorFlow会自动回收tape资源，简单场景下无需手动清理
• 确保所有张量操作都使用TensorFlow原生API，避免混入NumPy操作中断梯度追踪

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