如何用tfp.stats.quantiles()实现np.quantile()的等价功能？

np.quantile和tfp.stats.quantiles的核心差异在于分位数参数的传递方式，只需对应好参数即可实现完全一致的功能，以下是匹配你示例代码的实现方案：

完整对比代码 #

import tensorflow_probability as tfp
import tensorflow as tf 
import numpy as np 

# 生成测试数据，固定随机种子保证结果可复现
tf.random.set_seed(42)
np.random.seed(42)
inputs = tf.random.normal((1, 4096, 4))

# NumPy 实现
print("NumPy 结果：")
np_result = np.quantile(inputs.numpy(), q=0.9, axis=1, keepdims=False)
print(np_result)

# TensorFlow Probability 等价实现
print("\nTensorFlow Probability 结果：")
tfp_result = tfp.stats.quantiles(
    inputs,
    quantile_points=[0.9],  # 对应np.quantile的q参数
    axis=1,
    keepdims=False
)
# 移除因单个分位数产生的额外维度，对齐np结果的维度
tfp_result = tf.squeeze(tfp_result, axis=-1)
print(tfp_result.numpy())

# 验证结果一致性（浮点误差范围内）
print("\n结果一致性验证：")
print(np.allclose(np_result, tfp_result.numpy()))

关键参数说明 #

• quantile_points：传入与np.quantile中q相同的0-1数值列表，若需多个分位数可直接扩展列表（如[0.25, 0.5, 0.75]）
• axis：与np.quantile的axis参数完全一致，指定计算分位数的维度
• keepdims：同样和np.quantile的keepdims参数功能一致，控制是否保留原维度
• tf.squeeze：由于tfp.stats.quantiles会为每个分位点数新增一个维度，当仅计算单个分位数时，需要用此操作移除多余维度，保证结果维度与np.quantile一致

内容的提问来源于stack exchange，提问作者S. P