code: def piecewise_quadratic_transform(x, w_tilde, v_tilde, inverse=False):
      """Element-wise piecewise-quadratic transformation
      Args:
          x : torch.Tensor
              *, The variable spans the D-dim unit hypercube ([0,1))
          w_tilde : torch.Tensor
              * x K defined in the paper
          v_tilde : torch.Tensor
              * x (K+1) defined in the paper
          inverse : bool
              forward or inverse
      Returns:
          c : torch.Tensor
              *, transformed value
          log_j : torch.Tensor
              *, log determinant of the Jacobian matrix
      """
      w = torch.softmax(w_tilde, dim=-1)
      v = weighted_softmax(v_tilde, w)
      w_cumsum = torch.cumsum(w, dim=-1)
      # force sum = 1
      w_cumsum[..., -1] = 1.0
      w_cumsum_shift = F.pad(w_cumsum, (1, 0), "constant", 0)
      cdf = torch.cumsum((v[..., 1:] + v[..., :-1]) / 2 * w, dim=-1)
      # force sum = 1
      cdf[..., -1] = 1.0
      cdf_shift = F.pad(cdf, (1, 0), "constant", 0)
  
      if not inverse:
          # * x D x 1, (w_cumsum[idx-1] < x <= w_cumsum[idx])
          bin_index = torch.searchsorted(w_cumsum, x.unsqueeze(-1))
      else:
          # * x D x 1, (cdf[idx-1] < x <= cdf[idx])
          bin_index = torch.searchsorted(cdf, x.unsqueeze(-1))
  
      w_b = torch.gather(w, -1, bin_index).squeeze(-1)
      w_bn1 = torch.gather(w_cumsum_shift, -1, bin_index).squeeze(-1)
      v_b = torch.gather(v, -1, bin_index).squeeze(-1)
      v_bp1 = torch.gather(v, -1, bin_index + 1).squeeze(-1)
      cdf_bn1 = torch.gather(cdf_shift, -1, bin_index).squeeze(-1)
  
      if not inverse:
          alpha = (x - w_bn1) / w_b.clamp(
              min=torch.finfo(w_b.dtype).eps
          )
          c = (
              (alpha**2) / 2 * (v_bp1 - v_b) * w_b
              + alpha * v_b * w_b
              + cdf_bn1
          )
  
          # just sum of log pdfs
          log_j = (
              torch.lerp(v_b, v_bp1, alpha)
              .clamp(min=torch.finfo(c.dtype).eps)
              .log()
          )
  
          # make sure it falls into [0,1)
          c = c.clamp(
              min=torch.finfo(c.dtype).eps,
              max=1.0 - torch.finfo(c.dtype).eps,
          )
          return c, log_j
      else:
          # quadratic equation for alpha
          # alpha should fall into (0, 1]. Since a, b > 0, the symmetry axis -b/2a < 0 and we should pick the larger root
          # skip calculating the log_j in inverse since we don't need it
          a = (v_bp1 - v_b) * w_b / 2
          b = v_b * w_b
          c = cdf_bn1 - x
          alpha = (-b + torch.sqrt((b**2) - 4 * a * c)) / (2 * a)
          inv = alpha * w_b + w_bn1
  
          # make sure it falls into [0,1)
          inv = inv.clamp(
              min=torch.finfo(c.dtype).eps,
              max=1.0 - torch.finfo(inv.dtype).eps,
          )
          return inv, None
  
file_path: NeMo/nemo/collections/tts/parts/utils/splines.py
imports: import numpy
  import torch
  import torch.nn.functional
label_desc: Use image node_id 5 for calling a global function with example usage: piecewise_quadratic_transform(x, w_tilde, v_tilde, inverse) and returns: c,  log_j, inv,  None