如何修改TensorFlow Probability中分布的dtype类型？

实例化时直接指定dtype #

TensorFlow Probability（TFP）的大部分分布类支持在构造时通过dtype参数显式指定数据类型，同时确保传入的参数（如loc、scale等）与目标dtype一致即可：

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

# 实例化时指定tf.float64
dist = tfd.Normal(
    loc=tf.constant(0.0, dtype=tf.float64),
    scale=tf.constant(1.0, dtype=tf.float64),
    dtype=tf.float64
)
print(dist.dtype)  # 输出: tf.float64

如果传入的参数dtype与指定的dtype匹配，甚至可以省略dtype参数，分布会自动继承参数的dtype：

# 仅通过参数dtype隐式指定
dist = tfd.Normal(
    loc=tf.constant(0.0, dtype=tf.float64),
    scale=tf.constant(1.0, dtype=tf.float64)
)
print(dist.dtype)  # 输出: tf.float64

修改已实例化分布的dtype #

TFP的分布对象是不可变的，没有专门的setter方法用于修改已创建分布的dtype。如果需要将现有分布的dtype从float32转为float64，需要重新实例化一个新的分布对象，同时将原分布的参数张量转换为目标dtype：

# 原float32分布
dist_float32 = tfd.Normal(loc=0.0, scale=1.0)
print(dist_float32.dtype)  # 输出: tf.float32

# 转换参数并重新实例化为float64
dist_float64 = tfd.Normal(
    loc=tf.cast(dist_float32.loc, tf.float64),
    scale=tf.cast(dist_float32.scale, tf.float64),
    dtype=tf.float64
)
print(dist_float64.dtype)  # 输出: tf.float64

为什么tf.cast对分布对象无效？ #

tf.cast是针对TensorFlow张量的操作，而TFP分布对象是封装了参数、采样逻辑等的复合对象，并非张量。直接对分布对象调用tf.cast不会修改其内部参数的dtype，自然无法改变分布的整体dtype。必须针对分布的参数张量单独进行类型转换，再重新构建分布。

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