有没有一种体面的方法可以在 tensorflow_probability 中的双射器中植入一些双射器？

Question

假设我想构建一个双射器，其中有两个 tfp.bijector。因为它是一个 tfb.Bijector，我必须实现 _forward_log_det_jacobian 方法。

class Bij(tfb.Bijector):
    def __init___(...):
        super().__init__(..., forward_min_event_ndims=3)
        self.bi1 = Bijector1()
        self.bi2 = Bijector2()

    def _forward(self, x):
        x = self.bi1.forward(x)
        # some ops on x here
        return self.bi2.forward(x)
        # I hope the way it forwards is right, I can't just chain them up cuz some ops on x between bi1 and bi2. If there's better way to achieve this, pls let me know!

    ...
    
    def _forward_log_det_jacobian(self, x):
        # I am a bit confused here. I use bi1.forward_log_det_jacobian(x) to get
        # jaconbian then summed it with bi2.forward_log_det_jacobian(x) with parameter
        # event_ndims = forward_min_event_ndims in both forward_log_det_jacobian().

输出似乎是正确的，但我不太确定，因为bi1 是一个包含forward_min_event_ndims=1(

有没有推荐的方法来实现这样的目标？还是我写的正确？

Answer 1

你能用tfb.Chain吗？ 双射器还实现了一个__call__ 方法，该方法将为您提供一个Chain 实例，即chain = bij1(bij2)。