code: def piecewise_linear_transform(
      x, q_tilde, compute_jacobian=True, outlier_passthru=True
  ):
      """Apply an element-wise piecewise-linear transformation to some variables
  
      Args:
          x : torch.Tensor
              a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit
          hypercube
  
          q_tilde: torch.Tensor
                  is a tensor with shape (N,k,b) where b is the number of bins.
                  This contains the un-normalized heights of the bins of the piecewise-constant PDF for dimension k,
                  i.e. q_tilde lives in all of R and we don't impose a constraint on their sum yet.
                  Normalization is imposed in this function using softmax.
  
          compute_jacobian : bool, optional
                              determines whether the jacobian should be compute or None is returned
  
      Returns:
          tuple of torch.Tensor
              pair `(y,h)`.
              - `y` is a tensor with shape (N,k) living in the k-dimensional unit hypercube
              - `j` is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.
      """
      logj = None
  
      third_dimension_softmax = torch.nn.Softmax(dim=2)
  
      # Compute the bin width w
      N, k, b = q_tilde.shape
      Nx, kx = x.shape
      assert N == Nx and k == kx, "Shape mismatch"
  
      w = 1.0 / b
  
      # Compute normalized bin heights with softmax function on bin dimension
      q = 1.0 / w * third_dimension_softmax(q_tilde)
      # x is in the mx-th bin: x \in [0,1],
      # mx \in [[0,b-1]], so we clamp away the case x == 1
      mx = torch.clamp(torch.floor(b * x), 0, b - 1).to(torch.long)
      # Need special error handling because trying to index with mx
      # if it contains nans will lock the GPU. (device-side assert triggered)
      if (
          torch.any(torch.isnan(mx)).item()
          or torch.any(mx < 0)
          or torch.any(mx >= b)
      ):
          raise AvertedCUDARuntimeError(
              "NaN detected in PWLinear bin indexing"
          )
  
      # We compute the output variable in-place
      out = (
          x - mx * w
      )  # alpha (element of [0.,w], the position of x in its bin
  
      # Multiply by the slope
      # q has shape (N,k,b), mxu = mx.unsqueeze(-1) has shape (N,k) with entries that are a b-index
      # gather defines slope[i, j, k] = q[i, j, mxu[i, j, k]] with k taking only 0 as a value
      # i.e. we say slope[i, j] = q[i, j, mx [i, j]]
      slopes = torch.gather(q, 2, mx.unsqueeze(-1)).squeeze(-1)
      out = out * slopes
      # The jacobian is the product of the slopes in all dimensions
  
      # Compute the integral over the left-bins.
      # 1. Compute all integrals: cumulative sum of bin height * bin weight.
      # We want that index i contains the cumsum *strictly to the left* so we shift by 1
      # leaving the first entry null, which is achieved with a roll and assignment
      q_left_integrals = torch.roll(torch.cumsum(q, 2) * w, 1, 2)
      q_left_integrals[:, :, 0] = 0
  
      # 2. Access the correct index to get the left integral of each point and add it to our transformation
      out = out + torch.gather(
          q_left_integrals, 2, mx.unsqueeze(-1)
      ).squeeze(-1)
  
      # Regularization: points must be strictly within the unit hypercube
      # Use the dtype information from pytorch
      eps = torch.finfo(out.dtype).eps
      out = out.clamp(min=eps, max=1.0 - eps)
      oob_mask = torch.logical_or(x < 0.0, x > 1.0).detach().float()
      if outlier_passthru:
          out = out * (1 - oob_mask) + x * oob_mask
          slopes = slopes * (1 - oob_mask) + oob_mask
  
      if compute_jacobian:
          # logj = torch.log(torch.prod(slopes.float(), 1))
          logj = torch.sum(torch.log(slopes), 1)
      del slopes
  
      return out, logj
  
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 1 for calling a global function with example usage: piecewise_linear_transform(x, q_tilde, compute_jacobian, outlier_passthru) and returns: out,  logj