Method output

ioblocks.py:45–53

float tensor (B, L, D) -> Tuple of locs, scales, and weights

(self, h: T.Tensor)

        return hidden
    
    def output(self, h: T.Tensor) -> Tuple[T.Tensor, T.Tensor, T.Tensor]:
        """float tensor (B, L, D) -> Tuple of locs, scales, and weights"""
        batch_size, seq_len, _ = h.shape
        
        locs = self.fc_loc(h).view(batch_size, seq_len, self.num_components, self.latent_dim)
        scales = T.clamp(self._square_plus(self.fc_scale(h)), min=1e-6).view(batch_size, seq_len, self.num_components, self.latent_dim)
        weights = self.fc_weight(h).view(batch_size, seq_len, self.num_components)
        
        return (locs, scales, weights)

    def loss(self, data, dataHat):
        locs, scales, weights = dataHat