class_docstr: An implementation of the coupling layer with rational spline bijections of
  linear and quadratic order (Durkan et al., 2019; Dolatabadi et al., 2020).
  Rational splines are functions that are comprised of segments that are the ratio
  of two polynomials (see :class:`~pyro.distributions.transforms.Spline`).
  
  The spline coupling layer uses the transformation,
  
      :math:`\mathbf{y}_{1:d} = g_\theta(\mathbf{x}_{1:d})`
      :math:`\mathbf{y}_{(d+1):D} = h_\phi(\mathbf{x}_{(d+1):D};\mathbf{x}_{1:d})`
  
  where :math:`\mathbf{x}` are the inputs, :math:`\mathbf{y}` are the outputs,
  e.g. :math:`\mathbf{x}_{1:d}` represents the first :math:`d` elements of the
  inputs, :math:`g_\theta` is either the identity function or an elementwise
  rational monotonic spline with parameters :math:`\theta`, and :math:`h_\phi` is
  a conditional elementwise spline spline, conditioning on the first :math:`d`
  elements.
  
  Example usage:
  
  >>> from pyro.nn import DenseNN
  >>> input_dim = 10
  >>> split_dim = 6
  >>> count_bins = 8
  >>> base_dist = dist.Normal(torch.zeros(input_dim), torch.ones(input_dim))
  >>> param_dims = [(input_dim - split_dim) * count_bins,
  ... (input_dim - split_dim) * count_bins,
  ... (input_dim - split_dim) * (count_bins - 1),
  ... (input_dim - split_dim) * count_bins]
  >>> hypernet = DenseNN(split_dim, [10*input_dim], param_dims)
  >>> transform = SplineCoupling(input_dim, split_dim, hypernet)
  >>> pyro.module("my_transform", transform)  # doctest: +SKIP
  >>> flow_dist = dist.TransformedDistribution(base_dist, [transform])
  >>> flow_dist.sample()  # doctest: +SKIP
  
  :param input_dim: Dimension of the input vector. Despite operating element-wise,
      this is required so we know how many parameters to store.
  :type input_dim: int
  :param split_dim: Zero-indexed dimension :math:`d` upon which to perform input/
      output split for transformation.
  :param hypernet: a neural network whose forward call returns a tuple of spline
      parameters (see :class:`~pyro.distributions.transforms.ConditionalSpline`).
  :type hypernet: callable
  :param count_bins: The number of segments comprising the spline.
  :type count_bins: int
  :param bound: The quantity :math:`K` determining the bounding box,
      :math:`[-K,K]\times[-K,K]`, of the spline.
  :type bound: float
  :param order: One of ['linear', 'quadratic'] specifying the order of the spline.
  :type order: string
  
  References:
  
  Conor Durkan, Artur Bekasov, Iain Murray, George Papamakarios. Neural
  Spline Flows. NeurIPS 2019.
  
  Hadi M. Dolatabadi, Sarah Erfani, Christopher Leckie. Invertible Generative
  Modeling using Linear Rational Splines. AISTATS 2020.
