TensorFlow Probability中fit_surrogate_posterior参数不兼容报错求助

问题重现 #

运行TensorFlow Probability结构时间序列案例的变分损失优化代码时，触发如下类型错误：

# Build and optimize the variational loss function.
elbo_loss_curve = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn=co2_model.joint_distribution(
        observed_time_series=co2_by_month_training_data).log_prob,
    surrogate_posterior=variational_posteriors,
    optimizer=tf.optimizers.Adam(learning_rate=0.1),
    num_steps=num_variational_steps,
    jit_compile=True)

plt.plot(elbo_loss_curve)
plt.show()

报错信息：

TypeError: TF_TryEvaluateConstant_wrapper(): incompatible function arguments. The following argument types are supported:
    1. (arg0: tensorflow.python.client._pywrap_tf_session.TF_Graph, arg1: tensorflow.python.client._pywrap_tf_session.TF_Output) -> object

Invoked with: <tensorflow.python.framework.c_api_util.ScopedTFGraph object at 0x0000024EDCF90C70>, <tensorflow.python.client._pywrap_tf_session.TF_Output object at 0x0000024EDD26DBB0>

可能原因 #

1. 版本不兼容：TensorFlow（TF）与TensorFlow Probability（TFP）版本不匹配，导致底层C API调用的参数类型不兼容。
2. XLA编译冲突：jit_compile=True启用的XLA编译与TFP变分推断组件存在兼容性问题（Windows环境下更易触发）。
3. 输入数据类型错误：co2_by_month_training_data未转换为TensorFlow张量，使用numpy数组等非TF原生类型触发底层类型校验错误。

解决方案 #

方案1：禁用XLA编译 #

将jit_compile=True改为jit_compile=False，绕过XLA编译的底层冲突：

elbo_loss_curve = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn=co2_model.joint_distribution(
        observed_time_series=co2_by_month_training_data).log_prob,
    surrogate_posterior=variational_posteriors,
    optimizer=tf.optimizers.Adam(learning_rate=0.1),
    num_steps=num_variational_steps,
    jit_compile=False)  # 禁用XLA编译

方案2：确保输入为TensorFlow张量 #

将numpy数组转换为TF张量，避免类型不匹配：

# 转换输入数据为TF张量
co2_by_month_training_data = tf.convert_to_tensor(co2_by_month_training_data, dtype=tf.float32)

elbo_loss_curve = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn=co2_model.joint_distribution(
        observed_time_series=co2_by_month_training_data).log_prob,
    surrogate_posterior=variational_posteriors,
    optimizer=tf.optimizers.Adam(learning_rate=0.1),
    num_steps=num_variational_steps,
    jit_compile=True)

方案3：对齐TF与TFP版本 #

安装兼容的版本组合（例如TF 2.15.0对应TFP 0.23.0）：

pip install tensorflow==2.15.0 tensorflow-probability==0.23.0 --force-reinstall

方案4：调整优化器参数 #

适当调低学习率，避免优化过程中数值不稳定触发底层错误：

elbo_loss_curve = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn=co2_model.joint_distribution(
        observed_time_series=co2_by_month_training_data).log_prob,
    surrogate_posterior=variational_posteriors,
    optimizer=tf.optimizers.Adam(learning_rate=0.01),  # 调低学习率
    num_steps=num_variational_steps,
    jit_compile=True)

验证步骤 #

1. 优先尝试方案1和方案2，修改后运行代码观察错误是否解决。
2. 若问题仍存在，执行方案3重装兼容版本的TF和TFP。
3. 最后尝试方案4调整优化器参数，排除数值不稳定因素。

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