Apple M1上TensorFlow Probability二项分布类型错误求助

环境：Apple M1设备，tensorflow-macos 2.10.0、tensorflow-probability 0.18.0、numpy 1.23.3，其他TensorFlow模型可正常运行。

模型代码中，berndfs是字典，每个s对应包含样本和调查数据的两个pandas DataFrame，核心代码如下：

def tfdmrp_run():
    n_chains = 4
    dtype = tf.float32
    berndfs = load_tfdmrp_datacache()
    shortnames = list(berndfs.keys())
    
    sample = berndfs[s]['svy']
    age_shape = len(berndfs[s]['census'].agecat.unique())
    gender_shape = len(berndfs[s]['census'].sex.unique())
    edu_shape = len(berndfs[s]['census'].educat.unique())
    counts = sample['count'].values.tolist()
    agecatlist = sample.agecat.values.tolist()
    genderlist = sample.gender.values.tolist()
    edulist = sample.educat.values.tolist()

    modlist = [
        tfd.HalfNormal(1),
        lambda sigma_age: tfd.Sample(tfd.Normal(0,sigma_age),sample_shape=age_shape),
        tfd.HalfNormal(1),
        lambda sigma_gender: tfd.Sample(tfd.Normal(0,sigma_gender),sample_shape=gender_shape),
        tfd.HalfNormal(1),
        lambda sigma_edu: tfd.Sample(tfd.Normal(0,sigma_edu),sample_shape=edu_shape),
        tfd.Normal(0,1), #intercept
        lambda intercept,coef_edu,a,coef_gender,b,coef_age: tfd.Independent(
            tfd.Binomial(
                total_count=tf.cast(counts,tf.int32),
                logits=intercept
                + tf.squeeze(tf.gather(coef_age, tf.cast(agecatlist,tf.int32),axis=-1))
                + tf.squeeze(tf.gather(coef_gender,tf.cast(genderlist,tf.int32),axis=-1))
                + tf.squeeze(tf.gather(coef_edu,tf.cast(edulist,tf.int32),axis=-1))
            ),
            reinterpreted_batch_ndims=1
        )
    ]

    model = tfd.JointDistributionSequential(modlist)
    model.resolve_graph()

调用model.resolve_graph()或sample()方法时，抛出错误：

*** TypeError: Found incompatible dtypes, <class 'numpy.int32'> and <class 'numpy.float32'>. Seen so far: [<class 'numpy.int32'>, <class 'numpy.float32'>, ...]

完整报错栈：

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/jacobtucker/Documents/Projects/repos/rutracker/dmrp/dmrptfp.py", line 124, in tfdmrp_run
    model.resolve_graph()
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution_sequential.py", line 460, in resolve_graph
    distribution_names = self._flat_resolve_names(
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution_sequential.py", line 473, in _flat_resolve_names
    for d in self._get_single_sample_distributions()]
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 353, in _get_single_sample_distributions
    ds = self._execute_model(
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 1045, in _execute_model
    d = gen.send(next_value)
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution_sequential.py", line 399, in _model_coroutine
    dist = dist_fn(*xs)
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution_sequential.py", line 610, in dist_fn_wrapped
    return dist_fn(*reversed(xs[-len(args):]))
  File "/Users/jacobtucker/Documents/Projects/repos/rutracker/dmrp/dmrptfp.py", line 112, in <lambda>
    tfd.Binomial(
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/decorator.py", line 232, in fun
    return caller(func, *(extras + args), **kw)
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/distribution.py", line 342, in wrapped_init
    default_init(self_, *args, **kwargs)
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/distributions/binomial.py", line 371, in __init__
    dtype = dtype_util.common_dtype([total_count, logits, probs], tf.float32)
  File "/Users/jacobtucker/miniconda3/envs/tfdmrp/lib/python3.10/site-packages/tensorflow_probability/python/internal/dtype_util.py", line 104, in common_dtype
    raise TypeError(
TypeError: Found incompatible dtypes, <class 'numpy.int32'> and <class 'numpy.float32'>. Seen so far: [<class 'numpy.int32'>, <class 'numpy.float32'>, ...]

用户的困惑：模型输入看似都是Python原生类型，却出现dtype不兼容问题，依赖版本无兼容性问题，求原因及解决办法。

错误根源在tfd.Binomial的初始化步骤：

• counts = sample['count'].values.tolist()得到的列表中，元素是numpy.int32类型（而非Python原生int）。
• 尽管用了tf.cast(counts, tf.int32)，但在JointDistributionSequential解析模型图时，TFP会先尝试用原始numpy类型做类型推断，而logits部分是TensorFlow的float32张量，numpy类型和TensorFlow类型体系不兼容，导致无法找到共同dtype，抛出错误。

有两种可行方案：

方案一：直接将counts转为TensorFlow张量 #

把counts = sample['count'].values.tolist()替换为：

counts = tf.cast(sample['count'].values, tf.int32)

这样total_count直接是TensorFlow的int32张量，TFP能正确处理它和float32类型logits的类型兼容问题。

方案二：将counts转为Python原生int列表 #

如果需要保留列表形式，先把pandas列的类型转为原生int，再转列表：

counts = sample['count'].astype(int).tolist()

之后再用tf.cast(counts, tf.int32)转换，此时列表元素是Python原生int，TFP的类型推断能正常工作。

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