Why is forward method for ANP maximizing LogLikelihood for Context labels?

[yair-schiff]

In the ANP implementation, during training mode, the full set (i.e. context + target) is being passed to predict meaning that log likelihood is being maximized for both target and context labels.

        if self.training:
            pz = self.lenc(batch.xc, batch.yc)
            qz = self.lenc(batch.x, batch.y)
            z = qz.rsample() if num_samples is None else \
                    qz.rsample([num_samples])
            py = self.predict(batch.xc, batch.yc, batch.x,
                    z=z, num_samples=num_samples)


            if num_samples > 1:
                # K * B * N
                recon = py.log_prob(stack(batch.y, num_samples)).sum(-1)
                # K * B
                log_qz = qz.log_prob(z).sum(-1)
                log_pz = pz.log_prob(z).sum(-1)


                # K * B
                log_w = recon.sum(-1) + log_pz - log_qz


                outs.loss = -logmeanexp(log_w).mean() / batch.x.shape[-2]

Shouldn't this line:

            py = self.predict(batch.xc, batch.yc, batch.x,

be changed to this

            py = self.predict(batch.xc, batch.yc, batch.xt,

i.e. batch.x --> batch.xt in 3rd argument of predict?