TensorFlow中GradientTape.gradient的output_gradients参数作用解析

我想了解TensorFlow中GradientTape对象的gradient函数里的output_gradients参数的作用。官方文档指出该参数应包含“target每个可微元素对应的梯度列表”，默认值为None，但文档的解释不够清晰。

当省略该参数时，函数会计算某函数的雅可比矩阵，例如对于z <- f(x)，得到dz/dx。我原本以为传入output_gradients为dL/dz（损失L对z的梯度）时，会根据链式法则计算乘积，即dL/dx = dz/dx * dL/dz，但测试后得到的结果却不符合预期。output_gradients到底是如何起作用的？文档中没有相关的详细说明...

以下是一段示例代码（来自gpflow仓库）：

from typing import Optional
import tensorflow as tf
import numpy as np
import gpflow

def natgrad_apply_gradients(
        q_mu_grad: tf.Tensor,
        q_sqrt_grad: tf.Tensor,
        q_mu: gpflow.Parameter,
        q_sqrt: gpflow.Parameter,
        xi_transform: Optional[gpflow.optimizers.natgrad.XiTransform] = None,
    ) -> None:
        gamma = 1
        xi_transform = gpflow.optimizers.natgrad.XiNat()

        dL_dmean = gpflow.base._to_constrained(q_mu_grad, q_mu.transform)
        dL_dvarsqrt = gpflow.base._to_constrained(q_sqrt_grad, q_sqrt.transform)

        with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
            tape.watch([q_mu.unconstrained_variable, q_sqrt.unconstrained_variable])
            eta1, eta2 = gpflow.optimizers.natgrad.meanvarsqrt_to_expectation(q_mu, q_sqrt)
            meanvarsqrt = gpflow.optimizers.natgrad.expectation_to_meanvarsqrt(eta1, eta2)

        dL_deta1, dL_deta2 = tape.gradient(
            meanvarsqrt, [eta1, eta2], output_gradients=[dL_dmean, dL_dvarsqrt]
        )

        dtheta_deta1, dtheta_deta2 = tape.gradient(
            meanvarsqrt, [eta1, eta2], output_gradients=None
        )

        return dL_deta1, dL_deta2, dtheta_deta1, dtheta_deta2

X_data = tf.ones(5)
num_latent_gps = 1
static_num_data = X_data.shape[0]
q_sqrt_unconstrained_shape = (num_latent_gps, gpflow.utilities.triangular_size(static_num_data))
num_data = gpflow.Parameter(tf.shape(X_data)[0], shape=[], dtype=tf.int32, trainable=False)
dynamic_num_data = tf.convert_to_tensor(num_data)

mu = np.array([[0.93350756], [0.15833747], [0.23830378], [0.28742445], [0.14999759]])
q_mu = gpflow.Parameter(mu, shape=(static_num_data, num_latent_gps))
q_sqrt = tf.eye(dynamic_num_data, batch_shape=[num_latent_gps])
q_sqrt = gpflow.Parameter(
    q_sqrt,
    transform=gpflow.utilities.triangular(),
    unconstrained_shape=q_sqrt_unconstrained_shape,
    constrained_shape=(num_latent_gps, static_num_data, static_num_data),
)

q_mu_grad = q_mu.unconstrained_variable * 0.33
q_sqrt_grad = q_sqrt.unconstrained_variable

dL_deta1, dL_deta2, dtheta_deta1, dtheta_deta2 = natgrad_apply_gradients(q_mu_grad, q_sqrt_grad, q_mu, q_sqrt)

dL_deta1 !== dtheta_deta1 * q_mu_grad

希望有人能解答我的疑问，提前感谢！

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