Python tfd.Distribution.sample中的复杂事件形状

我正在尝试实现一个表示复杂状态转换模型（STM）的自定义tfd.Distribution。我需要从抽象方法tfd.Distribution的实现中返回一个由两个不同维度的数组组成的元组。但是，当包装器方法（tfd.Distribution.sample）尝试打包这些数组时，我遇到了麻烦

STM的特征是存在于许多相互排斥状态中的种群。随着时间的推移，种群中的个体根据一个随机过程在状态之间转移。为了表示STM的实现（即一个样本），最终得到一个长度为T的向量，其中包含发生转换的时间，以及一个多维形状数组[T，M，N]，其中T是时间步数，M是状态数，N是总体中的个体数

到目前为止，我已经：

Foo类（tfd.分布）：
定义初始化（自）：
super（），
#…这里有其他配置
)
def_样本n（自身、n、种子=无）：
#这里的采样算法
#t.shape=[t]
#y.shape=[T，M，N]
返回t，y
foo=foo（）
foo.sample（）

期望的结果：调用

foo.sample（）

应该返回一个元组（tf.tensor，tf.tensor），其形状分别为[T]和[T，M，N]

实际：

ValueError:形状必须具有相等的秩，但为1和3
将形状0与其他形状合并。对于输入形状为[11]，[11,31000]的“MyEpidemic/sample/Shape/packed”（op:“Pack”）。

看看

联合分布*

类；他们能处理这件事

注意，它们仅在夜间安装tfp的pip中

tfp.distributions.JointDistribution
tfp.distributions.JointDistributionCoroutine
tfp.distributions.JointDistributionCoroutine.Root
tfp.distributions.JointDistributionNamed
tfp.distributions.JointDistributionSequential

---

非常感谢，看起来它可能达到我想要的…除非有点奇怪的语义。是的，我同意它有点奇怪。您的时间步长是自回归的，还是可以并行地对所有[T]进行采样？我假设[T，M，N]样本是以[T]样本为条件的，但它们也是自回归的吗？你看过tfp.sts包吗？是的，自回归——这是一个流行的过程。我已经阅读了tfp.sts——它接近于我所需要的，因为事件形状当然需要有一个“num_time_steps”维度，但是输出的观测值（x_t）位需要是多维的，否则可能会使复杂的数据结构扁平化，从而混淆用户。此外，由于我的模型有一个离散的状态空间，仅仅为了让离散值与时间步长（是浮点，不一定是相等的间隔）处于相同的张量，而将离散值表示为浮点值感觉有点错误。tfd.HiddenMarkovModel适合吗？比sts更适合离散观测。