class_docstr: A bijection that generalizes a permutation on the channels of a batch of 2D
  image in :math:`[\ldots,C,H,W]` format. Specifically this transform performs
  the operation,
  
      :math:`\mathbf{y} = \text{torch.nn.functional.conv2d}(\mathbf{x}, W)`
  
  where :math:`\mathbf{x}` are the inputs, :math:`\mathbf{y}` are the outputs,
  and :math:`W\sim C\times C\times 1\times 1` is the filter matrix for a 1x1
  convolution with :math:`C` input and output channels.
  
  Ignoring the final two dimensions, :math:`W` is restricted to be the matrix
  product,
  
      :math:`W = PLU`
  
  where :math:`P\sim C\times C` is a permutation matrix on the channel
  dimensions, :math:`L\sim C\times C` is a lower triangular matrix with ones on
  the diagonal, and :math:`U\sim C\times C` is an upper triangular matrix.
  :math:`W` is initialized to a random orthogonal matrix. Then, :math:`P` is fixed
  and the learnable parameters set to :math:`L,U`.
  
  The input :math:`\mathbf{x}` and output :math:`\mathbf{y}` both have shape
  `[...,C,H,W]`, where `C` is the number of channels set at initialization.
  
  This operation was introduced in [1] for Glow normalizing flow, and is also
  known as 1x1 invertible convolution. It appears in other notable work such as
  [2,3], and corresponds to the class `tfp.bijectors.MatvecLU` of TensorFlow
  Probability.
  
  Example usage:
  
  >>> channels = 3
  >>> base_dist = dist.Normal(torch.zeros(channels, 32, 32),
  ... torch.ones(channels, 32, 32))
  >>> inv_conv = GeneralizedChannelPermute(channels=channels)
  >>> flow_dist = dist.TransformedDistribution(base_dist, [inv_conv])
  >>> flow_dist.sample()  # doctest: +SKIP
  
  :param channels: Number of channel dimensions in the input.
  :type channels: int
  
  [1] Diederik P. Kingma, Prafulla Dhariwal. Glow: Generative Flow with Invertible
  1x1 Convolutions. [arXiv:1807.03039]
  
  [2] Ryan Prenger, Rafael Valle, Bryan Catanzaro. WaveGlow: A Flow-based
  Generative Network for Speech Synthesis. [arXiv:1811.00002]
  
  [3] Conor Durkan, Artur Bekasov, Iain Murray, George Papamakarios. Neural Spline
  Flows. [arXiv:1906.04032]
class_docstr_tok: ["A","bijection","that","generalizes","a","permutation","on","the","channels","of","a","batch","of","2D","image","in",":","math",":","`","[","\\ldots",",","C",",","H",",","W","]","`","format",".","Specifically","this","transform","performs","the","operation",",",":","math",":","`","\\mathbf","{","y","}","=","\\text","{","torch.nn.functional.conv2d","}","(","\\mathbf","{","x","}",",","W",")","`","where",":","math",":","`","\\mathbf","{","x","}","`","are","the","inputs",",",":","math",":","`","\\mathbf","{","y","}","`","are","the","outputs",",","and",":","math",":","`","W\\sim","C\\times","C\\times","1\\times","1","`","is","the","filter","matrix","for","a","1x1","convolution","with",":","math",":","`","C","`","input","and","output","channels",".","Ignoring","the","final","two","dimensions",",",":","math",":","`","W","`","is","restricted","to","be","the","matrix","product",",",":","math",":","`","W","=","PLU","`","where",":","math",":","`","P\\sim","C\\times","C","`","is","a","permutation","matrix","on","the","channel","dimensions",",",":","math",":","`","L\\sim","C\\times","C","`","is","a","lower","triangular","matrix","with","ones","on","the","diagonal",",","and",":","math",":","`","U\\sim","C\\times","C","`","is","an","upper","triangular","matrix",".",":","math",":","`","W","`","is","initialized","to","a","random","orthogonal","matrix",".","Then",",",":","math",":","`","P","`","is","fixed","and","the","learnable","parameters","set","to",":","math",":","`","L",",","U","`",".","The","input",":","math",":","`","\\mathbf","{","x","}","`","and","output",":","math",":","`","\\mathbf","{","y","}","`","both","have","shape","`","[","...",",","C",",","H",",","W","]","`",",","where","`","C","`","is","the","number","of","channels","set","at","initialization",".","This","operation","was","introduced","in","[","1","]","for","Glow","normalizing","flow",",","and","is","also","known","as","1x1","invertible","convolution",".","It","appears","in","other","notable","work","such","as","[","2,3","]",",","and","corresponds","to","the","class","`","tfp.bijectors.MatvecLU","`","of","TensorFlow","Probability",".","Example","usage",":",">",">",">","channels","=","3",">",">",">","base_dist","=","dist.Normal","(","torch.zeros","(","channels",",","32",",","32",")",",","...","torch.ones","(","channels",",","32",",","32",")",")",">",">",">","inv_conv","=","GeneralizedChannelPermute","(","channels=channels",")",">",">",">","flow_dist","=","dist.TransformedDistribution","(","base_dist",",","[","inv_conv","]",")",">",">",">","flow_dist.sample","(",")","#","doctest",":","+SKIP",":","param","channels",":","Number","of","channel","dimensions","in","the","input",".",":","type","channels",":","int","[","1","]","Diederik","P.","Kingma",",","Prafulla","Dhariwal",".","Glow",":","Generative","Flow","with","Invertible","1x1","Convolutions",".","[","arXiv:1807.03039","]","[","2","]","Ryan","Prenger",",","Rafael","Valle",",","Bryan","Catanzaro",".","WaveGlow",":","A","Flow-based","Generative","Network","for","Speech","Synthesis",".","[","arXiv:1811.00002","]","[","3","]","Conor","Durkan",",","Artur","Bekasov",",","Iain","Murray",",","George","Papamakarios",".","Neural","Spline","Flows",".","[","arXiv:1906.04032","]"]
code: def __init__(self, channels=3, permutation=None):
      super(GeneralizedChannelPermute, self).__init__()
      self.__delattr__("permutation")
  
      # Sample a random orthogonal matrix
      W, _ = torch.linalg.qr(torch.randn(channels, channels))
  
      # Construct the partially pivoted LU-form and the pivots
      LU, pivots = torch.linalg.lu_factor(W)
  
      # Convert the pivots into the permutation matrix
      if permutation is None:
          P, _, _ = torch.lu_unpack(LU, pivots)
      else:
          if len(permutation) != channels:
              raise ValueError(
                  'Keyword argument "permutation" expected to have {} elements but {} found.'.format(
                      channels, len(permutation)
                  )
              )
          P = torch.eye(channels, channels)[
              permutation.type(dtype=torch.int64)
          ]
  
      # We register the permutation matrix so that the model can be serialized
      self.register_buffer("permutation", P)
  
      # NOTE: For this implementation I have chosen to store the parameters densely, rather than
      # storing L, U, and s separately
      self.LU = torch.nn.Parameter(LU)
  
imports: import torch
  from torch.distributions import Transform
  from pyro.nn import DenseNN
  from ..None import constraints
  from ..conditional import ConditionalTransformModule
  from ..torch_transform import TransformModule
  from ..util import copy_docs_from
label_desc: Use image node_id 1 to create a new GeneralizedChannelPermute object from inherited base classes: ConditionedGeneralizedChannelPermute, TransformModule with example: obj = GeneralizedChannelPermute(channels, permutation)