class_docstr: A 2D FIR downsampling layer with an optional convolution.
  
  Parameters:
      channels (`int`):
          number of channels in the inputs and outputs.
      use_conv (`bool`, default `False`):
          option to use a convolution.
      out_channels (`int`, optional):
          number of output channels. Defaults to `channels`.
      fir_kernel (`tuple`, default `(1, 3, 3, 1)`):
          kernel for the FIR filter.
class_docstr_tok: ["A","2D","FIR","downsampling","layer","with","an","optional","convolution",".","Parameters",":","channels","(","`","int","`",")",":","number","of","channels","in","the","inputs","and","outputs",".","use_conv","(","`","bool","`",",","default","`","False","`",")",":","option","to","use","a","convolution",".","out_channels","(","`","int","`",",","optional",")",":","number","of","output","channels",".","Defaults","to","`","channels","`",".","fir_kernel","(","`","tuple","`",",","default","`","(","1",",","3",",","3",",","1",")","`",")",":","kernel","for","the","FIR","filter","."]
docstr: Fused `Conv2d()` followed by `downsample_2d()`.
  Padding is performed only once at the beginning, not between the operations. The fused op is considerably more
  efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of
  arbitrary order.
  
  Args:
      hidden_states (`torch.FloatTensor`):
          Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
      weight (`torch.FloatTensor`, *optional*):
          Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. Grouped convolution can be
          performed by `inChannels = x.shape[0] // numGroups`.
      kernel (`torch.FloatTensor`, *optional*):
          FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which
          corresponds to average pooling.
      factor (`int`, *optional*, default to `2`):
          Integer downsampling factor.
      gain (`float`, *optional*, default to `1.0`):
          Scaling factor for signal magnitude.
  
  Returns:
      output (`torch.FloatTensor`):
          Tensor of the shape `[N, C, H // factor, W // factor]` or `[N, H // factor, W // factor, C]`, and same
          datatype as `x`.
docstr_tok: ["Fused","`","Conv2d","(",")","`","followed","by","`","downsample_2d","(",")","`",".","Padding","is","performed","only","once","at","the","beginning",",","not","between","the","operations",".","The","fused","op","is","considerably","more","efficient","than","performing","the","same","calculation","using","standard","TensorFlow","ops",".","It","supports","gradients","of","arbitrary","order",".","Args",":","hidden_states","(","`","torch.FloatTensor","`",")",":","Input","tensor","of","the","shape","`","[","N",",","C",",","H",",","W","]","`","or","`","[","N",",","H",",","W",",","C","]","`",".","weight","(","`","torch.FloatTensor","`",",","*","optional","*",")",":","Weight","tensor","of","the","shape","`","[","filterH",",","filterW",",","inChannels",",","outChannels","]","`",".","Grouped","convolution","can","be","performed","by","`","inChannels","=","x.shape","[","0","]","\/\/","numGroups","`",".","kernel","(","`","torch.FloatTensor","`",",","*","optional","*",")",":","FIR","filter","of","the","shape","`","[","firH",",","firW","]","`","or","`","[","firN","]","`","(","separable",")",".","The","default","is","`","[","1","]","*","factor","`",",","which","corresponds","to","average","pooling",".","factor","(","`","int","`",",","*","optional","*",",","default","to","`","2","`",")",":","Integer","downsampling","factor",".","gain","(","`","float","`",",","*","optional","*",",","default","to","`","1.0","`",")",":","Scaling","factor","for","signal","magnitude",".","Returns",":","output","(","`","torch.FloatTensor","`",")",":","Tensor","of","the","shape","`","[","N",",","C",",","H","\/\/","factor",",","W","\/\/","factor","]","`","or","`","[","N",",","H","\/\/","factor",",","W","\/\/","factor",",","C","]","`",",","and","same","datatype","as","`","x","`","."]
code: def _downsample_2d(
      self,
      hidden_states: torch.FloatTensor,
      weight: Optional[torch.FloatTensor] = None,
      kernel: Optional[torch.FloatTensor] = None,
      factor: int = 2,
      gain: float = 1,
  ) -> torch.FloatTensor:
      """Fused `Conv2d()` followed by `downsample_2d()`.
      Padding is performed only once at the beginning, not between the operations. The fused op is considerably more
      efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of
      arbitrary order.
  
      Args:
          hidden_states (`torch.FloatTensor`):
              Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
          weight (`torch.FloatTensor`, *optional*):
              Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. Grouped convolution can be
              performed by `inChannels = x.shape[0] // numGroups`.
          kernel (`torch.FloatTensor`, *optional*):
              FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which
              corresponds to average pooling.
          factor (`int`, *optional*, default to `2`):
              Integer downsampling factor.
          gain (`float`, *optional*, default to `1.0`):
              Scaling factor for signal magnitude.
  
      Returns:
          output (`torch.FloatTensor`):
              Tensor of the shape `[N, C, H // factor, W // factor]` or `[N, H // factor, W // factor, C]`, and same
              datatype as `x`.
      """
  
      assert isinstance(factor, int) and factor >= 1
      if kernel is None:
          kernel = [1] * factor
  
      # setup kernel
      kernel = torch.tensor(kernel, dtype=torch.float32)
      if kernel.ndim == 1:
          kernel = torch.outer(kernel, kernel)
      kernel /= torch.sum(kernel)
  
      kernel = kernel * gain
  
      if self.use_conv:
          _, _, convH, convW = weight.shape
          pad_value = (kernel.shape[0] - factor) + (convW - 1)
          stride_value = [factor, factor]
          upfirdn_input = upfirdn2d_native(
              hidden_states,
              torch.tensor(kernel, device=hidden_states.device),
              pad=((pad_value + 1) // 2, pad_value // 2),
          )
          output = F.conv2d(
              upfirdn_input, weight, stride=stride_value, padding=0
          )
      else:
          pad_value = kernel.shape[0] - factor
          output = upfirdn2d_native(
              hidden_states,
              torch.tensor(kernel, device=hidden_states.device),
              down=factor,
              pad=((pad_value + 1) // 2, pad_value // 2),
          )
  
      return output
  
imports: from typing import Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from ..utils import USE_PEFT_BACKEND
  from .lora import LoRACompatibleConv
  from .normalization import RMSNorm
  from .upsampling import upfirdn2d_native
label_desc: Use image node_id 2 for calling the FirDownsample2D obj's underlying member method code with example usage: obj._downsample_2d(hidden_states, weight, kernel, factor, gain) and returns: output