code: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import torch
  
  
  def calc_mean_invstddev(feature):
      if len(feature.size()) != 2:
          raise ValueError("We expect the input feature to be 2-D tensor")
      mean = feature.mean(0)
      var = feature.var(0)
      # avoid division by ~zero
      eps = 1e-8
      if (var < eps).any():
          return mean, 1.0 / (torch.sqrt(var) + eps)
      return mean, 1.0 / torch.sqrt(var)
  
  
  def apply_mv_norm(features):
      # If there is less than 2 spectrograms, the variance cannot be computed (is NaN)
      # and normalization is not possible, so return the item as it is
      if features.size(0) < 2:
          return features
      mean, invstddev = calc_mean_invstddev(features)
      res = (features - mean) * invstddev
      return res
  
  
  def lengths_to_encoder_padding_mask(lengths, batch_first=False):
      """
      convert lengths (a 1-D Long/Int tensor) to 2-D binary tensor
  
      Args:
          lengths: a (B, )-shaped tensor
  
      Return:
          max_length: maximum length of B sequences
          encoder_padding_mask: a (max_length, B) binary mask, where
          [t, b] = 0 for t < lengths[b] and 1 otherwise
  
      TODO:
          kernelize this function if benchmarking shows this function is slow
      """
      max_lengths = torch.max(lengths).item()
      bsz = lengths.size(0)
      encoder_padding_mask = torch.arange(
          max_lengths
      ).to(  # a (T, ) tensor with [0, ..., T-1]
          lengths.device
      ).view(  # move to the right device
          1, max_lengths
      ).expand(  # reshape to (1, T)-shaped tensor
          bsz, -1
      ) >= lengths.view(  # expand to (B, T)-shaped tensor
          bsz, 1
      ).expand(
          -1, max_lengths
      )
      if not batch_first:
          return encoder_padding_mask.t(), max_lengths
      else:
          return encoder_padding_mask, max_lengths
  
  
  def encoder_padding_mask_to_lengths(encoder_padding_mask, max_lengths, batch_size, device):
      """
      convert encoder_padding_mask (2-D binary tensor) to a 1-D tensor
  
      Conventionally, encoder output contains a encoder_padding_mask, which is
      a 2-D mask in a shape (T, B), whose (t, b) element indicate whether
      encoder_out[t, b] is a valid output (=0) or not (=1). Occasionally, we
      need to convert this mask tensor to a 1-D tensor in shape (B, ), where
      [b] denotes the valid length of b-th sequence
  
      Args:
          encoder_padding_mask: a (T, B)-shaped binary tensor or None; if None,
          indicating all are valid
      Return:
          seq_lengths: a (B,)-shaped tensor, where its (b, )-th element is the
          number of valid elements of b-th sequence
  
          max_lengths: maximum length of all sequence, if encoder_padding_mask is
          not None, max_lengths must equal to encoder_padding_mask.size(0)
  
          batch_size: batch size; if encoder_padding_mask is
          not None, max_lengths must equal to encoder_padding_mask.size(1)
  
          device: which device to put the result on
      """
      if encoder_padding_mask is None:
          return torch.Tensor([max_lengths] * batch_size).to(torch.int32).to(device)
  
      assert encoder_padding_mask.size(0) == max_lengths, "max_lengths does not match"
      assert encoder_padding_mask.size(1) == batch_size, "batch_size does not match"
  
      return max_lengths - torch.sum(encoder_padding_mask, dim=0)
  
apis: torch.sum torch.sqrt torch.arange torch.max torch.Tensor