class_docstr_tok: ["An","implementation","of","the","coupling","layer","with","rational","spline","bijections","of","linear","and","quadratic","order","(","Durkan","et","al.",",","2019",";","Dolatabadi","et","al.",",","2020",")",".","Rational","splines","are","functions","that","are","comprised","of","segments","that","are","the","ratio","of","two","polynomials","(","see",":","class",":","`","~pyro.distributions.transforms.Spline","`",")",".","The","spline","coupling","layer","uses","the","transformation",",",":","math",":","`","\\mathbf","{","y","}","_","{","1",":","d","}","=","g_\\theta","(","\\mathbf","{","x","}","_","{","1",":","d","}",")","`",":","math",":","`","\\mathbf","{","y","}","_","{","(","d+1",")",":","D","}","=","h_\\phi","(","\\mathbf","{","x","}","_","{","(","d+1",")",":","D","}",";","\\mathbf","{","x","}","_","{","1",":","d","}",")","`","where",":","math",":","`","\\mathbf","{","x","}","`","are","the","inputs",",",":","math",":","`","\\mathbf","{","y","}","`","are","the","outputs",",","e.g",".",":","math",":","`","\\mathbf","{","x","}","_","{","1",":","d","}","`","represents","the","first",":","math",":","`","d","`","elements","of","the","inputs",",",":","math",":","`","g_\\theta","`","is","either","the","identity","function","or","an","elementwise","rational","monotonic","spline","with","parameters",":","math",":","`","\\theta","`",",","and",":","math",":","`","h_\\phi","`","is","a","conditional","elementwise","spline","spline",",","conditioning","on","the","first",":","math",":","`","d","`","elements",".","Example","usage",":",">",">",">","from","pyro.nn","import","DenseNN",">",">",">","input_dim","=","10",">",">",">","split_dim","=","6",">",">",">","count_bins","=","8",">",">",">","base_dist","=","dist.Normal","(","torch.zeros","(","input_dim",")",",","torch.ones","(","input_dim",")",")",">",">",">","param_dims","=","[","(","input_dim","-","split_dim",")","*","count_bins",",","...","(","input_dim","-","split_dim",")","*","count_bins",",","...","(","input_dim","-","split_dim",")","*","(","count_bins","-","1",")",",","...","(","input_dim","-","split_dim",")","*","count_bins","]",">",">",">","hypernet","=","DenseNN","(","split_dim",",","[","10","*","input_dim","]",",","param_dims",")",">",">",">","transform","=","SplineCoupling","(","input_dim",",","split_dim",",","hypernet",")",">",">",">","pyro.module","(","``","my_transform","''",",","transform",")","#","doctest",":","+SKIP",">",">",">","flow_dist","=","dist.TransformedDistribution","(","base_dist",",","[","transform","]",")",">",">",">","flow_dist.sample","(",")","#","doctest",":","+SKIP",":","param","input_dim",":","Dimension","of","the","input","vector",".","Despite","operating","element-wise",",","this","is","required","so","we","know","how","many","parameters","to","store",".",":","type","input_dim",":","int",":","param","split_dim",":","Zero-indexed","dimension",":","math",":","`","d","`","upon","which","to","perform","input\/","output","split","for","transformation",".",":","param","hypernet",":","a","neural","network","whose","forward","call","returns","a","tuple","of","spline","parameters","(","see",":","class",":","`","~pyro.distributions.transforms.ConditionalSpline","`",")",".",":","type","hypernet",":","callable",":","param","count_bins",":","The","number","of","segments","comprising","the","spline",".",":","type","count_bins",":","int",":","param","bound",":","The","quantity",":","math",":","`","K","`","determining","the","bounding","box",",",":","math",":","`","[","-K",",","K","]","\\times","[","-K",",","K","]","`",",","of","the","spline",".",":","type","bound",":","float",":","param","order",":","One","of","[","'linear","'",",","'quadratic","'","]","specifying","the","order","of","the","spline",".",":","type","order",":","string","References",":","Conor","Durkan",",","Artur","Bekasov",",","Iain","Murray",",","George","Papamakarios",".","Neural","Spline","Flows",".","NeurIPS","2019",".","Hadi","M.","Dolatabadi",",","Sarah","Erfani",",","Christopher","Leckie",".","Invertible","Generative","Modeling","using","Linear","Rational","Splines",".","AISTATS","2020","."]
docstr: :param y: the output of the bijection
  :type y: torch.Tensor
  
  Inverts y => x. Uses a previously cached inverse if available,
  otherwise performs the inversion afresh.
docstr_tok: [":","param","y",":","the","output","of","the","bijection",":","type","y",":","torch.Tensor","Inverts","y","=",">","x",".","Uses","a","previously","cached","inverse","if","available",",","otherwise","performs","the","inversion","afresh","."]
code: def _inverse(self, y):
      """
      :param y: the output of the bijection
      :type y: torch.Tensor
  
      Inverts y => x. Uses a previously cached inverse if available,
      otherwise performs the inversion afresh.
      """
      y1, y2 = y[..., : self.split_dim], y[..., self.split_dim :]
  
      if not self.identity:
          x1 = self.lower_spline._inv_call(y1)
          log_detK = self.lower_spline._cache_log_detJ
      else:
          x1 = y1
  
      upper_spline = self.upper_spline.condition(x1)
      x2 = upper_spline._inv_call(y2)
      log_detJ = upper_spline._cache_log_detJ
  
      if not self.identity:
          log_detJ = torch.cat([log_detJ, log_detK], dim=-1)
      self._cache_log_detJ = log_detJ
  
      return torch.cat([x1, x2], dim=-1)
  
file_path: pyro/pyro/distributions/transforms/spline_coupling.py
imports: import torch
  from pyro.nn import DenseNN
  from ..None import constraints
  from ..torch_transform import TransformModule
  from ..util import copy_docs_from
  from .spline import ConditionalSpline, Spline
label_desc: Use image node_id 3 for calling the SplineCoupling obj's underlying member method code with example usage: obj._inverse(y) and returns: torch