Method input

ioblocks.py:40–43

Pre-sampled latents T.Tensor (B, L, Z) -> float tensor (B, L, D)

(self, sampled_latents: T.Tensor)

        return (x + T.sqrt(T.square(x) + 4)) / 2
    
    def input(self, sampled_latents: T.Tensor) -> T.Tensor:
        """Pre-sampled latents T.Tensor (B, L, Z) -> float tensor (B, L, D)"""
        hidden = self.input_projection(sampled_latents)
        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"""