class_docstr: Constructes a diagonal Gaussian distribution from the hidden states and samples from it using reparametrization.
class_docstr_tok: ["Constructes","a","diagonal","Gaussian","distribution","from","the","hidden","states","and","samples","from","it","using","reparametrization","."]
docstr: We expect here shapes to be [S x B x H] for Sequence, Batch, Hidden sizes (due to tensor parallel support).
  
  inputs:
      hiddens: accepts a tensor of shape [S x B x H]
  
  outputs:
      z: a sample from Gaussian a tensor of shape [S x B x H]
      z_mean: mean of Gaussian a tensor of shape [S x B x H]
      z_logvar: log variance of Gaussian a tensor of shape [S x B x H]
      z_log_prob: log probability of z over posterior log q(z|x) a tensor of shape [S x B x H]
docstr_tok: ["We","expect","here","shapes","to","be","[","S","x","B","x","H","]","for","Sequence",",","Batch",",","Hidden","sizes","(","due","to","tensor","parallel","support",")",".","inputs",":","hiddens",":","accepts","a","tensor","of","shape","[","S","x","B","x","H","]","outputs",":","z",":","a","sample","from","Gaussian","a","tensor","of","shape","[","S","x","B","x","H","]","z_mean",":","mean","of","Gaussian","a","tensor","of","shape","[","S","x","B","x","H","]","z_logvar",":","log","variance","of","Gaussian","a","tensor","of","shape","[","S","x","B","x","H","]","z_log_prob",":","log","probability","of","z","over","posterior","log","q","(","z|x",")","a","tensor","of","shape","[","S","x","B","x","H","]"]
code: def _transform(self, inputs, batch_data=None):
      """
      We expect here shapes to be [S x B x H] for Sequence, Batch, Hidden sizes (due to tensor parallel support).
  
      inputs:
          hiddens: accepts a tensor of shape [S x B x H]
  
      outputs:
          z: a sample from Gaussian a tensor of shape [S x B x H]
          z_mean: mean of Gaussian a tensor of shape [S x B x H]
          z_logvar: log variance of Gaussian a tensor of shape [S x B x H]
          z_log_prob: log probability of z over posterior log q(z|x) a tensor of shape [S x B x H]
      """
      hiddens = inputs["hiddens"]
      # compute distribution's parameters (or use cached ones)
      if "z_mean" in inputs and "z_logvar" in inputs:
          z_mean = inputs["z_mean"]
          z_logvar = inputs["z_logvar"]
      else:
          # ColumnLinear returns output and bias, we ignore bias here (already added to hiddens)
          z_mean, z_logvar = self.hiddens_to_mean_logvar(hiddens)[
              0
          ].chunk(2, dim=-1)
      # clamp logvar
      z_logvar = z_logvar.clamp(min=self.min_logvar)
      # sample z with reparametrization (or use cached one)
      if "z" in inputs:
          z = inputs["z"]
          z_log_prob = inputs.get("z_log_prob", None)
      else:
          e = torch.randn_like(hiddens)
          z = (z_logvar * 0.5).exp() * e + z_mean
          z_log_prob = None
  
      if z_log_prob is None:
          # compute log probability of z under a diagonal Gaussian distribution
          z_log_prob = -0.5 * (
              math.log(2 * math.pi)
              + z_logvar
              + (z - z_mean).pow(2) / z_logvar.exp()
          )
          # sum over the last dimension (hidden_size)
          z_log_prob = z_log_prob.sum(dim=-1)
  
      return {
          "z": z,  # [S x B x H]
          "z_mean": z_mean,  # [S x B x H]
          "z_logvar": z_logvar,  # [S x B x H]
          "z_log_prob": z_log_prob,  # [S x B]
      }
  
file_path: NeMo/nemo/collections/nlp/modules/common/megatron/hiddens/megatron_hidden_transform.py
imports: import math
  import torch
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, init_method_normal
label_desc: Use image node_id 4 for calling the MegatronGaussianHiddenTransform obj's underlying member method code with example usage: obj._transform(inputs, batch_data) and returns: dict