{"kind":"task","effective_mode":"full","benchmark":{"kind":"benchmark","effective_mode":"full","slug":"longbench-v2","formal_name":"LongBench v2","introduction":"長い資料の深い理解と推論を、多肢選択問題で評価するベンチマークです。公式紹介では503問を収録し、単一・複数文書の質問応答やコードリポジトリ理解などを扱います。\n\nLongBench v2 evaluates deep understanding and reasoning over long contexts through multiple-choice questions. Its official description lists 503 questions spanning tasks such as single-document and multi-document QA and code-repository understanding.","introduction_ja":"","introduction_en":"","category":"Category not supplied","task_count":null,"acquisition_status":"Acquisition status not supplied","official_url":"https://huggingface.co/datasets/zai-org/LongBench-v2","indexing_mode":"noindex"},"task_id":"3b801dd0-be9a-59e6-a415-983af4e86133","task_key":"train--66f1dac1821e116aacb27df1","task_revision_id":"1","upstream_id":"66f1dac1821e116aacb27df1","short_description":"This is the troch.nn modeule. In this module, there exists an implementation of…","config":"","split":"train","body":"{\"choice_A\":\"_DEFAULT_SPARSE_BLOCK_SIZE, _ModificationType.SCORE_MOD\",\"choice_B\":\"_LARGE_SPARSE_BLOCK_SIZE, _ModificationType.MASK_MOD\",\"choice_C\":\"_DEFAULT_SPARSE_BLOCK_SIZE, _ModificationType.MASK_MOD\",\"choice_D\":\"_LARGE_SPARSE_BLOCK_SIZE, _ModificationType.SCORE_MOD\",\"context\":\"# mypy: allow-untyped-defs\\n\\\"\\\"\\\"Functionality for Python <-> C++ frontend inter-op.\\\"\\\"\\\"\\n\\nfrom torch import nn\\n\\n\\nclass OrderedDictWrapper:\\n    \\\"\\\"\\\"A wrapper around a C++ OrderedDict.\\n\\n    It dynamically evaluates the OrderedDict getter on a bound C++ module, such\\n    that new changes on the C++ side are picked up. Otherwise accessing e.g.\\n    ``cpp_module._parameters`` just once would get a frozen copy of the parameters\\n    at the time of access. ``torch.nn.Module`` accesses ``_parameters`` et al. via ``self.__dict__``\\n    so using properties does not work.\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, cpp_module, attr):\\n        self.cpp_module = cpp_module\\n        self.attr = attr\\n\\n    @property\\n    def cpp_dict(self):\\n        return getattr(self.cpp_module, self.attr)\\n\\n    # Magic methods cannot be assigned dynamically and bypass ``getattr``, so we\\n    # must manually override them.\\n\\n    def items(self):\\n        return self.cpp_dict.items()\\n\\n    def keys(self):\\n        return self.cpp_dict.keys()\\n\\n    def values(self):\\n        return self.cpp_dict.values()\\n\\n    def __iter__(self):\\n        return self.cpp_dict.__iter__()\\n\\n    def __len__(self):\\n        return self.cpp_dict.__len__()\\n\\n    def __contains__(self, key):\\n        return self.cpp_dict.__contains__(key)\\n\\n    def __getitem__(self, key):\\n        return self.cpp_dict.__getitem__(key)\\n\\n\\nclass ModuleWrapper(nn.Module):\\n    \\\"\\\"\\\"A subclass of ``torch.nn.Module`` that wraps a C++ frontend module and delegates all access.\\\"\\\"\\\"\\n\\n    def __init__(self, cpp_module):\\n        # Assign before the super class constructor so ``self.training`` can be\\n        # assigned to in the super class constructor.\\n        self.cpp_module = cpp_module\\n        super().__init__()\\n        self._parameters = OrderedDictWrapper(cpp_module, \\\"_parameters\\\")  # type: ignore[assignment]\\n        self._buffers: OrderedDictWrapper = OrderedDictWrapper(cpp_module, \\\"_buffers\\\")  # type: ignore[assignment]\\n        self._modules: OrderedDictWrapper = OrderedDictWrapper(cpp_module, \\\"_modules\\\")  # type: ignore[assignment]\\n        for attr in dir(cpp_module):\\n            # Skip magic methods and the three attributes above.\\n            if not attr.startswith(\\\"_\\\"):\\n                setattr(self, attr, getattr(self.cpp_module, attr))\\n\\n    def _apply(self, fn, recurse=True):\\n        for param in self.parameters():\\n            # Tensors stored in modules are graph leaves, and we don't\\n            # want to create copy nodes, so we have to unpack the data.\\n            param.data = fn(param.data)\\n            if param._grad is not None:\\n                param._grad.data = fn(param._grad.data)\\n\\n        for buf in self.buffers():\\n            buf.data = fn(buf.data)\\n\\n        return self\\n\\n    # nn.Module defines training as a boolean\\n    @property  # type: ignore[override]\\n    def training(self):\\n        return self.cpp_module.training\\n\\n    @training.setter\\n    def training(self, mode):\\n        self.cpp_module.train(mode)\\n\\n    def __repr__(self):\\n        return self.cpp_module.__repr__()\\n\\n\\n\\\"\\\"\\\"Functional interface.\\\"\\\"\\\"\\n\\nimport importlib\\nimport math\\nimport warnings\\nfrom typing import Callable, List, Optional, Tuple, TYPE_CHECKING, Union\\n\\nimport torch\\nfrom torch import _VF, sym_int as _sym_int, Tensor\\nfrom torch._C import _add_docstr, _infer_size\\nfrom torch._jit_internal import (\\n    _overload,\\n    boolean_dispatch,\\n    BroadcastingList1,\\n    BroadcastingList2,\\n    BroadcastingList3,\\n)\\nfrom torch._torch_docs import reproducibility_notes, sparse_support_notes, tf32_notes\\nfrom torch.nn import _reduction as _Reduction, grad  # noqa: F401\\nfrom torch.nn.modules.utils import _list_with_default, _pair, _single, _triple\\nfrom torch.overrides import (\\n    handle_torch_function,\\n    has_torch_function,\\n    has_torch_function_unary,\\n    has_torch_function_variadic,\\n)\\n\\n\\nif TYPE_CHECKING:\\n    from torch.types import _dtype as DType\\nelse:\\n    # The JIT doesn't understand Union, nor torch.dtype here\\n    DType = int\\n\\ntry:\\n    import numpy as np\\nexcept ModuleNotFoundError:\\n    np = None\\n\\n\\nconv1d = _add_docstr(\\n    torch.conv1d,\\n    r\\\"\\\"\\\"\\nconv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor\\n\\nApplies a 1D convolution over an input signal composed of several input\\nplanes.\\n\\n{tf32_note}\\n\\nSee :class:`~torch.nn.Conv1d` for details and output shape.\\n\\nNote:\\n    {cudnn_reproducibility_note}\\n\\nNote:\\n    This operator supports complex data types i.e. ``complex32, complex64, complex128``.\\n\\\"\\\"\\\".format(\\n        **reproducibility_notes, **tf32_notes\\n    )\\n    + r\\\"\\\"\\\"\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iW)`\\n    weight: filters of shape :math:`(\\\\text{out\\\\_channels} , \\\\frac{\\\\text{in\\\\_channels}}{\\\\text{groups}} , kW)`\\n    bias: optional bias of shape :math:`(\\\\text{out\\\\_channels})`. Default: ``None``\\n    stride: the stride of the convolving kernel. Can be a single number or\\n      a one-element tuple `(sW,)`. Default: 1\\n    padding: implicit paddings on both sides of the input. Can be a string {'valid', 'same'},\\n      single number or a one-element tuple `(padW,)`. Default: 0\\n      ``padding='valid'`` is the same as no padding. ``padding='same'`` pads\\n      the input so the output has the same shape as the input. However, this mode\\n      doesn't support any stride values other than 1.\\n\\n      .. warning::\\n          For ``padding='same'``, if the ``weight`` is even-length and\\n          ``dilation`` is odd in any dimension, a full :func:`pad` operation\\n          may be needed internally. Lowering performance.\\n    dilation: the spacing between kernel elements. Can be a single number or\\n      a one-element tuple `(dW,)`. Default: 1\\n    groups: split input into groups, :math:`\\\\text{in\\\\_channels}` should be divisible by\\n      the number of groups. Default: 1\\n\\nExamples::\\n\\n    >>> inputs = torch.randn(33, 16, 30)\\n    >>> filters = torch.randn(20, 16, 5)\\n    >>> F.conv1d(inputs, filters)\\n\\\"\\\"\\\",\\n)\\n\\nconv2d = _add_docstr(\\n    torch.conv2d,\\n    r\\\"\\\"\\\"\\nconv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor\\n\\nApplies a 2D convolution over an input image composed of several input\\nplanes.\\n\\n{tf32_note}\\n\\nSee :class:`~torch.nn.Conv2d` for details and output shape.\\n\\nNote:\\n    {cudnn_reproducibility_note}\\n\\nNote:\\n    This operator supports complex data types i.e. ``complex32, complex64, complex128``.\\n\\\"\\\"\\\".format(\\n        **reproducibility_notes, **tf32_notes\\n    )\\n    + r\\\"\\\"\\\"\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iH , iW)`\\n    weight: filters of shape :math:`(\\\\text{out\\\\_channels} , \\\\frac{\\\\text{in\\\\_channels}}{\\\\text{groups}} , kH , kW)`\\n    bias: optional bias tensor of shape :math:`(\\\\text{out\\\\_channels})`. Default: ``None``\\n    stride: the stride of the convolving kernel. Can be a single number or a\\n      tuple `(sH, sW)`. Default: 1\\n    padding: implicit paddings on both sides of the input. Can be a string {'valid', 'same'},\\n      single number or a tuple `(padH, padW)`. Default: 0\\n      ``padding='valid'`` is the same as no padding. ``padding='same'`` pads\\n      the input so the output has the same shape as the input. However, this mode\\n      doesn't support any stride values other than 1.\\n\\n      .. warning::\\n          For ``padding='same'``, if the ``weight`` is even-length and\\n          ``dilation`` is odd in any dimension, a full :func:`pad` operation\\n          may be needed internally. Lowering performance.\\n\\n    dilation: the spacing between kernel elements. Can be a single number or\\n      a tuple `(dH, dW)`. Default: 1\\n    groups: split input into groups, both :math:`\\\\text{in\\\\_channels}` and :math:`\\\\text{out\\\\_channels}`\\n      should be divisible by the number of groups. Default: 1\\n\\nExamples::\\n\\n    >>> # With square kernels and equal stride\\n    >>> filters = torch.randn(8, 4, 3, 3)\\n    >>> inputs = torch.randn(1, 4, 5, 5)\\n    >>> F.conv2d(inputs, filters, padding=1)\\n\\\"\\\"\\\",\\n)  # noqa: E501\\n\\nconv3d = _add_docstr(\\n    torch.conv3d,\\n    r\\\"\\\"\\\"\\nconv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor\\n\\nApplies a 3D convolution over an input image composed of several input\\nplanes.\\n\\n{tf32_note}\\n\\nSee :class:`~torch.nn.Conv3d` for details and output shape.\\n\\nNote:\\n    {cudnn_reproducibility_note}\\n\\nNote:\\n    This operator supports complex data types i.e. ``complex32, complex64, complex128``.\\n\\\"\\\"\\\".format(\\n        **reproducibility_notes, **tf32_notes\\n    )\\n    + r\\\"\\\"\\\"\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iT , iH , iW)`\\n    weight: filters of shape :math:`(\\\\text{out\\\\_channels} , \\\\frac{\\\\text{in\\\\_channels}}{\\\\text{groups}} , kT , kH , kW)`\\n    bias: optional bias tensor of shape :math:`(\\\\text{out\\\\_channels})`. Default: None\\n    stride: the stride of the convolving kernel. Can be a single number or a\\n      tuple `(sT, sH, sW)`. Default: 1\\n    padding: implicit paddings on both sides of the input. Can be a string {'valid', 'same'},\\n      single number or a tuple `(padT, padH, padW)`. Default: 0\\n      ``padding='valid'`` is the same as no padding. ``padding='same'`` pads\\n      the input so the output has the same shape as the input. However, this mode\\n      doesn't support any stride values other than 1.\\n\\n      .. warning::\\n          For ``padding='same'``, if the ``weight`` is even-length and\\n          ``dilation`` is odd in any dimension, a full :func:`pad` operation\\n          may be needed internally. Lowering performance.\\n\\n    dilation: the spacing between kernel elements. Can be a single number or\\n      a tuple `(dT, dH, dW)`. Default: 1\\n    groups: split input into groups, :math:`\\\\text{in\\\\_channels}` should be divisible by\\n      the number of groups. Default: 1\\n\\nExamples::\\n\\n    >>> filters = torch.randn(33, 16, 3, 3, 3)\\n    >>> inputs = torch.randn(20, 16, 50, 10, 20)\\n    >>> F.conv3d(inputs, filters)\\n\\\"\\\"\\\",\\n)  # noqa: E501\\n\\nconv_transpose1d = _add_docstr(\\n    torch.conv_transpose1d,\\n    r\\\"\\\"\\\"\\nconv_transpose1d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor\\n\\nApplies a 1D transposed convolution operator over an input signal\\ncomposed of several input planes, sometimes also called \\\"deconvolution\\\".\\n\\n{tf32_note}\\n\\nSee :class:`~torch.nn.ConvTranspose1d` for details and output shape.\\n\\nNote:\\n    {cudnn_reproducibility_note}\\n\\\"\\\"\\\".format(\\n        **reproducibility_notes, **tf32_notes\\n    )\\n    + r\\\"\\\"\\\"\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iW)`\\n    weight: filters of shape :math:`(\\\\text{in\\\\_channels} , \\\\frac{\\\\text{out\\\\_channels}}{\\\\text{groups}} , kW)`\\n    bias: optional bias of shape :math:`(\\\\text{out\\\\_channels})`. Default: None\\n    stride: the stride of the convolving kernel. Can be a single number or a\\n      tuple ``(sW,)``. Default: 1\\n    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both\\n      sides of each dimension in the input. Can be a single number or a tuple\\n      ``(padW,)``. Default: 0\\n    output_padding: additional size added to one side of each dimension in the\\n      output shape. Can be a single number or a tuple ``(out_padW)``. Default: 0\\n    groups: split input into groups, :math:`\\\\text{in\\\\_channels}` should be divisible by the\\n      number of groups. Default: 1\\n    dilation: the spacing between kernel elements. Can be a single number or\\n      a tuple ``(dW,)``. Default: 1\\n\\nExamples::\\n\\n    >>> inputs = torch.randn(20, 16, 50)\\n    >>> weights = torch.randn(16, 33, 5)\\n    >>> F.conv_transpose1d(inputs, weights)\\n\\\"\\\"\\\",\\n)\\n\\nconv_transpose2d = _add_docstr(\\n    torch.conv_transpose2d,\\n    r\\\"\\\"\\\"\\nconv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor\\n\\nApplies a 2D transposed convolution operator over an input image\\ncomposed of several input planes, sometimes also called \\\"deconvolution\\\".\\n\\n{tf32_note}\\n\\nSee :class:`~torch.nn.ConvTranspose2d` for details and output shape.\\n\\nNote:\\n    {cudnn_reproducibility_note}\\n\\\"\\\"\\\".format(\\n        **reproducibility_notes, **tf32_notes\\n    )\\n    + r\\\"\\\"\\\"\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iH , iW)`\\n    weight: filters of shape :math:`(\\\\text{in\\\\_channels} , \\\\frac{\\\\text{out\\\\_channels}}{\\\\text{groups}} , kH , kW)`\\n    bias: optional bias of shape :math:`(\\\\text{out\\\\_channels})`. Default: None\\n    stride: the stride of the convolving kernel. Can be a single number or a\\n      tuple ``(sH, sW)``. Default: 1\\n    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both\\n      sides of each dimension in the input. Can be a single number or a tuple\\n      ``(padH, padW)``. Default: 0\\n    output_padding: additional size added to one side of each dimension in the\\n      output shape. Can be a single number or a tuple ``(out_padH, out_padW)``.\\n      Default: 0\\n    groups: split input into groups, :math:`\\\\text{in\\\\_channels}` should be divisible by the\\n      number of groups. Default: 1\\n    dilation: the spacing between kernel elements. Can be a single number or\\n      a tuple ``(dH, dW)``. Default: 1\\n\\nExamples::\\n\\n    >>> # With square kernels and equal stride\\n    >>> inputs = torch.randn(1, 4, 5, 5)\\n    >>> weights = torch.randn(4, 8, 3, 3)\\n    >>> F.conv_transpose2d(inputs, weights, padding=1)\\n\\\"\\\"\\\",\\n)  # noqa: E501\\n\\nconv_transpose3d = _add_docstr(\\n    torch.conv_transpose3d,\\n    r\\\"\\\"\\\"\\nconv_transpose3d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor\\n\\nApplies a 3D transposed convolution operator over an input image\\ncomposed of several input planes, sometimes also called \\\"deconvolution\\\"\\n\\n{tf32_note}\\n\\nSee :class:`~torch.nn.ConvTranspose3d` for details and output shape.\\n\\nNote:\\n    {cudnn_reproducibility_note}\\n\\\"\\\"\\\".format(\\n        **reproducibility_notes, **tf32_notes\\n    )\\n    + r\\\"\\\"\\\"\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iT , iH , iW)`\\n    weight: filters of shape :math:`(\\\\text{in\\\\_channels} , \\\\frac{\\\\text{out\\\\_channels}}{\\\\text{groups}} , kT , kH , kW)`\\n    bias: optional bias of shape :math:`(\\\\text{out\\\\_channels})`. Default: None\\n    stride: the stride of the convolving kernel. Can be a single number or a\\n      tuple ``(sT, sH, sW)``. Default: 1\\n    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both\\n      sides of each dimension in the input. Can be a single number or a tuple\\n      ``(padT, padH, padW)``. Default: 0\\n    output_padding: additional size added to one side of each dimension in the\\n      output shape. Can be a single number or a tuple\\n      ``(out_padT, out_padH, out_padW)``. Default: 0\\n    groups: split input into groups, :math:`\\\\text{in\\\\_channels}` should be divisible by the\\n      number of groups. Default: 1\\n    dilation: the spacing between kernel elements. Can be a single number or\\n      a tuple `(dT, dH, dW)`. Default: 1\\n\\nExamples::\\n\\n    >>> inputs = torch.randn(20, 16, 50, 10, 20)\\n    >>> weights = torch.randn(16, 33, 3, 3, 3)\\n    >>> F.conv_transpose3d(inputs, weights)\\n\\\"\\\"\\\",\\n)  # noqa: E501\\n\\nconv_tbc = _add_docstr(\\n    torch.conv_tbc,\\n    r\\\"\\\"\\\"\\nApplies a 1-dimensional sequence convolution over an input sequence.\\nInput and output dimensions are (Time, Batch, Channels) - hence TBC.\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{sequence length} \\\\times batch \\\\times \\\\text{in\\\\_channels})`\\n    weight: filter of shape (:math:`\\\\text{kernel width} \\\\times \\\\text{in\\\\_channels} \\\\times \\\\text{out\\\\_channels}`)\\n    bias: bias of shape (:math:`\\\\text{out\\\\_channels}`)\\n    pad: number of timesteps to pad. Default: 0\\n\\\"\\\"\\\",\\n)\\n\\n\\n# Pooling\\navg_pool1d = _add_docstr(\\n    torch.avg_pool1d,\\n    r\\\"\\\"\\\"\\navg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) -> Tensor\\n\\nApplies a 1D average pooling over an input signal composed of several\\ninput planes.\\n\\nSee :class:`~torch.nn.AvgPool1d` for details and output shape.\\n\\nArgs:\\n    input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iW)`\\n    kernel_size: the size of the window. Can be a single number or a\\n      tuple `(kW,)`\\n    stride: the stride of the window. Can be a single number or a tuple\\n      `(sW,)`. Default: :attr:`kernel_size`\\n    padding: implicit zero paddings on both sides of the input. Can be a\\n      single number or a tuple `(padW,)`. Default: 0\\n    ceil_mode: when True, will use `ceil` instead of `floor` to compute the\\n        output shape. Default: ``False``\\n    count_include_pad: when True, will include the zero-padding in the\\n        averaging calculation. Default: ``True``\\n\\nExamples::\\n\\n    >>> # pool of square window of size=3, stride=2\\n    >>> input = torch.tensor([[[1, 2, 3, 4, 5, 6, 7]]], dtype=torch.float32)\\n    >>> F.avg_pool1d(input, kernel_size=3, stride=2)\\n    tensor([[[ 2.,  4.,  6.]]])\\n\\n\\\"\\\"\\\",\\n)\\n\\n\\navg_pool2d = _add_docstr(\\n    torch._C._nn.avg_pool2d,\\n    r\\\"\\\"\\\"\\navg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -> Tensor\\n\\nApplies 2D average-pooling operation in :math:`kH \\\\times kW` regions by step size\\n:math:`sH \\\\times sW` steps. The number of output features is equal to the number of\\ninput planes.\\n\\nSee :class:`~torch.nn.AvgPool2d` for details and output shape.\\n\\nArgs:\\n    input: input tensor :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iH , iW)`\\n    kernel_size: size of the pooling region. Can be a single number or a\\n      tuple `(kH, kW)`\\n    stride: stride of the pooling operation. Can be a single number or a\\n      tuple `(sH, sW)`. Default: :attr:`kernel_size`\\n    padding: implicit zero paddings on both sides of the input. Can be a\\n      single number or a tuple `(padH, padW)`. Default: 0\\n    ceil_mode: when True, will use `ceil` instead of `floor` in the formula\\n        to compute the output shape. Default: ``False``\\n    count_include_pad: when True, will include the zero-padding in the\\n        averaging calculation. Default: ``True``\\n    divisor_override: if specified, it will be used as divisor, otherwise\\n         size of the pooling region will be used. Default: None\\n\\\"\\\"\\\",\\n)\\n\\navg_pool3d = _add_docstr(\\n    torch._C._nn.avg_pool3d,\\n    r\\\"\\\"\\\"\\navg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -> Tensor\\n\\nApplies 3D average-pooling operation in :math:`kT \\\\times kH \\\\times kW` regions by step\\nsize :math:`sT \\\\times sH \\\\times sW` steps. The number of output features is equal to\\n:math:`\\\\lfloor\\\\frac{\\\\text{input planes}}{sT}\\\\rfloor`.\\n\\nSee :class:`~torch.nn.AvgPool3d` for details and output shape.\\n\\nArgs:\\n    input: input tensor :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iT \\\\times iH , iW)`\\n    kernel_size: size of the pooling region. Can be a single number or a\\n      tuple `(kT, kH, kW)`\\n    stride: stride of the pooling operation. Can be a single number or a\\n      tuple `(sT, sH, sW)`. Default: :attr:`kernel_size`\\n    padding: implicit zero paddings on both sides of the input. Can be a\\n      single number or a tuple `(padT, padH, padW)`, Default: 0\\n    ceil_mode: when True, will use `ceil` instead of `floor` in the formula\\n        to compute the output shape\\n    count_include_pad: when True, will include the zero-padding in the\\n        averaging calculation\\n    divisor_override: if specified, it will be used as divisor, otherwise\\n        size of the pooling region will be used. Default: None\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef fractional_max_pool2d_with_indices(\\n    input: Tensor,\\n    kernel_size: BroadcastingList2[int],\\n    output_size: Optional[BroadcastingList2[int]] = None,\\n    output_ratio: Optional[BroadcastingList2[float]] = None,\\n    return_indices: bool = False,\\n    _random_samples: Optional[Tensor] = None,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    fractional_max_pool2d(input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)\\n\\n    Applies 2D fractional max pooling over an input signal composed of several input planes.\\n\\n    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham\\n\\n    The max-pooling operation is applied in :math:`kH \\\\times kW` regions by a stochastic\\n    step size determined by the target output size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        kernel_size: the size of the window to take a max over.\\n                     Can be a single number :math:`k` (for a square kernel of :math:`k \\\\times k`)\\n                     or a tuple `(kH, kW)`\\n        output_size: the target output size of the image of the form :math:`oH \\\\times oW`.\\n                     Can be a tuple `(oH, oW)` or a single number :math:`oH` for a square image :math:`oH \\\\times oH`\\n        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.\\n                      This has to be a number or tuple in the range (0, 1)\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to :func:`~torch.nn.functional.max_unpool2d`.\\n\\n    Examples::\\n        >>> input = torch.randn(20, 16, 50, 32)\\n        >>> # pool of square window of size=3, and target output size 13x12\\n        >>> F.fractional_max_pool2d(input, 3, output_size=(13, 12))\\n        >>> # pool of square window and target output size being half of input image size\\n        >>> F.fractional_max_pool2d(input, 3, output_ratio=(0.5, 0.5))\\n\\n    .. _Fractional MaxPooling:\\n        http://arxiv.org/abs/1412.6071\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, _random_samples):\\n        return handle_torch_function(\\n            fractional_max_pool2d_with_indices,\\n            (input, _random_samples),\\n            input,\\n            kernel_size,\\n            output_size=output_size,\\n            output_ratio=output_ratio,\\n            return_indices=return_indices,\\n            _random_samples=_random_samples,\\n        )\\n    if output_size is None and output_ratio is None:\\n        raise ValueError(\\n            \\\"fractional_max_pool2d requires specifying either an output_size or an output_ratio\\\"\\n        )\\n    if output_size is None:\\n        assert output_ratio is not None\\n        if len(output_ratio) > 2:\\n            raise ValueError(\\n                \\\"fractional_max_pool2d requires output_ratio to either be a single Int or tuple of Ints.\\\"\\n            )\\n        _output_ratio = _pair(output_ratio)\\n        output_size = [\\n            int(input.size(-2) * _output_ratio[0]),\\n            int(input.size(-1) * _output_ratio[1]),\\n        ]\\n\\n    if _random_samples is None:\\n        n_batch = 1 if input.dim() == 3 else input.size(0)\\n        _random_samples = torch.rand(\\n            n_batch, input.size(-3), 2, dtype=input.dtype, device=input.device\\n        )\\n    return torch._C._nn.fractional_max_pool2d(\\n        input, kernel_size, output_size, _random_samples\\n    )\\n\\n\\ndef _fractional_max_pool2d(\\n    input: Tensor,\\n    kernel_size: BroadcastingList2[int],\\n    output_size: Optional[BroadcastingList2[int]] = None,\\n    output_ratio: Optional[BroadcastingList2[float]] = None,\\n    return_indices: bool = False,\\n    _random_samples: Optional[Tensor] = None,\\n) -> Tensor:\\n    if has_torch_function_variadic(input, _random_samples):\\n        return handle_torch_function(\\n            fractional_max_pool2d,\\n            (input, _random_samples),\\n            input,\\n            kernel_size,\\n            output_size=output_size,\\n            output_ratio=output_ratio,\\n            return_indices=return_indices,\\n            _random_samples=_random_samples,\\n        )\\n    return fractional_max_pool2d_with_indices(\\n        input, kernel_size, output_size, output_ratio, return_indices, _random_samples\\n    )[0]\\n\\n\\nfractional_max_pool2d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=4,\\n    default=False,\\n    if_true=fractional_max_pool2d_with_indices,\\n    if_false=_fractional_max_pool2d,\\n    module_name=__name__,\\n    func_name=\\\"fractional_max_pool2d\\\",\\n)\\n\\n\\ndef fractional_max_pool3d_with_indices(\\n    input: Tensor,\\n    kernel_size: BroadcastingList3[int],\\n    output_size: Optional[BroadcastingList3[int]] = None,\\n    output_ratio: Optional[BroadcastingList3[float]] = None,\\n    return_indices: bool = False,\\n    _random_samples: Optional[Tensor] = None,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    fractional_max_pool3d(input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)\\n\\n    Applies 3D fractional max pooling over an input signal composed of several input planes.\\n\\n    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham\\n\\n    The max-pooling operation is applied in :math:`kT \\\\times kH \\\\times kW` regions by a stochastic\\n    step size determined by the target output size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        kernel_size: the size of the window to take a max over.\\n                     Can be a single number :math:`k` (for a square kernel of :math:`k \\\\times k \\\\times k`)\\n                     or a tuple `(kT, kH, kW)`\\n        output_size: the target output size of the form :math:`oT \\\\times oH \\\\times oW`.\\n                     Can be a tuple `(oT, oH, oW)` or a single number :math:`oH` for a cubic output\\n                     :math:`oH \\\\times oH \\\\times oH`\\n        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.\\n                      This has to be a number or tuple in the range (0, 1)\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to :func:`~torch.nn.functional.max_unpool3d`.\\n\\n    Shape:\\n        - Input: :math:`(N, C, T_{in}, H_{in}, W_{in})` or :math:`(C, T_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, T_{out}, H_{out}, W_{out})` or :math:`(C, T_{out}, H_{out}, W_{out})`, where\\n          :math:`(T_{out}, H_{out}, W_{out})=\\\\text{output\\\\_size}` or\\n          :math:`(T_{out}, H_{out}, W_{out})=\\\\text{output\\\\_ratio} \\\\times (T_{in}, H_{in}, W_{in})`\\n\\n    Examples::\\n        >>> input = torch.randn(20, 16, 50, 32, 16)\\n        >>> # pool of cubic window of size=3, and target output size 13x12x11\\n        >>> F.fractional_max_pool3d(input, 3, output_size=(13, 12, 11))\\n        >>> # pool of cubic window and target output size being half of input size\\n        >>> F.fractional_max_pool3d(input, 3, output_ratio=(0.5, 0.5, 0.5))\\n\\n    .. _Fractional MaxPooling:\\n        http://arxiv.org/abs/1412.6071\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, _random_samples):\\n        return handle_torch_function(\\n            fractional_max_pool3d_with_indices,\\n            (input, _random_samples),\\n            input,\\n            kernel_size,\\n            output_size=output_size,\\n            output_ratio=output_ratio,\\n            return_indices=return_indices,\\n            _random_samples=_random_samples,\\n        )\\n    if output_size is None and output_ratio is None:\\n        raise ValueError(\\n            \\\"fractional_max_pool3d requires specifying either an output_size or an output_ratio\\\"\\n        )\\n    if output_size is None:\\n        assert output_ratio is not None\\n        _output_ratio = _triple(output_ratio)\\n        output_size = [\\n            int(input.size(-3) * _output_ratio[0]),\\n            int(input.size(-2) * _output_ratio[1]),\\n            int(input.size(-1) * _output_ratio[2]),\\n        ]\\n\\n    if _random_samples is None:\\n        n_batch = 1 if input.dim() == 4 else input.size(0)\\n        _random_samples = torch.rand(\\n            n_batch, input.size(-4), 3, dtype=input.dtype, device=input.device\\n        )\\n    return torch._C._nn.fractional_max_pool3d(\\n        input, kernel_size, output_size, _random_samples\\n    )\\n\\n\\ndef _fractional_max_pool3d(\\n    input: Tensor,\\n    kernel_size: BroadcastingList3[int],\\n    output_size: Optional[BroadcastingList3[int]] = None,\\n    output_ratio: Optional[BroadcastingList3[float]] = None,\\n    return_indices: bool = False,\\n    _random_samples: Optional[Tensor] = None,\\n) -> Tensor:\\n    if has_torch_function_variadic(input, _random_samples):\\n        return handle_torch_function(\\n            fractional_max_pool3d,\\n            (input, _random_samples),\\n            input,\\n            kernel_size,\\n            output_size=output_size,\\n            output_ratio=output_ratio,\\n            return_indices=return_indices,\\n            _random_samples=_random_samples,\\n        )\\n    return fractional_max_pool3d_with_indices(\\n        input, kernel_size, output_size, output_ratio, return_indices, _random_samples\\n    )[0]\\n\\n\\nfractional_max_pool3d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=4,\\n    default=False,\\n    if_true=fractional_max_pool3d_with_indices,\\n    if_false=_fractional_max_pool3d,\\n    module_name=__name__,\\n    func_name=\\\"fractional_max_pool3d\\\",\\n)\\n\\n\\ndef max_pool1d_with_indices(\\n    input: Tensor,\\n    kernel_size: BroadcastingList1[int],\\n    stride: Optional[BroadcastingList1[int]] = None,\\n    padding: BroadcastingList1[int] = 0,\\n    dilation: BroadcastingList1[int] = 1,\\n    ceil_mode: bool = False,\\n    return_indices: bool = False,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)\\n\\n    Applies a 1D max pooling over an input signal composed of several input\\n    planes.\\n\\n    .. note::\\n        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from\\n        what seen in :class:`~torch.nn.MaxPool1d`, and will change in a future release.\\n\\n    See :class:`~torch.nn.MaxPool1d` for details.\\n\\n    Args:\\n        input: input tensor of shape :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iW)`, minibatch dim optional.\\n        kernel_size: the size of the window. Can be a single number or a\\n            tuple `(kW,)`\\n        stride: the stride of the window. Can be a single number or a tuple\\n            `(sW,)`. Default: :attr:`kernel_size`\\n        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.\\n        dilation: The stride between elements within a sliding window, must be > 0.\\n        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This\\n                   ensures that every element in the input tensor is covered by a sliding window.\\n        return_indices: If ``True``, will return the argmax along with the max values.\\n                        Useful for :class:`torch.nn.functional.max_unpool1d` later\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_pool1d_with_indices,\\n            (input,),\\n            input,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            dilation=dilation,\\n            ceil_mode=ceil_mode,\\n            return_indices=return_indices,\\n        )\\n    if stride is None:\\n        stride = torch.jit.annotate(List[int], [])\\n    return torch.max_pool1d_with_indices(\\n        input, kernel_size, stride, padding, dilation, ceil_mode\\n    )\\n\\n\\ndef _max_pool1d(\\n    input: Tensor,\\n    kernel_size: BroadcastingList1[int],\\n    stride: Optional[BroadcastingList1[int]] = None,\\n    padding: BroadcastingList1[int] = 0,\\n    dilation: BroadcastingList1[int] = 1,\\n    ceil_mode: bool = False,\\n    return_indices: bool = False,\\n) -> Tensor:\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_pool1d,\\n            (input,),\\n            input,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            dilation=dilation,\\n            ceil_mode=ceil_mode,\\n            return_indices=return_indices,\\n        )\\n    if stride is None:\\n        stride = torch.jit.annotate(List[int], [])\\n    return torch.max_pool1d(input, kernel_size, stride, padding, dilation, ceil_mode)\\n\\n\\nmax_pool1d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=6,\\n    default=False,\\n    if_true=max_pool1d_with_indices,\\n    if_false=_max_pool1d,\\n    module_name=__name__,\\n    func_name=\\\"max_pool1d\\\",\\n)\\n\\n\\ndef max_pool2d_with_indices(\\n    input: Tensor,\\n    kernel_size: BroadcastingList2[int],\\n    stride: Optional[BroadcastingList2[int]] = None,\\n    padding: BroadcastingList2[int] = 0,\\n    dilation: BroadcastingList2[int] = 1,\\n    ceil_mode: bool = False,\\n    return_indices: bool = False,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)\\n\\n    Applies a 2D max pooling over an input signal composed of several input\\n    planes.\\n\\n    .. note::\\n        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from\\n        what seen in :class:`~torch.nn.MaxPool2d`, and will change in a future release.\\n\\n    See :class:`~torch.nn.MaxPool2d` for details.\\n\\n    Args:\\n        input: input tensor :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iH , iW)`, minibatch dim optional.\\n        kernel_size: size of the pooling region. Can be a single number or a\\n            tuple `(kH, kW)`\\n        stride: stride of the pooling operation. Can be a single number or a\\n            tuple `(sH, sW)`. Default: :attr:`kernel_size`\\n        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.\\n        dilation: The stride between elements within a sliding window, must be > 0.\\n        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This\\n                   ensures that every element in the input tensor is covered by a sliding window.\\n        return_indices: If ``True``, will return the argmax along with the max values.\\n                        Useful for :class:`torch.nn.functional.max_unpool2d` later\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_pool2d_with_indices,\\n            (input,),\\n            input,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            dilation=dilation,\\n            ceil_mode=ceil_mode,\\n            return_indices=return_indices,\\n        )\\n    if stride is None:\\n        stride = torch.jit.annotate(List[int], [])\\n    return torch._C._nn.max_pool2d_with_indices(\\n        input, kernel_size, stride, padding, dilation, ceil_mode\\n    )\\n\\n\\ndef _max_pool2d(\\n    input: Tensor,\\n    kernel_size: BroadcastingList2[int],\\n    stride: Optional[BroadcastingList2[int]] = None,\\n    padding: BroadcastingList2[int] = 0,\\n    dilation: BroadcastingList2[int] = 1,\\n    ceil_mode: bool = False,\\n    return_indices: bool = False,\\n) -> Tensor:\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_pool2d,\\n            (input,),\\n            input,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            dilation=dilation,\\n            ceil_mode=ceil_mode,\\n            return_indices=return_indices,\\n        )\\n    if stride is None:\\n        stride = torch.jit.annotate(List[int], [])\\n    return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)\\n\\n\\nmax_pool2d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=6,\\n    default=False,\\n    if_true=max_pool2d_with_indices,\\n    if_false=_max_pool2d,\\n    module_name=__name__,\\n    func_name=\\\"max_pool2d\\\",\\n)\\n\\n\\ndef max_pool3d_with_indices(\\n    input: Tensor,\\n    kernel_size: BroadcastingList3[int],\\n    stride: Optional[BroadcastingList3[int]] = None,\\n    padding: BroadcastingList3[int] = 0,\\n    dilation: BroadcastingList3[int] = 1,\\n    ceil_mode: bool = False,\\n    return_indices: bool = False,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)\\n\\n    Applies a 3D max pooling over an input signal composed of several input\\n    planes.\\n\\n    .. note::\\n        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from\\n        what seen in :class:`~torch.nn.MaxPool3d`, and will change in a future release.\\n\\n    See :class:`~torch.nn.MaxPool3d` for details.\\n\\n    Args:\\n        input: input tensor :math:`(\\\\text{minibatch} , \\\\text{in\\\\_channels} , iD, iH , iW)`, minibatch dim optional.\\n        kernel_size: size of the pooling region. Can be a single number or a\\n                     tuple `(kT, kH, kW)`\\n        stride: stride of the pooling operation. Can be a single number or a\\n                tuple `(sT, sH, sW)`. Default: :attr:`kernel_size`\\n        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.\\n        dilation: The stride between elements within a sliding window, must be > 0.\\n        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This\\n                   ensures that every element in the input tensor is covered by a sliding window.\\n        return_indices: If ``True``, will return the argmax along with the max values.\\n                        Useful for :class:`torch.nn.functional.max_unpool3d` later\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_pool3d_with_indices,\\n            (input,),\\n            input,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            dilation=dilation,\\n            ceil_mode=ceil_mode,\\n            return_indices=return_indices,\\n        )\\n    if stride is None:\\n        stride = torch.jit.annotate(List[int], [])\\n    return torch._C._nn.max_pool3d_with_indices(\\n        input, kernel_size, stride, padding, dilation, ceil_mode\\n    )\\n\\n\\ndef _max_pool3d(\\n    input: Tensor,\\n    kernel_size: BroadcastingList3[int],\\n    stride: Optional[BroadcastingList3[int]] = None,\\n    padding: BroadcastingList3[int] = 0,\\n    dilation: BroadcastingList3[int] = 1,\\n    ceil_mode: bool = False,\\n    return_indices: bool = False,\\n) -> Tensor:\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_pool3d,\\n            (input,),\\n            input,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            dilation=dilation,\\n            ceil_mode=ceil_mode,\\n            return_indices=return_indices,\\n        )\\n    if stride is None:\\n        stride = torch.jit.annotate(List[int], [])\\n    return torch.max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode)\\n\\n\\nmax_pool3d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=6,\\n    default=False,\\n    if_true=max_pool3d_with_indices,\\n    if_false=_max_pool3d,\\n    module_name=__name__,\\n    func_name=\\\"max_pool3d\\\",\\n)\\n\\n\\ndef _unpool_output_size(\\n    input: Tensor,\\n    kernel_size: List[int],\\n    stride: List[int],\\n    padding: List[int],\\n    output_size: Optional[List[int]],\\n) -> List[int]:\\n    input_size = input.size()\\n    default_size = torch.jit.annotate(List[int], [])\\n    for d in range(len(kernel_size)):\\n        default_size.append(\\n            (input_size[-len(kernel_size) + d] - 1) * stride[d]\\n            + kernel_size[d]\\n            - 2 * padding[d]\\n        )\\n    if output_size is None:\\n        ret = default_size\\n    else:\\n        if len(output_size) == len(kernel_size) + 2:\\n            output_size = output_size[2:]\\n        if len(output_size) != len(kernel_size):\\n            raise ValueError(\\n                \\\"output_size should be a sequence containing \\\"\\n                f\\\"{len(kernel_size)} or {len(kernel_size) + 2} elements, but it has a length of '{len(output_size)}'\\\"\\n            )\\n        for d in range(len(kernel_size)):\\n            min_size = default_size[d] - stride[d]\\n            max_size = default_size[d] + stride[d]\\n            if not (min_size < output_size[d] < max_size):\\n                raise ValueError(\\n                    f'invalid output_size \\\"{output_size}\\\" (dim {d} must be between {min_size} and {max_size})'\\n                )\\n\\n        ret = output_size\\n    return ret\\n\\n\\ndef max_unpool1d(\\n    input: Tensor,\\n    indices: Tensor,\\n    kernel_size: BroadcastingList1[int],\\n    stride: Optional[BroadcastingList1[int]] = None,\\n    padding: BroadcastingList1[int] = 0,\\n    output_size: Optional[BroadcastingList1[int]] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute a partial inverse of :class:`MaxPool1d`.\\n\\n    See :class:`~torch.nn.MaxUnpool1d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_unpool1d,\\n            (input,),\\n            input,\\n            indices,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            output_size=output_size,\\n        )\\n    kernel_size = _single(kernel_size)\\n    if stride is not None:\\n        _stride = _single(stride)\\n    else:\\n        _stride = kernel_size\\n    padding = _single(padding)\\n    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)\\n    if isinstance(output_size, list):\\n        output_size = output_size + [1]\\n    else:\\n        output_size = output_size + (1,)\\n    return torch._C._nn.max_unpool2d(\\n        input.unsqueeze(-1), indices.unsqueeze(-1), output_size\\n    ).squeeze(-1)\\n\\n\\ndef max_unpool2d(\\n    input: Tensor,\\n    indices: Tensor,\\n    kernel_size: BroadcastingList2[int],\\n    stride: Optional[BroadcastingList2[int]] = None,\\n    padding: BroadcastingList2[int] = 0,\\n    output_size: Optional[BroadcastingList2[int]] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute a partial inverse of :class:`MaxPool2d`.\\n\\n    See :class:`~torch.nn.MaxUnpool2d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_unpool2d,\\n            (input,),\\n            input,\\n            indices,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            output_size=output_size,\\n        )\\n    kernel_size = _pair(kernel_size)\\n    if stride is not None:\\n        _stride = _pair(stride)\\n    else:\\n        _stride = kernel_size\\n    padding = _pair(padding)\\n    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)\\n    return torch._C._nn.max_unpool2d(input, indices, output_size)\\n\\n\\ndef max_unpool3d(\\n    input: Tensor,\\n    indices: Tensor,\\n    kernel_size: BroadcastingList3[int],\\n    stride: Optional[BroadcastingList3[int]] = None,\\n    padding: BroadcastingList3[int] = 0,\\n    output_size: Optional[BroadcastingList3[int]] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute a partial inverse of :class:`MaxPool3d`.\\n\\n    See :class:`~torch.nn.MaxUnpool3d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            max_unpool3d,\\n            (input,),\\n            input,\\n            indices,\\n            kernel_size,\\n            stride=stride,\\n            padding=padding,\\n            output_size=output_size,\\n        )\\n    kernel_size = _triple(kernel_size)\\n    if stride is not None:\\n        _stride = _triple(stride)\\n    else:\\n        _stride = kernel_size\\n    padding = _triple(padding)\\n    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)\\n    return torch._C._nn.max_unpool3d(input, indices, output_size, _stride, padding)\\n\\n\\ndef lp_pool3d(\\n    input: Tensor,\\n    norm_type: Union[int, float],\\n    kernel_size: BroadcastingList3[int],\\n    stride: Optional[BroadcastingList3[int]] = None,\\n    ceil_mode: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"\\n    Apply a 3D power-average pooling over an input signal composed of several input planes.\\n\\n    If the sum of all inputs to the power of `p` is\\n    zero, the gradient is set to zero as well.\\n\\n    See :class:`~torch.nn.LPPool3d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            lp_pool3d,\\n            (input,),\\n            input,\\n            norm_type,\\n            kernel_size,\\n            stride=stride,\\n            ceil_mode=ceil_mode,\\n        )\\n    kd, kw, kh = _triple(kernel_size)\\n    if stride is not None:\\n        out = avg_pool3d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)\\n    else:\\n        out = avg_pool3d(\\n            input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode\\n        )\\n\\n    return (\\n        (torch.sign(out) * relu(torch.abs(out))).mul(kd * kw * kh).pow(1.0 / norm_type)\\n    )\\n\\n\\ndef lp_pool2d(\\n    input: Tensor,\\n    norm_type: Union[int, float],\\n    kernel_size: BroadcastingList2[int],\\n    stride: Optional[BroadcastingList2[int]] = None,\\n    ceil_mode: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"\\n    Apply a 2D power-average pooling over an input signal composed of several input planes.\\n\\n    If the sum of all inputs to the power of `p` is\\n    zero, the gradient is set to zero as well.\\n\\n    See :class:`~torch.nn.LPPool2d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            lp_pool2d,\\n            (input,),\\n            input,\\n            norm_type,\\n            kernel_size,\\n            stride=stride,\\n            ceil_mode=ceil_mode,\\n        )\\n    kw, kh = _pair(kernel_size)\\n    if stride is not None:\\n        out = avg_pool2d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)\\n    else:\\n        out = avg_pool2d(\\n            input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode\\n        )\\n\\n    return (torch.sign(out) * relu(torch.abs(out))).mul(kw * kh).pow(1.0 / norm_type)\\n\\n\\ndef lp_pool1d(\\n    input: Tensor,\\n    norm_type: Union[int, float],\\n    kernel_size: int,\\n    stride: Optional[BroadcastingList1[int]] = None,\\n    ceil_mode: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply a 1D power-average pooling over an input signal composed of several input planes.\\n\\n    If the sum of all inputs to the power of `p` is\\n    zero, the gradient is set to zero as well.\\n\\n    See :class:`~torch.nn.LPPool1d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            lp_pool1d,\\n            (input,),\\n            input,\\n            norm_type,\\n            kernel_size,\\n            stride=stride,\\n            ceil_mode=ceil_mode,\\n        )\\n    if stride is not None:\\n        out = avg_pool1d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)\\n    else:\\n        out = avg_pool1d(\\n            input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode\\n        )\\n\\n    return (\\n        (torch.sign(out) * relu(torch.abs(out))).mul(kernel_size).pow(1.0 / norm_type)\\n    )\\n\\n\\ndef adaptive_max_pool1d_with_indices(\\n    input: Tensor,\\n    output_size: BroadcastingList1[int],\\n    return_indices: bool = False,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    adaptive_max_pool1d(input, output_size, return_indices=False)\\n\\n    Applies a 1D adaptive max pooling over an input signal composed of\\n    several input planes.\\n\\n    See :class:`~torch.nn.AdaptiveMaxPool1d` for details and output shape.\\n\\n    Args:\\n        output_size: the target output size (single integer)\\n        return_indices: whether to return pooling indices. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            adaptive_max_pool1d_with_indices,\\n            (input,),\\n            input,\\n            output_size,\\n            return_indices=return_indices,\\n        )\\n    return torch.adaptive_max_pool1d(input, output_size)\\n\\n\\ndef _adaptive_max_pool1d(\\n    input: Tensor,\\n    output_size: BroadcastingList1[int],\\n    return_indices: bool = False,\\n) -> Tensor:\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            adaptive_max_pool1d,\\n            (input,),\\n            input,\\n            output_size,\\n            return_indices=return_indices,\\n        )\\n    return adaptive_max_pool1d_with_indices(input, output_size)[0]\\n\\n\\nadaptive_max_pool1d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=2,\\n    default=False,\\n    if_true=adaptive_max_pool1d_with_indices,\\n    if_false=_adaptive_max_pool1d,\\n    module_name=__name__,\\n    func_name=\\\"adaptive_max_pool1d\\\",\\n)\\n\\n\\ndef adaptive_max_pool2d_with_indices(\\n    input: Tensor,\\n    output_size: BroadcastingList2[int],\\n    return_indices: bool = False,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"adaptive_max_pool2d(input, output_size, return_indices=False)\\n\\n    Applies a 2D adaptive max pooling over an input signal composed of\\n    several input planes.\\n\\n    See :class:`~torch.nn.AdaptiveMaxPool2d` for details and output shape.\\n\\n    Args:\\n        output_size: the target output size (single integer or\\n            double-integer tuple)\\n        return_indices: whether to return pooling indices. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            adaptive_max_pool2d_with_indices,\\n            (input,),\\n            input,\\n            output_size,\\n            return_indices=return_indices,\\n        )\\n    output_size = _list_with_default(output_size, input.size())\\n    return torch._C._nn.adaptive_max_pool2d(input, output_size)\\n\\n\\ndef _adaptive_max_pool2d(\\n    input: Tensor,\\n    output_size: BroadcastingList2[int],\\n    return_indices: bool = False,\\n) -> Tensor:\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            adaptive_max_pool2d,\\n            (input,),\\n            input,\\n            output_size,\\n            return_indices=return_indices,\\n        )\\n    return adaptive_max_pool2d_with_indices(input, output_size)[0]\\n\\n\\nadaptive_max_pool2d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=2,\\n    default=False,\\n    if_true=adaptive_max_pool2d_with_indices,\\n    if_false=_adaptive_max_pool2d,\\n    module_name=__name__,\\n    func_name=\\\"adaptive_max_pool2d\\\",\\n)\\n\\n\\ndef adaptive_max_pool3d_with_indices(\\n    input: Tensor,\\n    output_size: BroadcastingList3[int],\\n    return_indices: bool = False,\\n) -> Tuple[Tensor, Tensor]:  # noqa: D400\\n    r\\\"\\\"\\\"\\n    adaptive_max_pool3d(input, output_size, return_indices=False)\\n\\n    Applies a 3D adaptive max pooling over an input signal composed of\\n    several input planes.\\n\\n    See :class:`~torch.nn.AdaptiveMaxPool3d` for details and output shape.\\n\\n    Args:\\n        output_size: the target output size (single integer or\\n            triple-integer tuple)\\n        return_indices: whether to return pooling indices. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            adaptive_max_pool3d_with_indices,\\n            (input,),\\n            input,\\n            output_size,\\n            return_indices=return_indices,\\n        )\\n    output_size = _list_with_default(output_size, input.size())\\n    return torch._C._nn.adaptive_max_pool3d(input, output_size)\\n\\n\\ndef _adaptive_max_pool3d(\\n    input: Tensor,\\n    output_size: BroadcastingList3[int],\\n    return_indices: bool = False,\\n) -> Tensor:\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            adaptive_max_pool3d,\\n            (input,),\\n            input,\\n            output_size,\\n            return_indices=return_indices,\\n        )\\n    return adaptive_max_pool3d_with_indices(input, output_size)[0]\\n\\n\\nadaptive_max_pool3d = boolean_dispatch(\\n    arg_name=\\\"return_indices\\\",\\n    arg_index=2,\\n    default=False,\\n    if_true=adaptive_max_pool3d_with_indices,\\n    if_false=_adaptive_max_pool3d,\\n    module_name=__name__,\\n    func_name=\\\"adaptive_max_pool3d\\\",\\n)\\n\\n\\nadaptive_avg_pool1d = _add_docstr(\\n    torch.adaptive_avg_pool1d,\\n    r\\\"\\\"\\\"\\nadaptive_avg_pool1d(input, output_size) -> Tensor\\n\\nApplies a 1D adaptive average pooling over an input signal composed of\\nseveral input planes.\\n\\nSee :class:`~torch.nn.AdaptiveAvgPool1d` for details and output shape.\\n\\nArgs:\\n    output_size: the target output size (single integer)\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef adaptive_avg_pool2d(input: Tensor, output_size: BroadcastingList2[int]) -> Tensor:\\n    r\\\"\\\"\\\"Apply a 2D adaptive average pooling over an input signal composed of several input planes.\\n\\n    See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.\\n\\n    Args:\\n        output_size: the target output size (single integer or\\n            double-integer tuple)\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(adaptive_avg_pool2d, (input,), input, output_size)\\n    _output_size = _list_with_default(output_size, input.size())\\n    return torch._C._nn.adaptive_avg_pool2d(input, _output_size)\\n\\n\\ndef adaptive_avg_pool3d(input: Tensor, output_size: BroadcastingList3[int]) -> Tensor:\\n    r\\\"\\\"\\\"Apply a 3D adaptive average pooling over an input signal composed of several input planes.\\n\\n    See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.\\n\\n    Args:\\n        output_size: the target output size (single integer or\\n            triple-integer tuple)\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(adaptive_avg_pool3d, (input,), input, output_size)\\n    _output_size = _list_with_default(output_size, input.size())\\n    return torch._C._nn.adaptive_avg_pool3d(input, _output_size)\\n\\n\\n# Activation functions\\ndef dropout(\\n    input: Tensor,\\n    p: float = 0.5,\\n    training: bool = True,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"During training, randomly zeroes some elements of the input tensor with probability :attr:`p`.\\n\\n    Uses samples from a Bernoulli distribution.\\n\\n    See :class:`~torch.nn.Dropout` for details.\\n\\n    Args:\\n        p: probability of an element to be zeroed. Default: 0.5\\n        training: apply dropout if is ``True``. Default: ``True``\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            dropout, (input,), input, p=p, training=training, inplace=inplace\\n        )\\n    if p < 0.0 or p > 1.0:\\n        raise ValueError(f\\\"dropout probability has to be between 0 and 1, but got {p}\\\")\\n    return (\\n        _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training)\\n    )\\n\\n\\ndef alpha_dropout(\\n    input: Tensor,\\n    p: float = 0.5,\\n    training: bool = False,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply alpha dropout to the input.\\n\\n    See :class:`~torch.nn.AlphaDropout` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            alpha_dropout, (input,), input, p=p, training=training, inplace=inplace\\n        )\\n    if p < 0.0 or p > 1.0:\\n        raise ValueError(f\\\"dropout probability has to be between 0 and 1, but got {p}\\\")\\n    return (\\n        _VF.alpha_dropout_(input, p, training)\\n        if inplace\\n        else _VF.alpha_dropout(input, p, training)\\n    )\\n\\n\\ndef dropout1d(\\n    input: Tensor,\\n    p: float = 0.5,\\n    training: bool = True,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Randomly zero out entire channels (a channel is a 1D feature map).\\n\\n    For example, the :math:`j`-th channel of the :math:`i`-th sample in the\\n    batched input is a 1D tensor :math:`\\\\text{input}[i, j]` of the input tensor.\\n    Each channel will be zeroed out independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n\\n    See :class:`~torch.nn.Dropout1d` for details.\\n\\n    Args:\\n        p: probability of a channel to be zeroed. Default: 0.5\\n        training: apply dropout if is ``True``. Default: ``True``\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            dropout1d, (input,), input, p=p, training=training, inplace=inplace\\n        )\\n    if p < 0.0 or p > 1.0:\\n        raise ValueError(f\\\"dropout probability has to be between 0 and 1, but got {p}\\\")\\n    inp_dim = input.dim()\\n    if inp_dim not in (2, 3):\\n        raise RuntimeError(\\n            f\\\"dropout1d: Expected 2D or 3D input, but received a {inp_dim}D input. \\\"\\n            \\\"Note that dropout1d exists to provide channel-wise dropout on inputs with 1 \\\"\\n            \\\"spatial dimension, a channel dimension, and an optional batch dimension \\\"\\n            \\\"(i.e. 2D or 3D inputs).\\\"\\n        )\\n\\n    is_batched = inp_dim == 3\\n    if not is_batched:\\n        input = input.unsqueeze_(0) if inplace else input.unsqueeze(0)\\n\\n    result = (\\n        _VF.feature_dropout_(input, p, training)\\n        if inplace\\n        else _VF.feature_dropout(input, p, training)\\n    )\\n\\n    if not is_batched:\\n        result = result.squeeze_(0) if inplace else result.squeeze(0)\\n\\n    return result\\n\\n\\ndef dropout2d(\\n    input: Tensor,\\n    p: float = 0.5,\\n    training: bool = True,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Randomly zero out entire channels (a channel is a 2D feature map).\\n\\n    For example, the :math:`j`-th channel of the :math:`i`-th sample in the\\n    batched input is a 2D tensor :math:`\\\\text{input}[i, j]` of the input tensor.\\n    Each channel will be zeroed out independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n\\n    See :class:`~torch.nn.Dropout2d` for details.\\n\\n    Args:\\n        p: probability of a channel to be zeroed. Default: 0.5\\n        training: apply dropout if is ``True``. Default: ``True``\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            dropout2d, (input,), input, p=p, training=training, inplace=inplace\\n        )\\n    if p < 0.0 or p > 1.0:\\n        raise ValueError(f\\\"dropout probability has to be between 0 and 1, but got {p}\\\")\\n    inp_dim = input.dim()\\n    if inp_dim not in (3, 4):\\n        warn_msg = (\\n            f\\\"dropout2d: Received a {inp_dim}-D input to dropout2d, which is deprecated \\\"\\n            \\\"and will result in an error in a future release. To retain the behavior \\\"\\n            \\\"and silence this warning, please use dropout instead. Note that dropout2d \\\"\\n            \\\"exists to provide channel-wise dropout on inputs with 2 spatial dimensions, \\\"\\n            \\\"a channel dimension, and an optional batch dimension (i.e. 3D or 4D inputs).\\\"\\n        )\\n        warnings.warn(warn_msg)\\n\\n    # TODO: Properly support no-batch-dim inputs. For now, these are NOT supported; passing\\n    # a 3D input will perform dropout1d behavior instead. This was done historically and the\\n    # behavior is maintained here for now.\\n    # See https://github.com/pytorch/pytorch/issues/77081\\n    if inp_dim == 3:\\n        warnings.warn(\\n            \\\"dropout2d: Received a 3D input to dropout2d and assuming that channel-wise \\\"\\n            \\\"1D dropout behavior is desired - input is interpreted as shape (N, C, L), where C \\\"\\n            \\\"is the channel dim. This behavior will change in a future release to interpret the \\\"\\n            \\\"input as one without a batch dimension, i.e. shape (C, H, W). To maintain the 1D \\\"\\n            \\\"channel-wise dropout behavior, please switch to using dropout1d instead.\\\"\\n        )\\n\\n    result = (\\n        _VF.feature_dropout_(input, p, training)\\n        if inplace\\n        else _VF.feature_dropout(input, p, training)\\n    )\\n\\n    return result\\n\\n\\ndef dropout3d(\\n    input: Tensor,\\n    p: float = 0.5,\\n    training: bool = True,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Randomly zero out entire channels (a channel is a 3D feature map).\\n\\n    For example, the :math:`j`-th channel of the :math:`i`-th sample in the\\n    batched input is a 3D tensor :math:`\\\\text{input}[i, j]` of the input tensor.\\n    Each channel will be zeroed out independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n\\n    See :class:`~torch.nn.Dropout3d` for details.\\n\\n    Args:\\n        p: probability of a channel to be zeroed. Default: 0.5\\n        training: apply dropout if is ``True``. Default: ``True``\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            dropout3d, (input,), input, p=p, training=training, inplace=inplace\\n        )\\n    if p < 0.0 or p > 1.0:\\n        raise ValueError(f\\\"dropout probability has to be between 0 and 1, but got {p}\\\")\\n    inp_dim = input.dim()\\n    if inp_dim not in (4, 5):\\n        warn_msg = (\\n            f\\\"dropout3d: Received a {inp_dim}-D input to dropout3d, which is deprecated \\\"\\n            \\\"and will result in an error in a future release. To retain the behavior \\\"\\n            \\\"and silence this warning, please use dropout instead. Note that dropout3d \\\"\\n            \\\"exists to provide channel-wise dropout on inputs with 3 spatial dimensions, \\\"\\n            \\\"a channel dimension, and an optional batch dimension (i.e. 4D or 5D inputs).\\\"\\n        )\\n        warnings.warn(warn_msg)\\n\\n    is_batched = inp_dim == 5\\n    if not is_batched:\\n        input = input.unsqueeze_(0) if inplace else input.unsqueeze(0)\\n\\n    result = (\\n        _VF.feature_dropout_(input, p, training)\\n        if inplace\\n        else _VF.feature_dropout(input, p, training)\\n    )\\n\\n    if not is_batched:\\n        result = result.squeeze_(0) if inplace else result.squeeze(0)\\n    return result\\n\\n\\ndef feature_alpha_dropout(\\n    input: Tensor,\\n    p: float = 0.5,\\n    training: bool = False,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Randomly masks out entire channels (a channel is a feature map).\\n\\n    For example, the :math:`j`-th channel of the :math:`i`-th sample in the batch input\\n    is a tensor :math:`\\\\text{input}[i, j]` of the input tensor. Instead of\\n    setting activations to zero, as in regular Dropout, the activations are set\\n    to the negative saturation value of the SELU activation function.\\n\\n    Each element will be masked independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n    The elements to be masked are randomized on every forward call, and scaled\\n    and shifted to maintain zero mean and unit variance.\\n\\n    See :class:`~torch.nn.FeatureAlphaDropout` for details.\\n\\n    Args:\\n        p: dropout probability of a channel to be zeroed. Default: 0.5\\n        training: apply dropout if is ``True``. Default: ``True``\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            feature_alpha_dropout,\\n            (input,),\\n            input,\\n            p=p,\\n            training=training,\\n            inplace=inplace,\\n        )\\n    if p < 0.0 or p > 1.0:\\n        raise ValueError(f\\\"dropout probability has to be between 0 and 1, but got {p}\\\")\\n    return (\\n        _VF.feature_alpha_dropout_(input, p, training)\\n        if inplace\\n        else _VF.feature_alpha_dropout(input, p, training)\\n    )\\n\\n\\ndef _threshold(\\n    input: Tensor,\\n    threshold: float,\\n    value: float,\\n    inplace: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply a threshold to each element of the input Tensor.\\n\\n    See :class:`~torch.nn.Threshold` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            _threshold, (input,), input, threshold, value, inplace=inplace\\n        )\\n    if inplace:\\n        result = _VF.threshold_(input, threshold, value)\\n    else:\\n        result = _VF.threshold(input, threshold, value)\\n    return result\\n\\n\\n# We define this function as _threshold because it takes an argument\\n# named threshold, which clobbers the recursive reference to the\\n# function needed for __torch_function__ support\\nthreshold = _threshold\\n\\nthreshold_ = _add_docstr(\\n    _VF.threshold_,\\n    r\\\"\\\"\\\"\\nthreshold_(input, threshold, value) -> Tensor\\n\\nIn-place version of :func:`~threshold`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef relu(input: Tensor, inplace: bool = False) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"relu(input, inplace=False) -> Tensor\\n\\n    Applies the rectified linear unit function element-wise. See\\n    :class:`~torch.nn.ReLU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(relu, (input,), input, inplace=inplace)\\n    if inplace:\\n        result = torch.relu_(input)\\n    else:\\n        result = torch.relu(input)\\n    return result\\n\\n\\nrelu_ = _add_docstr(\\n    torch.relu_,\\n    r\\\"\\\"\\\"\\nrelu_(input) -> Tensor\\n\\nIn-place version of :func:`~relu`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef glu(input: Tensor, dim: int = -1) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"\\n    glu(input, dim=-1) -> Tensor\\n\\n    The gated linear unit. Computes:\\n\\n    .. math ::\\n        \\\\text{GLU}(a, b) = a \\\\otimes \\\\sigma(b)\\n\\n    where `input` is split in half along `dim` to form `a` and `b`, :math:`\\\\sigma`\\n    is the sigmoid function and :math:`\\\\otimes` is the element-wise product between matrices.\\n\\n    See `Language Modeling with Gated Convolutional Networks <https://arxiv.org/abs/1612.08083>`_.\\n\\n    Args:\\n        input (Tensor): input tensor\\n        dim (int): dimension on which to split the input. Default: -1\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(glu, (input,), input, dim=dim)\\n    if input.dim() == 0:\\n        raise RuntimeError(\\n            \\\"glu does not support scalars because halving size must be even\\\"\\n        )\\n    return torch._C._nn.glu(input, dim)\\n\\n\\ndef hardtanh(\\n    input: Tensor,\\n    min_val: float = -1.0,\\n    max_val: float = 1.0,\\n    inplace: bool = False,\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"\\n    hardtanh(input, min_val=-1., max_val=1., inplace=False) -> Tensor\\n\\n    Applies the HardTanh function element-wise. See :class:`~torch.nn.Hardtanh` for more\\n    details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            hardtanh, (input,), input, min_val=min_val, max_val=max_val, inplace=inplace\\n        )\\n    if min_val > max_val:\\n        raise ValueError(\\\"min_val cannot be greater than max_val\\\")\\n    if inplace:\\n        result = torch._C._nn.hardtanh_(input, min_val, max_val)\\n    else:\\n        result = torch._C._nn.hardtanh(input, min_val, max_val)\\n    return result\\n\\n\\nhardtanh_ = _add_docstr(\\n    torch._C._nn.hardtanh_,\\n    r\\\"\\\"\\\"\\nhardtanh_(input, min_val=-1., max_val=1.) -> Tensor\\n\\nIn-place version of :func:`~hardtanh`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef relu6(input: Tensor, inplace: bool = False) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"relu6(input, inplace=False) -> Tensor\\n\\n    Applies the element-wise function :math:`\\\\text{ReLU6}(x) = \\\\min(\\\\max(0,x), 6)`.\\n\\n    See :class:`~torch.nn.ReLU6` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(relu6, (input,), input, inplace=inplace)\\n    if inplace:\\n        result = torch._C._nn.relu6_(input)\\n    else:\\n        result = torch._C._nn.relu6(input)\\n    return result\\n\\n\\ndef elu(input: Tensor, alpha: float = 1.0, inplace: bool = False) -> Tensor:\\n    r\\\"\\\"\\\"Apply the Exponential Linear Unit (ELU) function element-wise.\\n\\n    See :class:`~torch.nn.ELU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(elu, (input,), input, alpha=alpha, inplace=inplace)\\n    if inplace:\\n        result = torch._C._nn.elu_(input, alpha)\\n    else:\\n        result = torch._C._nn.elu(input, alpha)\\n    return result\\n\\n\\nelu_ = _add_docstr(\\n    torch._C._nn.elu_,\\n    r\\\"\\\"\\\"\\nelu_(input, alpha=1.) -> Tensor\\n\\nIn-place version of :func:`~elu`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef selu(input: Tensor, inplace: bool = False) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"selu(input, inplace=False) -> Tensor\\n\\n    Applies element-wise,\\n    :math:`\\\\text{SELU}(x) = scale * (\\\\max(0,x) + \\\\min(0, \\\\alpha * (\\\\exp(x) - 1)))`,\\n    with :math:`\\\\alpha=1.6732632423543772848170429916717` and\\n    :math:`scale=1.0507009873554804934193349852946`.\\n\\n    See :class:`~torch.nn.SELU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(selu, (input,), input, inplace=inplace)\\n    if inplace:\\n        result = torch.selu_(input)\\n    else:\\n        result = torch.selu(input)\\n    return result\\n\\n\\nselu_ = _add_docstr(\\n    torch.selu_,\\n    r\\\"\\\"\\\"\\nselu_(input) -> Tensor\\n\\nIn-place version of :func:`~selu`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef celu(\\n    input: Tensor,\\n    alpha: float = 1.0,\\n    inplace: bool = False,\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"celu(input, alpha=1., inplace=False) -> Tensor\\n\\n    Applies element-wise,\\n    :math:`\\\\text{CELU}(x) = \\\\max(0,x) + \\\\min(0, \\\\alpha * (\\\\exp(x/\\\\alpha) - 1))`.\\n\\n    See :class:`~torch.nn.CELU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            celu, (input,), input, alpha=alpha, inplace=inplace\\n        )\\n    if inplace:\\n        result = torch.celu_(input, alpha)\\n    else:\\n        result = torch.celu(input, alpha)\\n    return result\\n\\n\\ncelu_ = _add_docstr(\\n    torch.celu_,\\n    r\\\"\\\"\\\"\\ncelu_(input, alpha=1.) -> Tensor\\n\\nIn-place version of :func:`~celu`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef leaky_relu(\\n    input: Tensor,\\n    negative_slope: float = 0.01,\\n    inplace: bool = False,\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"\\n    leaky_relu(input, negative_slope=0.01, inplace=False) -> Tensor\\n\\n    Applies element-wise,\\n    :math:`\\\\text{LeakyReLU}(x) = \\\\max(0, x) + \\\\text{negative\\\\_slope} * \\\\min(0, x)`\\n\\n    See :class:`~torch.nn.LeakyReLU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            leaky_relu, (input,), input, negative_slope=negative_slope, inplace=inplace\\n        )\\n    if inplace:\\n        result = torch._C._nn.leaky_relu_(input, negative_slope)\\n    else:\\n        result = torch._C._nn.leaky_relu(input, negative_slope)\\n    return result\\n\\n\\nleaky_relu_ = _add_docstr(\\n    torch._C._nn.leaky_relu_,\\n    r\\\"\\\"\\\"\\nleaky_relu_(input, negative_slope=0.01) -> Tensor\\n\\nIn-place version of :func:`~leaky_relu`.\\n\\\"\\\"\\\",\\n)\\n\\n\\nprelu = _add_docstr(\\n    torch.prelu,\\n    r\\\"\\\"\\\"prelu(input, weight) -> Tensor\\n\\nApplies element-wise the function\\n:math:`\\\\text{PReLU}(x) = \\\\max(0,x) + \\\\text{weight} * \\\\min(0,x)` where weight is a\\nlearnable parameter.\\n\\n.. note::\\n    `weight` is expected to be a scalar or 1-D tensor. If `weight` is 1-D,\\n    its size must match the number of input channels, determined by\\n    `input.size(1)` when `input.dim() >= 2`, otherwise 1.\\n    In the 1-D case, note that when `input` has dim > 2, `weight` can be expanded\\n    to the shape of `input` in a way that is not possible using normal\\n    :ref:`broadcasting semantics<broadcasting-semantics>`.\\n\\nSee :class:`~torch.nn.PReLU` for more details.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef rrelu(\\n    input: Tensor,\\n    lower: float = 1.0 / 8,\\n    upper: float = 1.0 / 3,\\n    training: bool = False,\\n    inplace: bool = False,\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"rrelu(input, lower=1./8, upper=1./3, training=False, inplace=False) -> Tensor\\n\\n    Randomized leaky ReLU.\\n\\n    See :class:`~torch.nn.RReLU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            rrelu,\\n            (input,),\\n            input,\\n            lower=lower,\\n            upper=upper,\\n            training=training,\\n            inplace=inplace,\\n        )\\n    if inplace:\\n        result = torch.rrelu_(input, lower, upper, training)\\n    else:\\n        result = torch.rrelu(input, lower, upper, training)\\n    return result\\n\\n\\nrrelu_ = _add_docstr(\\n    torch.rrelu_,\\n    r\\\"\\\"\\\"\\nrrelu_(input, lower=1./8, upper=1./3, training=False) -> Tensor\\n\\nIn-place version of :func:`~rrelu`.\\n\\\"\\\"\\\",\\n)\\n\\nlogsigmoid = _add_docstr(\\n    torch._C._nn.log_sigmoid,\\n    r\\\"\\\"\\\"\\nlogsigmoid(input) -> Tensor\\n\\nApplies element-wise :math:`\\\\text{LogSigmoid}(x_i) = \\\\log \\\\left(\\\\frac{1}{1 + \\\\exp(-x_i)}\\\\right)`\\n\\nSee :class:`~torch.nn.LogSigmoid` for more details.\\n\\\"\\\"\\\",\\n)\\n\\ngelu = _add_docstr(\\n    torch._C._nn.gelu,\\n    r\\\"\\\"\\\"\\ngelu(input, approximate = 'none') -> Tensor\\n\\nWhen the approximate argument is 'none', it applies element-wise the function\\n:math:`\\\\text{GELU}(x) = x * \\\\Phi(x)`\\n\\nwhere :math:`\\\\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.\\n\\nWhen the approximate argument is 'tanh', Gelu is estimated with\\n\\n.. math::\\n    \\\\text{GELU}(x) = 0.5 * x * (1 + \\\\text{Tanh}(\\\\sqrt{2 / \\\\pi} * (x + 0.044715 * x^3)))\\n\\nSee `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_.\\n\\\"\\\"\\\",\\n)\\n\\nhardshrink = _add_docstr(\\n    torch.hardshrink,\\n    r\\\"\\\"\\\"\\nhardshrink(input, lambd=0.5) -> Tensor\\n\\nApplies the hard shrinkage function element-wise\\n\\nSee :class:`~torch.nn.Hardshrink` for more details.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef tanhshrink(input):  # noqa: D400,D402\\n    r\\\"\\\"\\\"tanhshrink(input) -> Tensor\\n\\n    Applies element-wise, :math:`\\\\text{Tanhshrink}(x) = x - \\\\text{Tanh}(x)`\\n\\n    See :class:`~torch.nn.Tanhshrink` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(tanhshrink, (input,), input)\\n    return input - input.tanh()\\n\\n\\ndef softsign(input):  # noqa: D400,D402\\n    r\\\"\\\"\\\"softsign(input) -> Tensor\\n\\n    Applies element-wise, the function :math:`\\\\text{SoftSign}(x) = \\\\frac{x}{1 + |x|}`\\n\\n    See :class:`~torch.nn.Softsign` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(softsign, (input,), input)\\n    return input / (input.abs() + 1)\\n\\n\\nsoftplus = _add_docstr(\\n    torch._C._nn.softplus,\\n    r\\\"\\\"\\\"\\nsoftplus(input, beta=1, threshold=20) -> Tensor\\n\\nApplies element-wise, the function :math:`\\\\text{Softplus}(x) = \\\\frac{1}{\\\\beta} * \\\\log(1 + \\\\exp(\\\\beta * x))`.\\n\\nFor numerical stability the implementation reverts to the linear function\\nwhen :math:`input \\\\times \\\\beta > threshold`.\\n\\nSee :class:`~torch.nn.Softplus` for more details.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef _get_softmax_dim(name: str, ndim: int, stacklevel: int) -> int:\\n    warnings.warn(\\n        f\\\"Implicit dimension choice for {name} has been deprecated. \\\"\\n        \\\"Change the call to include dim=X as an argument.\\\",\\n        stacklevel=stacklevel,\\n    )\\n    if ndim == 0 or ndim == 1 or ndim == 3:\\n        ret = 0\\n    else:\\n        ret = 1\\n    return ret\\n\\n\\ndef softmin(\\n    input: Tensor,\\n    dim: Optional[int] = None,\\n    _stacklevel: int = 3,\\n    dtype: Optional[DType] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply a softmin function.\\n\\n    Note that :math:`\\\\text{Softmin}(x) = \\\\text{Softmax}(-x)`. See softmax definition for mathematical formula.\\n\\n    See :class:`~torch.nn.Softmin` for more details.\\n\\n    Args:\\n        input (Tensor): input\\n        dim (int): A dimension along which softmin will be computed (so every slice\\n            along dim will sum to 1).\\n        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.\\n          If specified, the input tensor is casted to :attr:`dtype` before the operation\\n          is performed. This is useful for preventing data type overflows. Default: None.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            softmin, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype\\n        )\\n    if dim is None:\\n        dim = _get_softmax_dim(\\\"softmin\\\", input.dim(), _stacklevel)\\n    if dtype is None:\\n        ret = (-input).softmax(dim)\\n    else:\\n        ret = (-input).softmax(dim, dtype=dtype)\\n    return ret\\n\\n\\ndef softmax(\\n    input: Tensor,\\n    dim: Optional[int] = None,\\n    _stacklevel: int = 3,\\n    dtype: Optional[DType] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply a softmax function.\\n\\n    Softmax is defined as:\\n\\n    :math:`\\\\text{Softmax}(x_{i}) = \\\\frac{\\\\exp(x_i)}{\\\\sum_j \\\\exp(x_j)}`\\n\\n    It is applied to all slices along dim, and will re-scale them so that the elements\\n    lie in the range `[0, 1]` and sum to 1.\\n\\n    See :class:`~torch.nn.Softmax` for more details.\\n\\n    Args:\\n        input (Tensor): input\\n        dim (int): A dimension along which softmax will be computed.\\n        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.\\n          If specified, the input tensor is casted to :attr:`dtype` before the operation\\n          is performed. This is useful for preventing data type overflows. Default: None.\\n\\n    .. note::\\n        This function doesn't work directly with NLLLoss,\\n        which expects the Log to be computed between the Softmax and itself.\\n        Use log_softmax instead (it's faster and has better numerical properties).\\n\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            softmax, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype\\n        )\\n    if dim is None:\\n        dim = _get_softmax_dim(\\\"softmax\\\", input.dim(), _stacklevel)\\n    if dtype is None:\\n        ret = input.softmax(dim)\\n    else:\\n        ret = input.softmax(dim, dtype=dtype)\\n    return ret\\n\\n\\ndef gumbel_softmax(\\n    logits: Tensor,\\n    tau: float = 1,\\n    hard: bool = False,\\n    eps: float = 1e-10,\\n    dim: int = -1,\\n) -> Tensor:\\n    r\\\"\\\"\\\"\\n    Sample from the Gumbel-Softmax distribution (`Link 1`_  `Link 2`_) and optionally discretize.\\n\\n    Args:\\n      logits: `[..., num_features]` unnormalized log probabilities\\n      tau: non-negative scalar temperature\\n      hard: if ``True``, the returned samples will be discretized as one-hot vectors,\\n            but will be differentiated as if it is the soft sample in autograd\\n      dim (int): A dimension along which softmax will be computed. Default: -1.\\n\\n    Returns:\\n      Sampled tensor of same shape as `logits` from the Gumbel-Softmax distribution.\\n      If ``hard=True``, the returned samples will be one-hot, otherwise they will\\n      be probability distributions that sum to 1 across `dim`.\\n\\n    .. note::\\n      This function is here for legacy reasons, may be removed from nn.Functional in the future.\\n\\n    .. note::\\n      The main trick for `hard` is to do  `y_hard - y_soft.detach() + y_soft`\\n\\n      It achieves two things:\\n      - makes the output value exactly one-hot\\n      (since we add then subtract y_soft value)\\n      - makes the gradient equal to y_soft gradient\\n      (since we strip all other gradients)\\n\\n    Examples::\\n        >>> logits = torch.randn(20, 32)\\n        >>> # Sample soft categorical using reparametrization trick:\\n        >>> F.gumbel_softmax(logits, tau=1, hard=False)\\n        >>> # Sample hard categorical using \\\"Straight-through\\\" trick:\\n        >>> F.gumbel_softmax(logits, tau=1, hard=True)\\n\\n    .. _Link 1:\\n        https://arxiv.org/abs/1611.00712\\n    .. _Link 2:\\n        https://arxiv.org/abs/1611.01144\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(logits):\\n        return handle_torch_function(\\n            gumbel_softmax, (logits,), logits, tau=tau, hard=hard, eps=eps, dim=dim\\n        )\\n    if eps != 1e-10:\\n        warnings.warn(\\\"`eps` parameter is deprecated and has no effect.\\\")\\n\\n    gumbels = (\\n        -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format)\\n        .exponential_()\\n        .log()\\n    )  # ~Gumbel(0,1)\\n    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)\\n    y_soft = gumbels.softmax(dim)\\n\\n    if hard:\\n        # Straight through.\\n        index = y_soft.max(dim, keepdim=True)[1]\\n        y_hard = torch.zeros_like(\\n            logits, memory_format=torch.legacy_contiguous_format\\n        ).scatter_(dim, index, 1.0)\\n        ret = y_hard - y_soft.detach() + y_soft\\n    else:\\n        # Reparametrization trick.\\n        ret = y_soft\\n    return ret\\n\\n\\ndef log_softmax(\\n    input: Tensor,\\n    dim: Optional[int] = None,\\n    _stacklevel: int = 3,\\n    dtype: Optional[DType] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply a softmax followed by a logarithm.\\n\\n    While mathematically equivalent to log(softmax(x)), doing these two\\n    operations separately is slower and numerically unstable. This function\\n    uses an alternative formulation to compute the output and gradient correctly.\\n\\n    See :class:`~torch.nn.LogSoftmax` for more details.\\n\\n    Args:\\n        input (Tensor): input\\n        dim (int): A dimension along which log_softmax will be computed.\\n        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.\\n          If specified, the input tensor is cast to :attr:`dtype` before the operation\\n          is performed. This is useful for preventing data type overflows. Default: None.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            log_softmax, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype\\n        )\\n    if dim is None:\\n        dim = _get_softmax_dim(\\\"log_softmax\\\", input.dim(), _stacklevel)\\n    if dtype is None:\\n        ret = input.log_softmax(dim)\\n    else:\\n        ret = input.log_softmax(dim, dtype=dtype)\\n    return ret\\n\\n\\nsoftshrink = _add_docstr(\\n    torch._C._nn.softshrink,\\n    r\\\"\\\"\\\"\\nsoftshrink(input, lambd=0.5) -> Tensor\\n\\nApplies the soft shrinkage function elementwise\\n\\nSee :class:`~torch.nn.Softshrink` for more details.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef tanh(input):  # noqa: D400,D402\\n    r\\\"\\\"\\\"tanh(input) -> Tensor\\n\\n    Applies element-wise,\\n    :math:`\\\\text{Tanh}(x) = \\\\tanh(x) = \\\\frac{\\\\exp(x) - \\\\exp(-x)}{\\\\exp(x) + \\\\exp(-x)}`\\n\\n    See :class:`~torch.nn.Tanh` for more details.\\n    \\\"\\\"\\\"\\n    return input.tanh()\\n\\n\\ndef sigmoid(input):  # noqa: D400,D402\\n    r\\\"\\\"\\\"sigmoid(input) -> Tensor\\n\\n    Applies the element-wise function :math:`\\\\text{Sigmoid}(x) = \\\\frac{1}{1 + \\\\exp(-x)}`\\n\\n    See :class:`~torch.nn.Sigmoid` for more details.\\n    \\\"\\\"\\\"\\n    return input.sigmoid()\\n\\n\\ndef hardsigmoid(input: Tensor, inplace: bool = False) -> Tensor:\\n    r\\\"\\\"\\\"Apply the Hardsigmoid function element-wise.\\n\\n    .. math::\\n        \\\\text{Hardsigmoid}(x) = \\\\begin{cases}\\n            0 & \\\\text{if~} x \\\\le -3, \\\\\\\\\\n            1 & \\\\text{if~} x \\\\ge +3, \\\\\\\\\\n            x / 6 + 1 / 2 & \\\\text{otherwise}\\n        \\\\end{cases}\\n\\n    Args:\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n\\n    See :class:`~torch.nn.Hardsigmoid` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(hardsigmoid, (input,), input, inplace=inplace)\\n    if inplace:\\n        return torch._C._nn.hardsigmoid_(input)\\n    return torch._C._nn.hardsigmoid(input)\\n\\n\\nlinear = _add_docstr(\\n    torch._C._nn.linear,\\n    r\\\"\\\"\\\"\\nlinear(input, weight, bias=None) -> Tensor\\n\\nApplies a linear transformation to the incoming data: :math:`y = xA^T + b`.\\n\\nThis operation supports 2-D :attr:`weight` with :ref:`sparse layout<sparse-docs>`\\n\\n{sparse_beta_warning}\\n\\nThis operator supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\nShape:\\n\\n    - Input: :math:`(*, in\\\\_features)` where `*` means any number of\\n      additional dimensions, including none\\n    - Weight: :math:`(out\\\\_features, in\\\\_features)` or :math:`(in\\\\_features)`\\n    - Bias: :math:`(out\\\\_features)` or :math:`()`\\n    - Output: :math:`(*, out\\\\_features)` or :math:`(*)`, based on the shape of the weight\\n\\\"\\\"\\\".format(\\n        **sparse_support_notes\\n    ),\\n)\\n\\n\\nbilinear = _add_docstr(\\n    torch.bilinear,\\n    r\\\"\\\"\\\"\\nbilinear(input1, input2, weight, bias=None) -> Tensor\\n\\nApplies a bilinear transformation to the incoming data:\\n:math:`y = x_1^T A x_2 + b`\\n\\nShape:\\n\\n    - input1: :math:`(N, *, H_{in1})` where :math:`H_{in1}=\\\\text{in1\\\\_features}`\\n      and :math:`*` means any number of additional dimensions.\\n      All but the last dimension of the inputs should be the same.\\n    - input2: :math:`(N, *, H_{in2})` where :math:`H_{in2}=\\\\text{in2\\\\_features}`\\n    - weight: :math:`(\\\\text{out\\\\_features}, \\\\text{in1\\\\_features},\\n      \\\\text{in2\\\\_features})`\\n    - bias: :math:`(\\\\text{out\\\\_features})`\\n    - output: :math:`(N, *, H_{out})` where :math:`H_{out}=\\\\text{out\\\\_features}`\\n      and all but the last dimension are the same shape as the input.\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef silu(input: Tensor, inplace: bool = False) -> Tensor:\\n    r\\\"\\\"\\\"Apply the Sigmoid Linear Unit (SiLU) function, element-wise.\\n\\n    The SiLU function is also known as the swish function.\\n\\n    .. math::\\n        \\\\text{silu}(x) = x * \\\\sigma(x), \\\\text{where } \\\\sigma(x) \\\\text{ is the logistic sigmoid.}\\n\\n    .. note::\\n        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_\\n        where the SiLU (Sigmoid Linear Unit) was originally coined, and see\\n        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation\\n        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:\\n        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_\\n        where the SiLU was experimented with later.\\n\\n    See :class:`~torch.nn.SiLU` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(silu, (input,), input, inplace=inplace)\\n    if inplace:\\n        return torch._C._nn.silu_(input)\\n    return torch._C._nn.silu(input)\\n\\n\\ndef mish(input: Tensor, inplace: bool = False) -> Tensor:\\n    r\\\"\\\"\\\"Apply the Mish function, element-wise.\\n\\n    Mish: A Self Regularized Non-Monotonic Neural Activation Function.\\n\\n    .. math::\\n        \\\\text{Mish}(x) = x * \\\\text{Tanh}(\\\\text{Softplus}(x))\\n\\n    .. note::\\n        See `Mish: A Self Regularized Non-Monotonic Neural Activation Function <https://arxiv.org/abs/1908.08681>`_\\n\\n    See :class:`~torch.nn.Mish` for more details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(mish, (input,), input, inplace=inplace)\\n    if inplace:\\n        return torch._C._nn.mish_(input)\\n    return torch._C._nn.mish(input)\\n\\n\\ndef hardswish(input: Tensor, inplace: bool = False) -> Tensor:\\n    r\\\"\\\"\\\"Apply hardswish function, element-wise.\\n\\n    Follows implementation as described in the paper:\\n    `Searching for MobileNetV3`_.\\n\\n    .. math::\\n        \\\\text{Hardswish}(x) = \\\\begin{cases}\\n            0 & \\\\text{if~} x \\\\le -3, \\\\\\\\\\n            x & \\\\text{if~} x \\\\ge +3, \\\\\\\\\\n            x \\\\cdot (x + 3) /6 & \\\\text{otherwise}\\n        \\\\end{cases}\\n\\n    See :class:`~torch.nn.Hardswish` for more details.\\n\\n    .. _`Searching for MobileNetV3`:\\n        https://arxiv.org/abs/1905.02244\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(hardswish, (input,), input, inplace=inplace)\\n    if inplace:\\n        return torch._C._nn.hardswish_(input)\\n    return torch._C._nn.hardswish(input)\\n\\n\\ndef _no_grad_embedding_renorm_(\\n    weight: Tensor,\\n    input: Tensor,\\n    max_norm: float,\\n    norm_type: float,\\n) -> Tuple[Tensor, Tensor]:\\n    torch.embedding_renorm_(weight.detach(), input, max_norm, norm_type)\\n\\n\\ndef embedding(\\n    input: Tensor,\\n    weight: Tensor,\\n    padding_idx: Optional[int] = None,\\n    max_norm: Optional[float] = None,\\n    norm_type: float = 2.0,\\n    scale_grad_by_freq: bool = False,\\n    sparse: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Generate a simple lookup table that looks up embeddings in a fixed dictionary and size.\\n\\n    This module is often used to retrieve word embeddings using indices.\\n    The input to the module is a list of indices, and the embedding matrix,\\n    and the output is the corresponding word embeddings.\\n\\n    See :class:`torch.nn.Embedding` for more details.\\n\\n    .. note::\\n        Note that the analytical gradients of this function with respect to\\n        entries in :attr:`weight` at the row specified by :attr:`padding_idx`\\n        are expected to differ from the numerical ones.\\n\\n    .. note::\\n        Note that `:class:`torch.nn.Embedding` differs from this function in\\n        that it initializes the row of :attr:`weight` specified by\\n        :attr:`padding_idx` to all zeros on construction.\\n\\n    Args:\\n        input (LongTensor): Tensor containing indices into the embedding matrix\\n        weight (Tensor): The embedding matrix with number of rows equal to the maximum possible index + 1,\\n            and number of columns equal to the embedding size\\n        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;\\n                                     therefore, the embedding vector at :attr:`padding_idx` is not updated during training,\\n                                     i.e. it remains as a fixed \\\"pad\\\".\\n        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`\\n                                    is renormalized to have norm :attr:`max_norm`.\\n                                    Note: this will modify :attr:`weight` in-place.\\n        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.\\n        scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse of frequency of\\n                                                the words in the mini-batch. Default ``False``.\\n        sparse (bool, optional): If ``True``, gradient w.r.t. :attr:`weight` will be a sparse tensor. See Notes under\\n                                 :class:`torch.nn.Embedding` for more details regarding sparse gradients.\\n\\n    Shape:\\n        - Input: LongTensor of arbitrary shape containing the indices to extract\\n        - Weight: Embedding matrix of floating point type with shape `(V, embedding_dim)`,\\n          where V = maximum index + 1 and embedding_dim = the embedding size\\n        - Output: `(*, embedding_dim)`, where `*` is the input shape\\n\\n    Examples::\\n\\n        >>> # a batch of 2 samples of 4 indices each\\n        >>> input = torch.tensor([[1, 2, 4, 5], [4, 3, 2, 9]])\\n        >>> # an embedding matrix containing 10 tensors of size 3\\n        >>> embedding_matrix = torch.rand(10, 3)\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> F.embedding(input, embedding_matrix)\\n        tensor([[[ 0.8490,  0.9625,  0.6753],\\n                 [ 0.9666,  0.7761,  0.6108],\\n                 [ 0.6246,  0.9751,  0.3618],\\n                 [ 0.4161,  0.2419,  0.7383]],\\n\\n                [[ 0.6246,  0.9751,  0.3618],\\n                 [ 0.0237,  0.7794,  0.0528],\\n                 [ 0.9666,  0.7761,  0.6108],\\n                 [ 0.3385,  0.8612,  0.1867]]])\\n\\n        >>> # example with padding_idx\\n        >>> weights = torch.rand(10, 3)\\n        >>> weights[0, :].zero_()\\n        >>> embedding_matrix = weights\\n        >>> input = torch.tensor([[0, 2, 0, 5]])\\n        >>> F.embedding(input, embedding_matrix, padding_idx=0)\\n        tensor([[[ 0.0000,  0.0000,  0.0000],\\n                 [ 0.5609,  0.5384,  0.8720],\\n                 [ 0.0000,  0.0000,  0.0000],\\n                 [ 0.6262,  0.2438,  0.7471]]])\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, weight):\\n        return handle_torch_function(\\n            embedding,\\n            (input, weight),\\n            input,\\n            weight,\\n            padding_idx=padding_idx,\\n            max_norm=max_norm,\\n            norm_type=norm_type,\\n            scale_grad_by_freq=scale_grad_by_freq,\\n            sparse=sparse,\\n        )\\n    if padding_idx is not None:\\n        if padding_idx > 0:\\n            assert padding_idx < weight.size(\\n                0\\n            ), \\\"Padding_idx must be within num_embeddings\\\"\\n        elif padding_idx < 0:\\n            assert padding_idx >= -weight.size(\\n                0\\n            ), \\\"Padding_idx must be within num_embeddings\\\"\\n            padding_idx = weight.size(0) + padding_idx\\n    else:\\n        padding_idx = -1\\n    if max_norm is not None:\\n        # Note [embedding_renorm contiguous]\\n        # `embedding_renorm_` will call .contiguous() on input anyways, so we\\n        # call it here and take advantage of the improved locality in the\\n        # `embedding` call below too.\\n        input = input.contiguous()\\n        # Note [embedding_renorm set_grad_enabled]\\n        # XXX: equivalent to\\n        # with torch.no_grad():\\n        #   torch.embedding_renorm_\\n        # remove once script supports set_grad_enabled\\n        _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)\\n    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)\\n\\n\\ndef embedding_bag(\\n    input: Tensor,\\n    weight: Tensor,\\n    offsets: Optional[Tensor] = None,\\n    max_norm: Optional[float] = None,\\n    norm_type: float = 2,\\n    scale_grad_by_freq: bool = False,\\n    mode: str = \\\"mean\\\",\\n    sparse: bool = False,\\n    per_sample_weights: Optional[Tensor] = None,\\n    include_last_offset: bool = False,\\n    padding_idx: Optional[int] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute sums, means or maxes of `bags` of embeddings.\\n\\n    Calculation is done without instantiating the intermediate embeddings.\\n    See :class:`torch.nn.EmbeddingBag` for more details.\\n\\n    Note:\\n        {backward_reproducibility_note}\\n\\n    Args:\\n        input (LongTensor): Tensor containing bags of indices into the embedding matrix\\n        weight (Tensor): The embedding matrix with number of rows equal to the maximum possible index + 1,\\n            and number of columns equal to the embedding size\\n        offsets (LongTensor, optional): Only used when :attr:`input` is 1D. :attr:`offsets` determines\\n                             the starting index position of each bag (sequence) in :attr:`input`.\\n        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`\\n                                    is renormalized to have norm :attr:`max_norm`.\\n                                    Note: this will modify :attr:`weight` in-place.\\n        norm_type (float, optional): The ``p`` in the ``p``-norm to compute for the :attr:`max_norm` option.\\n                                     Default ``2``.\\n        scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of\\n                                                the words in the mini-batch. Default ``False``.\\n                                                Note: this option is not supported when ``mode=\\\"max\\\"``.\\n        mode (str, optional): ``\\\"sum\\\"``, ``\\\"mean\\\"`` or ``\\\"max\\\"``. Specifies the way to reduce the bag.\\n                                 Default: ``\\\"mean\\\"``\\n        sparse (bool, optional): if ``True``, gradient w.r.t. :attr:`weight` will be a sparse tensor. See Notes under\\n                                 :class:`torch.nn.Embedding` for more details regarding sparse gradients.\\n                                 Note: this option is not supported when ``mode=\\\"max\\\"``.\\n        per_sample_weights (Tensor, optional): a tensor of float / double weights, or None\\n            to indicate all weights should be taken to be 1. If specified, :attr:`per_sample_weights`\\n            must have exactly the same shape as input and is treated as having the same\\n            :attr:`offsets`, if those are not None.\\n\\n        include_last_offset (bool, optional): if ``True``, the size of offsets is equal to the number of bags + 1.\\n            The last element is the size of the input, or the ending index position of the last bag (sequence).\\n\\n        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the\\n                                     gradient; therefore, the embedding vector at :attr:`padding_idx` is not updated\\n                                     during training, i.e. it remains as a fixed \\\"pad\\\". Note that the embedding\\n                                     vector at :attr:`padding_idx` is excluded from the reduction.\\n\\n    Shape:\\n        - :attr:`input` (LongTensor) and :attr:`offsets` (LongTensor, optional)\\n\\n          - If :attr:`input` is 2D of shape `(B, N)`, it will be treated as ``B`` bags (sequences)\\n            each of fixed length ``N``, and this will return ``B`` values aggregated in a way\\n            depending on the :attr:`mode`. :attr:`offsets` is ignored and required to be ``None`` in this case.\\n\\n          - If :attr:`input` is 1D of shape `(N)`, it will be treated as a concatenation of\\n            multiple bags (sequences). :attr:`offsets` is required to be a 1D tensor containing\\n            the starting index positions of each bag in :attr:`input`. Therefore, for :attr:`offsets`\\n            of shape `(B)`, :attr:`input` will be viewed as having ``B`` bags.\\n            Empty bags (i.e., having 0-length) will have returned vectors filled by zeros.\\n\\n        - :attr:`weight` (Tensor): the learnable weights of the module of shape `(num_embeddings, embedding_dim)`\\n\\n        - :attr:`per_sample_weights` (Tensor, optional). Has the same shape as :attr:`input`.\\n\\n        - :attr:`output`: aggregated embedding values of shape `(B, embedding_dim)`\\n\\n    Examples::\\n\\n        >>> # an Embedding module containing 10 tensors of size 3\\n        >>> embedding_matrix = torch.rand(10, 3)\\n        >>> # a batch of 2 samples of 4 indices each\\n        >>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9])\\n        >>> offsets = torch.tensor([0, 4])\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> F.embedding_bag(input, embedding_matrix, offsets)\\n        tensor([[ 0.3397,  0.3552,  0.5545],\\n                [ 0.5893,  0.4386,  0.5882]])\\n\\n        >>> # example with padding_idx\\n        >>> embedding_matrix = torch.rand(10, 3)\\n        >>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9])\\n        >>> offsets = torch.tensor([0, 4])\\n        >>> F.embedding_bag(input, embedding_matrix, offsets, padding_idx=2, mode='sum')\\n        tensor([[ 0.0000,  0.0000,  0.0000],\\n                [-0.7082,  3.2145, -2.6251]])\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, weight, offsets, per_sample_weights):\\n        return handle_torch_function(\\n            embedding_bag,\\n            (input, weight, offsets, per_sample_weights),\\n            input,\\n            weight,\\n            offsets=offsets,\\n            max_norm=max_norm,\\n            norm_type=norm_type,\\n            scale_grad_by_freq=scale_grad_by_freq,\\n            mode=mode,\\n            sparse=sparse,\\n            per_sample_weights=per_sample_weights,\\n            include_last_offset=include_last_offset,\\n            padding_idx=padding_idx,\\n        )\\n    # Check for backward compatibility.\\n    # Used to be embedding_bag(weight, input, ...)\\n    # Now is     embedding_bag(input, weight, ...)\\n    if weight.dtype == torch.long and input.is_floating_point():\\n        warnings.warn(\\n            \\\"Argument order of nn.functional.embedding_bag was changed. \\\"\\n            \\\"Usage `embedding_bag(weight, input, ...)` is deprecated, \\\"\\n            \\\"and should now be `embedding_bag(input, weight, ...)`.\\\"\\n        )\\n        weight, input = input, weight\\n\\n    if per_sample_weights is not None and input.size() != per_sample_weights.size():\\n        raise ValueError(\\n            f\\\"embedding_bag: If per_sample_weights ({per_sample_weights.shape}) is not None, \\\"\\n            f\\\"then it must have the same shape as the input ({input.shape})\\\"\\n        )\\n\\n    if not weight.dim() == 2:\\n        raise ValueError(\\n            f\\\"weight has to be a 2D Tensor, but got Tensor of dimension {weight.dim()}\\\"\\n        )\\n\\n    if input.dim() == 2:\\n        if offsets is not None:\\n            type_str = \\\"<unknown>\\\"\\n            # TODO: Remove this once script supports type() calls\\n            if not torch.jit.is_scripting():\\n                type_str = str(type(offsets))\\n            raise ValueError(\\n                \\\"if input is 2D, then offsets has to be None\\\"\\n                \\\", as input is treated is a mini-batch of\\\"\\n                \\\" fixed length sequences. However, found \\\"\\n                f\\\"offsets of type {type_str}\\\"\\n            )\\n        offsets = torch.arange(\\n            0, input.numel(), input.size(1), dtype=input.dtype, device=input.device\\n        )\\n\\n        input = input.reshape(-1)\\n        if per_sample_weights is not None:\\n            per_sample_weights = per_sample_weights.reshape(-1)\\n    elif input.dim() == 1:\\n        if offsets is None:\\n            raise ValueError(\\\"offsets has to be a 1D Tensor but got None\\\")\\n        if offsets.dim() != 1:\\n            raise ValueError(\\\"offsets has to be a 1D Tensor\\\")\\n    else:\\n        raise ValueError(\\n            f\\\"input has to be 1D or 2D Tensor, but got Tensor of dimension {input.dim()}\\\"\\n        )\\n    if mode == \\\"sum\\\":\\n        mode_enum = 0\\n    elif mode == \\\"mean\\\":\\n        mode_enum = 1\\n    elif mode == \\\"max\\\":\\n        mode_enum = 2\\n\\n        if scale_grad_by_freq:\\n            raise ValueError(\\n                \\\"max mode does not support scaling the gradient by the frequency\\\"\\n            )\\n\\n        if sparse:\\n            raise ValueError(\\\"max mode does not support sparse weights\\\")\\n\\n    else:\\n        raise ValueError(\\\"mode has to be one of sum, mean or max\\\")\\n\\n    if max_norm is not None:\\n        # XXX: equivalent to\\n        # with torch.no_grad():\\n        #   torch.nembedding_renorm_\\n        # remove once script supports set_grad_enabled\\n        _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)\\n\\n    if per_sample_weights is not None and mode != \\\"sum\\\":\\n        raise NotImplementedError(\\n            \\\"embedding_bag: per_sample_weights was not None. \\\"\\n            \\\"per_sample_weights is only supported for mode='sum' \\\"\\n            f\\\"(got mode='{mode}'). Please open a feature request on GitHub.\\\"\\n        )\\n\\n    ret, _, _, _ = torch.embedding_bag(\\n        weight,\\n        input,\\n        offsets,\\n        scale_grad_by_freq,\\n        mode_enum,\\n        sparse,\\n        per_sample_weights,\\n        include_last_offset,\\n        padding_idx,\\n    )\\n    return ret\\n\\n\\nif embedding_bag.__doc__:\\n    embedding_bag.__doc__ = embedding_bag.__doc__.format(**reproducibility_notes)\\n\\n\\ndef _verify_batch_size(size: List[int]) -> None:\\n    # XXX: JIT script does not support the reduce from functools, and mul op is a\\n    # builtin, which cannot be used as a value to a func yet, so rewrite this size\\n    # check to a simple equivalent for loop\\n    #\\n    # TODO: make use of reduce like below when JIT is ready with the missing features:\\n    # from operator import mul\\n    # from functools import reduce\\n    #\\n    #   if reduce(mul, size[2:], size[0]) == 1\\n    size_prods = size[0]\\n    for i in range(len(size) - 2):\\n        size_prods *= size[i + 2]\\n    if size_prods == 1:\\n        raise ValueError(\\n            f\\\"Expected more than 1 value per channel when training, got input size {size}\\\"\\n        )\\n\\n\\ndef batch_norm(\\n    input: Tensor,\\n    running_mean: Optional[Tensor],\\n    running_var: Optional[Tensor],\\n    weight: Optional[Tensor] = None,\\n    bias: Optional[Tensor] = None,\\n    training: bool = False,\\n    momentum: float = 0.1,\\n    eps: float = 1e-5,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply Batch Normalization for each channel across a batch of data.\\n\\n    See :class:`~torch.nn.BatchNorm1d`, :class:`~torch.nn.BatchNorm2d`,\\n    :class:`~torch.nn.BatchNorm3d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, running_mean, running_var, weight, bias):\\n        return handle_torch_function(\\n            batch_norm,\\n            (input, running_mean, running_var, weight, bias),\\n            input,\\n            running_mean,\\n            running_var,\\n            weight=weight,\\n            bias=bias,\\n            training=training,\\n            momentum=momentum,\\n            eps=eps,\\n        )\\n    if training:\\n        _verify_batch_size(input.size())\\n\\n    return torch.batch_norm(\\n        input,\\n        weight,\\n        bias,\\n        running_mean,\\n        running_var,\\n        training,\\n        momentum,\\n        eps,\\n        torch.backends.cudnn.enabled,\\n    )\\n\\n\\ndef _verify_spatial_size(size: List[int]) -> None:\\n    # Verify that there is > 1 spatial element for instance norm calculation.\\n    size_prods = 1\\n    for i in range(2, len(size)):\\n        size_prods *= size[i]\\n    if size_prods == 1:\\n        raise ValueError(\\n            f\\\"Expected more than 1 spatial element when training, got input size {size}\\\"\\n        )\\n\\n\\ndef instance_norm(\\n    input: Tensor,\\n    running_mean: Optional[Tensor] = None,\\n    running_var: Optional[Tensor] = None,\\n    weight: Optional[Tensor] = None,\\n    bias: Optional[Tensor] = None,\\n    use_input_stats: bool = True,\\n    momentum: float = 0.1,\\n    eps: float = 1e-5,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply Instance Normalization independently for each channel in every data sample within a batch.\\n\\n    See :class:`~torch.nn.InstanceNorm1d`, :class:`~torch.nn.InstanceNorm2d`,\\n    :class:`~torch.nn.InstanceNorm3d` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, running_mean, running_var, weight, bias):\\n        return handle_torch_function(\\n            instance_norm,\\n            (input, running_mean, running_var, weight, bias),\\n            input,\\n            running_mean=running_mean,\\n            running_var=running_var,\\n            weight=weight,\\n            bias=bias,\\n            use_input_stats=use_input_stats,\\n            momentum=momentum,\\n            eps=eps,\\n        )\\n    if use_input_stats:\\n        _verify_spatial_size(input.size())\\n    return torch.instance_norm(\\n        input,\\n        weight,\\n        bias,\\n        running_mean,\\n        running_var,\\n        use_input_stats,\\n        momentum,\\n        eps,\\n        torch.backends.cudnn.enabled,\\n    )\\n\\n\\ndef layer_norm(\\n    input: Tensor,\\n    normalized_shape: List[int],\\n    weight: Optional[Tensor] = None,\\n    bias: Optional[Tensor] = None,\\n    eps: float = 1e-5,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply Layer Normalization for last certain number of dimensions.\\n\\n    See :class:`~torch.nn.LayerNorm` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, weight, bias):\\n        return handle_torch_function(\\n            layer_norm,\\n            (input, weight, bias),\\n            input,\\n            normalized_shape,\\n            weight=weight,\\n            bias=bias,\\n            eps=eps,\\n        )\\n    return torch.layer_norm(\\n        input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled\\n    )\\n\\n\\ndef rms_norm(\\n    input: Tensor,\\n    normalized_shape: List[int],\\n    weight: Optional[Tensor] = None,\\n    eps: Optional[float] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply Root Mean Square Layer Normalization.\\n\\n    See :class:`~torch.nn.RMSNorm` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, weight):\\n        return handle_torch_function(\\n            rms_norm, (input, weight), input, normalized_shape, weight=weight, eps=eps\\n        )\\n    return torch.rms_norm(input, normalized_shape, weight, eps)\\n\\n\\ndef group_norm(\\n    input: Tensor,\\n    num_groups: int,\\n    weight: Optional[Tensor] = None,\\n    bias: Optional[Tensor] = None,\\n    eps: float = 1e-5,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply Group Normalization for last certain number of dimensions.\\n\\n    See :class:`~torch.nn.GroupNorm` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, weight, bias):\\n        return handle_torch_function(\\n            group_norm,\\n            (\\n                input,\\n                weight,\\n                bias,\\n            ),\\n            input,\\n            num_groups,\\n            weight=weight,\\n            bias=bias,\\n            eps=eps,\\n        )\\n    if input.dim() < 2:\\n        raise RuntimeError(\\n            f\\\"Expected at least 2 dimensions for input tensor but received {input.dim()}\\\"\\n        )\\n    _verify_batch_size(\\n        [input.size(0) * input.size(1) // num_groups, num_groups]\\n        + list(input.size()[2:])\\n    )\\n    return torch.group_norm(\\n        input, num_groups, weight, bias, eps, torch.backends.cudnn.enabled\\n    )\\n\\n\\ndef local_response_norm(\\n    input: Tensor,\\n    size: int,\\n    alpha: float = 1e-4,\\n    beta: float = 0.75,\\n    k: float = 1.0,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply local response normalization over an input signal.\\n\\n    The input signal is composed of several input planes, where channels occupy the second dimension.\\n    Normalization is applied across channels.\\n\\n    See :class:`~torch.nn.LocalResponseNorm` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            local_response_norm, (input,), input, size, alpha=alpha, beta=beta, k=k\\n        )\\n    dim = input.dim()\\n    if dim < 3:\\n        raise ValueError(\\n            f\\\"Expected 3D or higher dimensionality                          input (got {dim} dimensions)\\\"\\n        )\\n\\n    if input.numel() == 0:\\n        return input\\n\\n    div = input.mul(input)\\n    if dim == 3:\\n        div = div.unsqueeze(1)\\n        div = pad(div, (0, 0, size // 2, (size - 1) // 2))\\n        div = avg_pool2d(div, (size, 1), stride=1).squeeze(1)\\n    else:\\n        sizes = input.size()\\n        div = div.view(sizes[0], 1, sizes[1], sizes[2], -1)\\n        div = pad(div, (0, 0, 0, 0, size // 2, (size - 1) // 2))\\n        div = avg_pool3d(div, (size, 1, 1), stride=1).squeeze(1)\\n        div = div.view(sizes)\\n    div = div.mul(alpha).add(k).pow(beta)\\n    return input / div\\n\\n\\n# loss\\n\\n\\ndef ctc_loss(\\n    log_probs: Tensor,\\n    targets: Tensor,\\n    input_lengths: Tensor,\\n    target_lengths: Tensor,\\n    blank: int = 0,\\n    reduction: str = \\\"mean\\\",\\n    zero_infinity: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Apply the Connectionist Temporal Classification loss.\\n\\n    See :class:`~torch.nn.CTCLoss` for details.\\n\\n    Note:\\n        {cudnn_reproducibility_note}\\n\\n    Note:\\n        {backward_reproducibility_note}\\n\\n    Args:\\n        log_probs: :math:`(T, N, C)` or :math:`(T, C)` where `C = number of characters in alphabet including blank`,\\n            `T = input length`, and `N = batch size`.\\n            The logarithmized probabilities of the outputs\\n            (e.g. obtained with :func:`torch.nn.functional.log_softmax`).\\n        targets: :math:`(N, S)` or `(sum(target_lengths))`.\\n            Targets cannot be blank. In the second form, the targets are assumed to be concatenated.\\n        input_lengths: :math:`(N)` or :math:`()`.\\n            Lengths of the inputs (must each be :math:`\\\\leq T`)\\n        target_lengths: :math:`(N)` or :math:`()`.\\n            Lengths of the targets\\n        blank (int, optional):\\n            Blank label. Default :math:`0`.\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the output losses will be divided by the target lengths and\\n            then the mean over the batch is taken, ``'sum'``: the output will be\\n            summed. Default: ``'mean'``\\n        zero_infinity (bool, optional):\\n            Whether to zero infinite losses and the associated gradients.\\n            Default: ``False``\\n            Infinite losses mainly occur when the inputs are too short\\n            to be aligned to the targets.\\n\\n    Example::\\n\\n        >>> log_probs = torch.randn(50, 16, 20).log_softmax(2).detach().requires_grad_()\\n        >>> targets = torch.randint(1, 20, (16, 30), dtype=torch.long)\\n        >>> input_lengths = torch.full((16,), 50, dtype=torch.long)\\n        >>> target_lengths = torch.randint(10, 30, (16,), dtype=torch.long)\\n        >>> loss = F.ctc_loss(log_probs, targets, input_lengths, target_lengths)\\n        >>> loss.backward()\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(log_probs, targets, input_lengths, target_lengths):\\n        return handle_torch_function(\\n            ctc_loss,\\n            (log_probs, targets, input_lengths, target_lengths),\\n            log_probs,\\n            targets,\\n            input_lengths,\\n            target_lengths,\\n            blank=blank,\\n            reduction=reduction,\\n            zero_infinity=zero_infinity,\\n        )\\n    return torch.ctc_loss(\\n        log_probs,\\n        targets,\\n        input_lengths,\\n        target_lengths,\\n        blank,\\n        _Reduction.get_enum(reduction),\\n        zero_infinity,\\n    )\\n\\n\\nif ctc_loss.__doc__:\\n    ctc_loss.__doc__ = ctc_loss.__doc__.format(**reproducibility_notes)\\n\\n\\ndef nll_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    weight: Optional[Tensor] = None,\\n    size_average: Optional[bool] = None,\\n    ignore_index: int = -100,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the negative log likelihood loss.\\n\\n    See :class:`~torch.nn.NLLLoss` for details.\\n\\n    Args:\\n        input: :math:`(N, C)` where `C = number of classes` or :math:`(N, C, H, W)`\\n            in case of 2D Loss, or :math:`(N, C, d_1, d_2, ..., d_K)` where :math:`K \\\\geq 1`\\n            in the case of K-dimensional loss. `input` is expected to be log-probabilities.\\n        target: :math:`(N)` where each value is :math:`0 \\\\leq \\\\text{targets}[i] \\\\leq C-1`,\\n            or :math:`(N, d_1, d_2, ..., d_K)` where :math:`K \\\\geq 1` for\\n            K-dimensional loss.\\n        weight (Tensor, optional): a manual rescaling weight given to each\\n            class. If given, has to be a Tensor of size `C`\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when reduce is ``False``. Default: ``True``\\n        ignore_index (int, optional): Specifies a target value that is ignored\\n            and does not contribute to the input gradient. When :attr:`size_average` is\\n            ``True``, the loss is averaged over non-ignored targets. Default: -100\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Example::\\n\\n        >>> # input is of size N x C = 3 x 5\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> # each element in target has to have 0 <= value < C\\n        >>> target = torch.tensor([1, 0, 4])\\n        >>> output = F.nll_loss(F.log_softmax(input, dim=1), target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, weight):\\n        return handle_torch_function(\\n            nll_loss,\\n            (input, target, weight),\\n            input,\\n            target,\\n            weight=weight,\\n            size_average=size_average,\\n            ignore_index=ignore_index,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n    return torch._C._nn.nll_loss_nd(\\n        input, target, weight, _Reduction.get_enum(reduction), ignore_index\\n    )\\n\\n\\ndef poisson_nll_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    log_input: bool = True,\\n    full: bool = False,\\n    size_average: Optional[bool] = None,\\n    eps: float = 1e-8,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Poisson negative log likelihood loss.\\n\\n    See :class:`~torch.nn.PoissonNLLLoss` for details.\\n\\n    Args:\\n        input: expectation of underlying Poisson distribution.\\n        target: random sample :math:`target \\\\sim \\\\text{Poisson}(input)`.\\n        log_input: if ``True`` the loss is computed as\\n            :math:`\\\\exp(\\\\text{input}) - \\\\text{target} * \\\\text{input}`, if ``False`` then loss is\\n            :math:`\\\\text{input} - \\\\text{target} * \\\\log(\\\\text{input}+\\\\text{eps})`. Default: ``True``\\n        full: whether to compute full loss, i. e. to add the Stirling\\n            approximation term. Default: ``False``\\n            :math:`\\\\text{target} * \\\\log(\\\\text{target}) - \\\\text{target} + 0.5 * \\\\log(2 * \\\\pi * \\\\text{target})`.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when reduce is ``False``. Default: ``True``\\n        eps (float, optional): Small value to avoid evaluation of :math:`\\\\log(0)` when\\n            :attr:`log_input`\\\\ =\\\\ ``False``. Default: 1e-8\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            poisson_nll_loss,\\n            (input, target),\\n            input,\\n            target,\\n            log_input=log_input,\\n            full=full,\\n            size_average=size_average,\\n            eps=eps,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n    if reduction != \\\"none\\\" and reduction != \\\"mean\\\" and reduction != \\\"sum\\\":\\n        ret = input\\n        raise ValueError(reduction + \\\" is not a valid value for reduction\\\")\\n\\n    ret = torch.poisson_nll_loss(\\n        input, target, log_input, full, eps, _Reduction.get_enum(reduction)\\n    )\\n    return ret\\n\\n\\ndef gaussian_nll_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    var: Tensor,\\n    full: bool = False,\\n    eps: float = 1e-6,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Gaussian negative log likelihood loss.\\n\\n    See :class:`~torch.nn.GaussianNLLLoss` for details.\\n\\n    Args:\\n        input: expectation of the Gaussian distribution.\\n        target: sample from the Gaussian distribution.\\n        var: tensor of positive variance(s), one for each of the expectations\\n            in the input (heteroscedastic), or a single one (homoscedastic).\\n        full (bool, optional): include the constant term in the loss calculation. Default: ``False``.\\n        eps (float, optional): value added to var, for stability. Default: 1e-6.\\n        reduction (str, optional): specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the output is the average of all batch member losses,\\n            ``'sum'``: the output is the sum of all batch member losses.\\n            Default: ``'mean'``.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, var):\\n        return handle_torch_function(\\n            gaussian_nll_loss,\\n            (input, target, var),\\n            input,\\n            target,\\n            var,\\n            full=full,\\n            eps=eps,\\n            reduction=reduction,\\n        )\\n\\n    # Check var size\\n    # If var.size == input.size, the case is heteroscedastic and no further checks are needed.\\n    # Otherwise:\\n    if var.size() != input.size():\\n        # If var is one dimension short of input, but the sizes match otherwise, then this is a homoscedastic case.\\n        # e.g. input.size = (10, 2, 3), var.size = (10, 2)\\n        # -> unsqueeze var so that var.shape = (10, 2, 1)\\n        # this is done so that broadcasting can happen in the loss calculation\\n        if input.size()[:-1] == var.size():\\n            var = torch.unsqueeze(var, -1)\\n\\n        # This checks if the sizes match up to the final dimension, and the final dimension of var is of size 1.\\n        # This is also a homoscedastic case.\\n        # e.g. input.size = (10, 2, 3), var.size = (10, 2, 1)\\n        elif (\\n            input.size()[:-1] == var.size()[:-1] and var.size(-1) == 1\\n        ):  # Heteroscedastic case\\n            pass\\n\\n        # If none of the above pass, then the size of var is incorrect.\\n        else:\\n            raise ValueError(\\\"var is of incorrect size\\\")\\n\\n    # Check validity of reduction mode\\n    if reduction != \\\"none\\\" and reduction != \\\"mean\\\" and reduction != \\\"sum\\\":\\n        raise ValueError(reduction + \\\" is not valid\\\")\\n\\n    # Entries of var must be non-negative\\n    if torch.any(var < 0):\\n        raise ValueError(\\\"var has negative entry/entries\\\")\\n\\n    # Clamp for stability\\n    var = var.clone()\\n    with torch.no_grad():\\n        var.clamp_(min=eps)\\n\\n    # Calculate the loss\\n    loss = 0.5 * (torch.log(var) + (input - target) ** 2 / var)\\n    if full:\\n        loss += 0.5 * math.log(2 * math.pi)\\n\\n    if reduction == \\\"mean\\\":\\n        return loss.mean()\\n    elif reduction == \\\"sum\\\":\\n        return loss.sum()\\n    else:\\n        return loss\\n\\n\\ndef kl_div(\\n    input: Tensor,\\n    target: Tensor,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n    log_target: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the KL Divergence loss.\\n\\n    Refer - The `Kullback-Leibler divergence Loss\\n    <https://en.wikipedia.org/wiki/Kullback-Leibler_divergence>`__\\n\\n    See :class:`~torch.nn.KLDivLoss` for details.\\n\\n    Args:\\n        input: Tensor of arbitrary shape in log-probabilities.\\n        target: Tensor of the same shape as input. See :attr:`log_target` for\\n            the target's interpretation.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when reduce is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'batchmean'`` | ``'sum'`` | ``'mean'``.\\n            ``'none'``: no reduction will be applied\\n            ``'batchmean'``: the sum of the output will be divided by the batchsize\\n            ``'sum'``: the output will be summed\\n            ``'mean'``: the output will be divided by the number of elements in the output\\n            Default: ``'mean'``\\n        log_target (bool): A flag indicating whether ``target`` is passed in the log space.\\n            It is recommended to pass certain distributions (like ``softmax``)\\n            in the log space to avoid numerical issues caused by explicit ``log``.\\n            Default: ``False``\\n\\n    .. note::\\n        :attr:`size_average` and :attr:`reduce` are in the process of being deprecated,\\n        and in the meantime, specifying either of those two args will override :attr:`reduction`.\\n\\n    .. warning::\\n        :attr:`reduction` = ``'mean'`` doesn't return the true kl divergence value, please use\\n        :attr:`reduction` = ``'batchmean'`` which aligns with KL math definition.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            kl_div,\\n            (input, target),\\n            input,\\n            target,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n            log_target=log_target,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        if reduction == \\\"mean\\\":\\n            warnings.warn(\\n                \\\"reduction: 'mean' divides the total loss by both the batch size and the support size.\\\"\\n                \\\"'batchmean' divides only by the batch size, and aligns with the KL div math definition.\\\"\\n                \\\"'mean' will be changed to behave the same as 'batchmean' in the next major release.\\\"\\n            )\\n\\n        # special case for batchmean\\n        if reduction == \\\"batchmean\\\":\\n            reduction_enum = _Reduction.get_enum(\\\"sum\\\")\\n        else:\\n            reduction_enum = _Reduction.get_enum(reduction)\\n\\n    reduced = torch.kl_div(input, target, reduction_enum, log_target=log_target)\\n\\n    if reduction == \\\"batchmean\\\" and input.dim() != 0:\\n        reduced = reduced / input.size()[0]\\n\\n    return reduced\\n\\n\\ndef cross_entropy(\\n    input: Tensor,\\n    target: Tensor,\\n    weight: Optional[Tensor] = None,\\n    size_average: Optional[bool] = None,\\n    ignore_index: int = -100,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n    label_smoothing: float = 0.0,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the cross entropy loss between input logits and target.\\n\\n    See :class:`~torch.nn.CrossEntropyLoss` for details.\\n\\n    Args:\\n        input (Tensor) : Predicted unnormalized logits;\\n            see Shape section below for supported shapes.\\n        target (Tensor) : Ground truth class indices or class probabilities;\\n            see Shape section below for supported shapes.\\n        weight (Tensor, optional): a manual rescaling weight given to each\\n            class. If given, has to be a Tensor of size `C`\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when reduce is ``False``. Default: ``True``\\n        ignore_index (int, optional): Specifies a target value that is ignored\\n            and does not contribute to the input gradient. When :attr:`size_average` is\\n            ``True``, the loss is averaged over non-ignored targets. Note that\\n            :attr:`ignore_index` is only applicable when the target contains class indices.\\n            Default: -100\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n        label_smoothing (float, optional): A float in [0.0, 1.0]. Specifies the amount\\n            of smoothing when computing the loss, where 0.0 means no smoothing. The targets\\n            become a mixture of the original ground truth and a uniform distribution as described in\\n            `Rethinking the Inception Architecture for Computer Vision <https://arxiv.org/abs/1512.00567>`__. Default: :math:`0.0`.\\n\\n    Shape:\\n        - Input: Shape :math:`(C)`, :math:`(N, C)` or :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1`\\n          in the case of `K`-dimensional loss.\\n        - Target: If containing class indices, shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with\\n          :math:`K \\\\geq 1` in the case of K-dimensional loss where each value should be between :math:`[0, C)`.\\n          If containing class probabilities, same shape as the input and each value should be between :math:`[0, 1]`.\\n\\n        where:\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                C ={} & \\\\text{number of classes} \\\\\\\\\\n                N ={} & \\\\text{batch size} \\\\\\\\\\n            \\\\end{aligned}\\n\\n    Examples::\\n\\n        >>> # Example of target with class indices\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.randint(5, (3,), dtype=torch.int64)\\n        >>> loss = F.cross_entropy(input, target)\\n        >>> loss.backward()\\n        >>>\\n        >>> # Example of target with class probabilities\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.randn(3, 5).softmax(dim=1)\\n        >>> loss = F.cross_entropy(input, target)\\n        >>> loss.backward()\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, weight):\\n        return handle_torch_function(\\n            cross_entropy,\\n            (input, target, weight),\\n            input,\\n            target,\\n            weight=weight,\\n            size_average=size_average,\\n            ignore_index=ignore_index,\\n            reduce=reduce,\\n            reduction=reduction,\\n            label_smoothing=label_smoothing,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n    return torch._C._nn.cross_entropy_loss(\\n        input,\\n        target,\\n        weight,\\n        _Reduction.get_enum(reduction),\\n        ignore_index,\\n        label_smoothing,\\n    )\\n\\n\\ndef binary_cross_entropy(\\n    input: Tensor,\\n    target: Tensor,\\n    weight: Optional[Tensor] = None,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Measure Binary Cross Entropy between the target and input probabilities.\\n\\n    See :class:`~torch.nn.BCELoss` for details.\\n\\n    Args:\\n        input: Tensor of arbitrary shape as probabilities.\\n        target: Tensor of the same shape as input with values between 0 and 1.\\n        weight (Tensor, optional): a manual rescaling weight\\n                if provided it's repeated to match input tensor shape\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when reduce is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Examples::\\n\\n        >>> input = torch.randn(3, 2, requires_grad=True)\\n        >>> target = torch.rand(3, 2, requires_grad=False)\\n        >>> loss = F.binary_cross_entropy(torch.sigmoid(input), target)\\n        >>> loss.backward()\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, weight):\\n        return handle_torch_function(\\n            binary_cross_entropy,\\n            (input, target, weight),\\n            input,\\n            target,\\n            weight=weight,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    if target.size() != input.size():\\n        raise ValueError(\\n            f\\\"Using a target size ({target.size()}) that is different to the input size ({input.size()}) is deprecated. \\\"\\n            \\\"Please ensure they have the same size.\\\"\\n        )\\n\\n    if weight is not None:\\n        new_size = _infer_size(target.size(), weight.size())\\n        weight = weight.expand(new_size)\\n\\n    return torch._C._nn.binary_cross_entropy(input, target, weight, reduction_enum)\\n\\n\\ndef binary_cross_entropy_with_logits(\\n    input: Tensor,\\n    target: Tensor,\\n    weight: Optional[Tensor] = None,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n    pos_weight: Optional[Tensor] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Calculate Binary Cross Entropy between target and input logits.\\n\\n    See :class:`~torch.nn.BCEWithLogitsLoss` for details.\\n\\n    Args:\\n        input: Tensor of arbitrary shape as unnormalized scores (often referred to as logits).\\n        target: Tensor of the same shape as input with values between 0 and 1\\n        weight (Tensor, optional): a manual rescaling weight\\n            if provided it's repeated to match input tensor shape\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when reduce is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n        pos_weight (Tensor, optional): a weight of positive examples to be broadcasted with target.\\n            Must be a tensor with equal size along the class dimension to the number of classes.\\n            Pay close attention to PyTorch's broadcasting semantics in order to achieve the desired\\n            operations. For a target of size [B, C, H, W] (where B is batch size) pos_weight of\\n            size [B, C, H, W] will apply different pos_weights to each element of the batch or\\n            [C, H, W] the same pos_weights across the batch. To apply the same positive weight\\n            along all spatial dimensions for a 2D multi-class target [C, H, W] use: [C, 1, 1].\\n            Default: ``None``\\n\\n    Examples::\\n\\n         >>> input = torch.randn(3, requires_grad=True)\\n         >>> target = torch.empty(3).random_(2)\\n         >>> loss = F.binary_cross_entropy_with_logits(input, target)\\n         >>> loss.backward()\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, weight, pos_weight):\\n        return handle_torch_function(\\n            binary_cross_entropy_with_logits,\\n            (input, target, weight, pos_weight),\\n            input,\\n            target,\\n            weight=weight,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n            pos_weight=pos_weight,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n\\n    if not (target.size() == input.size()):\\n        raise ValueError(\\n            f\\\"Target size ({target.size()}) must be the same as input size ({input.size()})\\\"\\n        )\\n\\n    return torch.binary_cross_entropy_with_logits(\\n        input, target, weight, pos_weight, reduction_enum\\n    )\\n\\n\\ndef smooth_l1_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n    beta: float = 1.0,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the Smooth L1 loss.\\n\\n    Function uses a squared term if the absolute\\n    element-wise error falls below beta and an L1 term otherwise.\\n\\n    See :class:`~torch.nn.SmoothL1Loss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            smooth_l1_loss,\\n            (input, target),\\n            input,\\n            target,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n            beta=beta,\\n        )\\n    if not (target.size() == input.size()):\\n        warnings.warn(\\n            f\\\"Using a target size ({target.size()}) that is different to the input size ({input.size()}). \\\"\\n            \\\"This will likely lead to incorrect results due to broadcasting. \\\"\\n            \\\"Please ensure they have the same size.\\\",\\n            stacklevel=2,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n\\n    expanded_input, expanded_target = torch.broadcast_tensors(input, target)\\n\\n    if beta == 0.0:\\n        return torch._C._nn.l1_loss(\\n            expanded_input, expanded_target, _Reduction.get_enum(reduction)\\n        )\\n    else:\\n        return torch._C._nn.smooth_l1_loss(\\n            expanded_input, expanded_target, _Reduction.get_enum(reduction), beta\\n        )\\n\\n\\ndef huber_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    reduction: str = \\\"mean\\\",\\n    delta: float = 1.0,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the Huber loss.\\n\\n    Function uses a squared term if the absolute\\n    element-wise error falls below delta and a delta-scaled L1 term otherwise.\\n\\n    When delta equals 1, this loss is equivalent to SmoothL1Loss.\\n    In general, Huber loss differs from SmoothL1Loss by a factor of delta (AKA beta in Smooth L1).\\n\\n    See :class:`~torch.nn.HuberLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            huber_loss,\\n            (input, target),\\n            input,\\n            target,\\n            reduction=reduction,\\n            delta=delta,\\n        )\\n    if not (target.size() == input.size()):\\n        warnings.warn(\\n            f\\\"Using a target size ({target.size()}) that is different to the input size ({input.size()}). \\\"\\n            \\\"This will likely lead to incorrect results due to broadcasting. \\\"\\n            \\\"Please ensure they have the same size.\\\",\\n            stacklevel=2,\\n        )\\n\\n    expanded_input, expanded_target = torch.broadcast_tensors(input, target)\\n    return torch._C._nn.huber_loss(\\n        expanded_input, expanded_target, _Reduction.get_enum(reduction), delta\\n    )\\n\\n\\ndef l1_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"l1_loss(input, target, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    Function that takes the mean element-wise absolute value difference.\\n\\n    See :class:`~torch.nn.L1Loss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            l1_loss,\\n            (input, target),\\n            input,\\n            target,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if not (target.size() == input.size()):\\n        warnings.warn(\\n            f\\\"Using a target size ({target.size()}) that is different to the input size ({input.size()}). \\\"\\n            \\\"This will likely lead to incorrect results due to broadcasting. \\\"\\n            \\\"Please ensure they have the same size.\\\",\\n            stacklevel=2,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n\\n    expanded_input, expanded_target = torch.broadcast_tensors(input, target)\\n    return torch._C._nn.l1_loss(\\n        expanded_input, expanded_target, _Reduction.get_enum(reduction)\\n    )\\n\\n\\ndef mse_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"mse_loss(input, target, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    Measures the element-wise mean squared error.\\n    See :class:`~torch.nn.MSELoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            mse_loss,\\n            (input, target),\\n            input,\\n            target,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if not (target.size() == input.size()):\\n        warnings.warn(\\n            f\\\"Using a target size ({target.size()}) that is different to the input size ({input.size()}). \\\"\\n            \\\"This will likely lead to incorrect results due to broadcasting. \\\"\\n            \\\"Please ensure they have the same size.\\\",\\n            stacklevel=2,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n\\n    expanded_input, expanded_target = torch.broadcast_tensors(input, target)\\n    return torch._C._nn.mse_loss(\\n        expanded_input, expanded_target, _Reduction.get_enum(reduction)\\n    )\\n\\n\\ndef margin_ranking_loss(\\n    input1: Tensor,\\n    input2: Tensor,\\n    target: Tensor,\\n    margin: float = 0,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"margin_ranking_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.MarginRankingLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input1, input2, target):\\n        return handle_torch_function(\\n            margin_ranking_loss,\\n            (input1, input2, target),\\n            input1,\\n            input2,\\n            target,\\n            margin=margin,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    if input1.dim() != input2.dim() or input1.dim() != target.dim():\\n        raise RuntimeError(\\n            f\\\"margin_ranking_loss : All input tensors should have same dimension but got sizes: \\\"\\n            f\\\"input1: {input1.size()}, input2: {input2.size()}, target: {target.size()} \\\"\\n        )\\n    return torch.margin_ranking_loss(input1, input2, target, margin, reduction_enum)\\n\\n\\ndef hinge_embedding_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    margin: float = 1.0,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"hinge_embedding_loss(input, target, margin=1.0, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.HingeEmbeddingLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            hinge_embedding_loss,\\n            (input, target),\\n            input,\\n            target,\\n            margin=margin,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    return torch.hinge_embedding_loss(input, target, margin, reduction_enum)\\n\\n\\ndef multilabel_margin_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"multilabel_margin_loss(input, target, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.MultiLabelMarginLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            multilabel_margin_loss,\\n            (input, target),\\n            input,\\n            target,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    return torch._C._nn.multilabel_margin_loss(input, target, reduction_enum)\\n\\n\\ndef soft_margin_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"\\n    soft_margin_loss(input, target, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.SoftMarginLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target):\\n        return handle_torch_function(\\n            soft_margin_loss,\\n            (input, target),\\n            input,\\n            target,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    return torch._C._nn.soft_margin_loss(input, target, reduction_enum)\\n\\n\\ndef multilabel_soft_margin_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    weight: Optional[Tensor] = None,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"multilabel_soft_margin_loss(input, target, weight=None, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.MultiLabelSoftMarginLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, weight):\\n        return handle_torch_function(\\n            multilabel_soft_margin_loss,\\n            (input, target, weight),\\n            input,\\n            target,\\n            weight=weight,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction = _Reduction.legacy_get_string(size_average, reduce)\\n\\n    loss = -(target * logsigmoid(input) + (1 - target) * logsigmoid(-input))\\n\\n    if weight is not None:\\n        loss = loss * weight\\n\\n    class_dim = input.dim() - 1\\n    C = input.size(class_dim)\\n    loss = loss.sum(dim=class_dim) / C  # only return N loss values\\n\\n    if reduction == \\\"none\\\":\\n        ret = loss\\n    elif reduction == \\\"mean\\\":\\n        ret = loss.mean()\\n    elif reduction == \\\"sum\\\":\\n        ret = loss.sum()\\n    else:\\n        ret = input\\n        raise ValueError(reduction + \\\" is not valid\\\")\\n    return ret\\n\\n\\ndef cosine_embedding_loss(\\n    input1: Tensor,\\n    input2: Tensor,\\n    target: Tensor,\\n    margin: float = 0,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"cosine_embedding_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.CosineEmbeddingLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input1, input2, target):\\n        return handle_torch_function(\\n            cosine_embedding_loss,\\n            (input1, input2, target),\\n            input1,\\n            input2,\\n            target,\\n            margin=margin,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    return torch.cosine_embedding_loss(input1, input2, target, margin, reduction_enum)\\n\\n\\ndef multi_margin_loss(\\n    input: Tensor,\\n    target: Tensor,\\n    p: int = 1,\\n    margin: float = 1.0,\\n    weight: Optional[Tensor] = None,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:  # noqa: D400,D402\\n    r\\\"\\\"\\\"multi_margin_loss(input, target, p=1, margin=1, weight=None, size_average=None, reduce=None, reduction='mean') -> Tensor\\n\\n    See :class:`~torch.nn.MultiMarginLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, target, weight):\\n        return handle_torch_function(\\n            multi_margin_loss,\\n            (input, target, weight),\\n            input,\\n            target,\\n            p=p,\\n            margin=margin,\\n            weight=weight,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    if p != 1 and p != 2:\\n        raise ValueError(\\\"only p == 1 and p == 2 supported\\\")\\n    if weight is not None:\\n        if weight.dim() != 1:\\n            raise ValueError(\\\"weight must be one-dimensional\\\")\\n\\n    return torch._C._nn.multi_margin_loss(\\n        input, target, p, margin, weight, reduction_enum\\n    )\\n\\n\\npixel_shuffle = _add_docstr(\\n    torch.pixel_shuffle,\\n    r\\\"\\\"\\\"\\npixel_shuffle(input, upscale_factor) -> Tensor\\n\\nRearranges elements in a tensor of shape :math:`(*, C \\\\times r^2, H, W)` to a\\ntensor of shape :math:`(*, C, H \\\\times r, W \\\\times r)`, where r is the :attr:`upscale_factor`.\\n\\nSee :class:`~torch.nn.PixelShuffle` for details.\\n\\nArgs:\\n    input (Tensor): the input tensor\\n    upscale_factor (int): factor to increase spatial resolution by\\n\\nExamples::\\n\\n    >>> input = torch.randn(1, 9, 4, 4)\\n    >>> output = torch.nn.functional.pixel_shuffle(input, 3)\\n    >>> print(output.size())\\n    torch.Size([1, 1, 12, 12])\\n\\\"\\\"\\\",\\n)\\n\\npixel_unshuffle = _add_docstr(\\n    torch.pixel_unshuffle,\\n    r\\\"\\\"\\\"\\npixel_unshuffle(input, downscale_factor) -> Tensor\\n\\nReverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a\\ntensor of shape :math:`(*, C, H \\\\times r, W \\\\times r)` to a tensor of shape\\n:math:`(*, C \\\\times r^2, H, W)`, where r is the :attr:`downscale_factor`.\\n\\nSee :class:`~torch.nn.PixelUnshuffle` for details.\\n\\nArgs:\\n    input (Tensor): the input tensor\\n    downscale_factor (int): factor to increase spatial resolution by\\n\\nExamples::\\n\\n    >>> input = torch.randn(1, 1, 12, 12)\\n    >>> output = torch.nn.functional.pixel_unshuffle(input, 3)\\n    >>> print(output.size())\\n    torch.Size([1, 9, 4, 4])\\n\\\"\\\"\\\",\\n)\\n\\nchannel_shuffle = _add_docstr(\\n    torch.channel_shuffle,\\n    r\\\"\\\"\\\"\\nchannel_shuffle(input, groups) -> Tensor\\n\\nDivide the channels in a tensor of shape :math:`(*, C , H, W)`\\ninto g groups and rearrange them as :math:`(*, C \\\\frac g, g, H, W)`,\\nwhile keeping the original tensor shape.\\n\\nSee :class:`~torch.nn.ChannelShuffle` for details.\\n\\nArgs:\\n    input (Tensor): the input tensor\\n    groups (int): number of groups to divide channels in and rearrange.\\n\\nExamples::\\n\\n    >>> input = torch.randn(1, 4, 2, 2)\\n    >>> print(input)\\n    [[[[1, 2],\\n       [3, 4]],\\n      [[5, 6],\\n       [7, 8]],\\n      [[9, 10],\\n       [11, 12]],\\n      [[13, 14],\\n       [15, 16]],\\n     ]]\\n    >>> output = torch.nn.functional.channel_shuffle(input, 2)\\n    >>> print(output)\\n    [[[[1, 2],\\n       [3, 4]],\\n      [[9, 10],\\n       [11, 12]],\\n      [[5, 6],\\n       [7, 8]],\\n      [[13, 14],\\n       [15, 16]],\\n     ]]\\n\\\"\\\"\\\",\\n)\\n\\nnative_channel_shuffle = _add_docstr(\\n    torch.native_channel_shuffle,\\n    r\\\"\\\"\\\"\\nnative_channel_shuffle(input, groups) -> Tensor\\n\\nNative kernel level implementation of the `channel_shuffle`.\\nThis function might become private in future releases, use with caution.\\n\\nDivide the channels in a tensor of shape :math:`(*, C , H, W)`\\ninto g groups and rearrange them as :math:`(*, C \\\\frac g, g, H, W)`,\\nwhile keeping the original tensor shape.\\n\\nSee :class:`~torch.nn.ChannelShuffle` for details.\\n\\nArgs:\\n    input (Tensor): the input tensor\\n    groups (int): number of groups to divide channels in and rearrange.\\n\\nExamples::\\n\\n    >>> input = torch.randn(1, 4, 2, 2)\\n    >>> print(input)\\n    [[[[1, 2],\\n       [3, 4]],\\n      [[5, 6],\\n       [7, 8]],\\n      [[9, 10],\\n       [11, 12]],\\n      [[13, 14],\\n       [15, 16]],\\n     ]]\\n    >>> output = torch.nn.functional.native_channel_shuffle(input, 2)\\n    >>> print(output)\\n    [[[[1, 2],\\n       [3, 4]],\\n      [[9, 10],\\n       [11, 12]],\\n      [[5, 6],\\n       [7, 8]],\\n      [[13, 14],\\n       [15, 16]],\\n     ]]\\n\\\"\\\"\\\",\\n)\\n\\n\\n@_overload\\ndef upsample(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[float] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n) -> Tensor:  # noqa: B950\\n    pass\\n\\n\\n@_overload\\ndef upsample(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[List[int]] = None,\\n    scale_factor: Optional[float] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n) -> Tensor:  # noqa: B950\\n    pass\\n\\n\\ndef upsample(  # noqa: F811\\n    input,\\n    size=None,\\n    scale_factor=None,\\n    mode=\\\"nearest\\\",\\n    align_corners=None,\\n):\\n    r\\\"\\\"\\\"Upsample input.\\n\\n    Provided tensor is upsampled to either the given :attr:`size` or the given\\n    :attr:`scale_factor`\\n\\n    .. warning::\\n        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.\\n        This is equivalent with ``nn.functional.interpolate(...)``.\\n\\n    Note:\\n        {backward_reproducibility_note}\\n\\n    The algorithm used for upsampling is determined by :attr:`mode`.\\n\\n    Currently temporal, spatial and volumetric upsampling are supported, i.e.\\n    expected inputs are 3-D, 4-D or 5-D in shape.\\n\\n    The input dimensions are interpreted in the form:\\n    `mini-batch x channels x [optional depth] x [optional height] x width`.\\n\\n    The modes available for upsampling are: `nearest`, `linear` (3D-only),\\n    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only)\\n\\n    Args:\\n        input (Tensor): the input tensor\\n        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):\\n            output spatial size.\\n        scale_factor (float or Tuple[float]): multiplier for spatial size. Has to match input size if it is a tuple.\\n        mode (str): algorithm used for upsampling:\\n            ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |\\n            ``'trilinear'``. Default: ``'nearest'``\\n        align_corners (bool, optional): Geometrically, we consider the pixels of the\\n            input and output as squares rather than points.\\n            If set to ``True``, the input and output tensors are aligned by the\\n            center points of their corner pixels, preserving the values at the corner pixels.\\n            If set to ``False``, the input and output tensors are aligned by the corner\\n            points of their corner pixels, and the interpolation uses edge value padding\\n            for out-of-boundary values, making this operation *independent* of input size\\n            when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`\\n            is ``'linear'``, ``'bilinear'``, ``'bicubic'`` or ``'trilinear'``.\\n            Default: ``False``\\n\\n    .. note::\\n        With ``mode='bicubic'``, it's possible to cause overshoot, in other words it can produce\\n        negative values or values greater than 255 for images.\\n        Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot\\n        when displaying the image.\\n\\n    .. warning::\\n        With ``align_corners = True``, the linearly interpolating modes\\n        (`linear`, `bilinear`, and `trilinear`) don't proportionally align the\\n        output and input pixels, and thus the output values can depend on the\\n        input size. This was the default behavior for these modes up to version\\n        0.3.1. Since then, the default behavior is ``align_corners = False``.\\n        See :class:`~torch.nn.Upsample` for concrete examples on how this\\n        affects the outputs.\\n\\n    \\\"\\\"\\\"\\n    warnings.warn(\\n        \\\"`nn.functional.upsample` is deprecated. \\\"\\n        \\\"Use `nn.functional.interpolate` instead.\\\",\\n        stacklevel=2,\\n    )\\n    return interpolate(input, size, scale_factor, mode, align_corners)\\n\\n\\nif upsample.__doc__:\\n    upsample.__doc__ = upsample.__doc__.format(**reproducibility_notes)\\n\\n\\ndef _is_integer(x) -> bool:\\n    r\\\"\\\"\\\"Type check the input number is an integer.\\n\\n    Will return True for int, SymInt, Numpy integers and Tensors with integer elements.\\n    \\\"\\\"\\\"\\n    if isinstance(x, (int, torch.SymInt)):\\n        return True\\n    if np is not None and isinstance(x, np.integer):\\n        return True\\n    return isinstance(x, Tensor) and not x.is_floating_point()\\n\\n\\n@_overload\\ndef interpolate(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[List[float]] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n    recompute_scale_factor: Optional[bool] = None,\\n    antialias: bool = False,\\n) -> Tensor:  # noqa: B950\\n    pass\\n\\n\\n@_overload\\ndef interpolate(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[List[int]] = None,\\n    scale_factor: Optional[List[float]] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n    recompute_scale_factor: Optional[bool] = None,\\n    antialias: bool = False,\\n) -> Tensor:  # noqa: B950\\n    pass\\n\\n\\n@_overload\\ndef interpolate(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[float] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n    recompute_scale_factor: Optional[bool] = None,\\n    antialias: bool = False,\\n) -> Tensor:  # noqa: B950\\n    pass\\n\\n\\n@_overload\\ndef interpolate(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[List[int]] = None,\\n    scale_factor: Optional[float] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n    recompute_scale_factor: Optional[bool] = None,\\n    antialias: bool = False,\\n) -> Tensor:\\n    pass\\n\\n\\ndef interpolate(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[List[float]] = None,\\n    mode: str = \\\"nearest\\\",\\n    align_corners: Optional[bool] = None,\\n    recompute_scale_factor: Optional[bool] = None,\\n    antialias: bool = False,\\n) -> Tensor:  # noqa: B950\\n    r\\\"\\\"\\\"Down/up samples the input.\\n\\n    Tensor interpolated to either the given :attr:`size` or the given\\n    :attr:`scale_factor`\\n\\n    The algorithm used for interpolation is determined by :attr:`mode`.\\n\\n    Currently temporal, spatial and volumetric sampling are supported, i.e.\\n    expected inputs are 3-D, 4-D or 5-D in shape.\\n\\n    The input dimensions are interpreted in the form:\\n    `mini-batch x channels x [optional depth] x [optional height] x width`.\\n\\n    The modes available for resizing are: `nearest`, `linear` (3D-only),\\n    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `area`, `nearest-exact`\\n\\n    Args:\\n        input (Tensor): the input tensor\\n        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):\\n            output spatial size.\\n        scale_factor (float or Tuple[float]): multiplier for spatial size. If `scale_factor` is a tuple,\\n            its length has to match the number of spatial dimensions; `input.dim() - 2`.\\n        mode (str): algorithm used for upsampling:\\n            ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |\\n            ``'trilinear'`` | ``'area'`` | ``'nearest-exact'``. Default: ``'nearest'``\\n        align_corners (bool, optional): Geometrically, we consider the pixels of the\\n            input and output as squares rather than points.\\n            If set to ``True``, the input and output tensors are aligned by the\\n            center points of their corner pixels, preserving the values at the corner pixels.\\n            If set to ``False``, the input and output tensors are aligned by the corner\\n            points of their corner pixels, and the interpolation uses edge value padding\\n            for out-of-boundary values, making this operation *independent* of input size\\n            when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`\\n            is ``'linear'``, ``'bilinear'``, ``'bicubic'`` or ``'trilinear'``.\\n            Default: ``False``\\n        recompute_scale_factor (bool, optional): recompute the scale_factor for use in the\\n            interpolation calculation. If `recompute_scale_factor` is ``True``, then\\n            `scale_factor` must be passed in and `scale_factor` is used to compute the\\n            output `size`. The computed output `size` will be used to infer new scales for\\n            the interpolation. Note that when `scale_factor` is floating-point, it may differ\\n            from the recomputed `scale_factor` due to rounding and precision issues.\\n            If `recompute_scale_factor` is ``False``, then `size` or `scale_factor` will\\n            be used directly for interpolation. Default: ``None``.\\n        antialias (bool, optional): flag to apply anti-aliasing. Default: ``False``. Using anti-alias\\n            option together with ``align_corners=False``, interpolation result would match Pillow\\n            result for downsampling operation. Supported modes: ``'bilinear'``, ``'bicubic'``.\\n\\n    .. note::\\n        With ``mode='bicubic'``, it's possible to cause overshoot, in other words it can produce\\n        negative values or values greater than 255 for images.\\n        Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot\\n        when displaying the image.\\n\\n    .. note::\\n        Mode ``mode='nearest-exact'`` matches Scikit-Image and PIL nearest neighbours interpolation\\n        algorithms and fixes known issues with ``mode='nearest'``. This mode is introduced to keep\\n        backward compatibility.\\n        Mode ``mode='nearest'`` matches buggy OpenCV's ``INTER_NEAREST`` interpolation algorithm.\\n\\n    .. note::\\n        The gradients for the dtype ``float16`` on CUDA may be inaccurate in the upsample operation\\n        when using modes ``['linear', 'bilinear', 'bicubic', 'trilinear', 'area']``.\\n        For more details, please refer to the discussion in\\n        `issue#104157 <https://github.com/pytorch/pytorch/issues/104157>`_.\\n\\n    Note:\\n        {backward_reproducibility_note}\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            interpolate,\\n            (input,),\\n            input,\\n            size=size,\\n            scale_factor=scale_factor,\\n            mode=mode,\\n            align_corners=align_corners,\\n            recompute_scale_factor=recompute_scale_factor,\\n            antialias=antialias,\\n        )\\n\\n    if mode in (\\\"nearest\\\", \\\"area\\\", \\\"nearest-exact\\\"):\\n        if align_corners is not None:\\n            raise ValueError(\\n                \\\"align_corners option can only be set with the \\\"\\n                \\\"interpolating modes: linear | bilinear | bicubic | trilinear\\\"\\n            )\\n    else:\\n        if align_corners is None:\\n            align_corners = False\\n\\n    dim = input.dim() - 2  # Number of spatial dimensions.\\n\\n    # Process size and scale_factor.  Validate that exactly one is set.\\n    # Validate its length if it is a list, or expand it if it is a scalar.\\n    # After this block, exactly one of output_size and scale_factors will\\n    # be non-None, and it will be a list (or tuple).\\n    if size is not None and scale_factor is not None:\\n        raise ValueError(\\\"only one of size or scale_factor should be defined\\\")\\n    elif size is not None:\\n        assert scale_factor is None\\n        scale_factors = None\\n        if isinstance(size, (list, tuple)):\\n            if len(size) != dim:\\n                raise ValueError(\\n                    \\\"Input and output must have the same number of spatial dimensions, but got \\\"\\n                    f\\\"input with spatial dimensions of {list(input.shape[2:])} and output size of {size}. \\\"\\n                    \\\"Please provide input tensor in (N, C, d1, d2, ...,dK) format and \\\"\\n                    \\\"output size in (o1, o2, ...,oK) format.\\\"\\n                )\\n            if not torch.jit.is_scripting():\\n                if not all(_is_integer(x) for x in size):\\n                    raise TypeError(\\n                        \\\"expected size to be one of int or Tuple[int] or Tuple[int, int] or \\\"\\n                        f\\\"Tuple[int, int, int], but got size with types {[type(x) for x in size]}\\\"\\n                    )\\n            output_size = size\\n        else:\\n            output_size = [size for _ in range(dim)]\\n    elif scale_factor is not None:\\n        assert size is None\\n        output_size = None\\n        if isinstance(scale_factor, (list, tuple)):\\n            if len(scale_factor) != dim:\\n                raise ValueError(\\n                    \\\"Input and scale_factor must have the same number of spatial dimensions, but \\\"\\n                    f\\\"got input with spatial dimensions of {list(input.shape[2:])} and \\\"\\n                    f\\\"scale_factor of shape {scale_factor}. \\\"\\n                    \\\"Please provide input tensor in (N, C, d1, d2, ...,dK) format and \\\"\\n                    \\\"scale_factor in (s1, s2, ...,sK) format.\\\"\\n                )\\n            scale_factors = scale_factor\\n        else:\\n            scale_factors = [scale_factor for _ in range(dim)]\\n    else:\\n        raise ValueError(\\\"either size or scale_factor should be defined\\\")\\n\\n    if (\\n        recompute_scale_factor is not None\\n        and recompute_scale_factor\\n        and size is not None\\n    ):\\n        raise ValueError(\\n            \\\"recompute_scale_factor is not meaningful with an explicit size.\\\"\\n        )\\n\\n    # \\\"area\\\" mode always requires an explicit size rather than scale factor.\\n    # Re-use the recompute_scale_factor code path.\\n    if mode == \\\"area\\\" and output_size is None:\\n        recompute_scale_factor = True\\n\\n    if recompute_scale_factor is not None and recompute_scale_factor:\\n        # We compute output_size here, then un-set scale_factors.\\n        # The C++ code will recompute it based on the (integer) output size.\\n        assert scale_factors is not None\\n        if not torch.jit.is_scripting() and torch._C._get_tracing_state():\\n            # make scale_factor a tensor in tracing so constant doesn't get baked in\\n            output_size = [\\n                (\\n                    torch.floor(\\n                        (\\n                            input.size(i + 2).float()\\n                            * torch.tensor(scale_factors[i], dtype=torch.float32)\\n                        ).float()\\n                    )\\n                )\\n                for i in range(dim)\\n            ]\\n        elif torch.jit.is_scripting():\\n            output_size = [\\n                int(math.floor(float(input.size(i + 2)) * scale_factors[i]))\\n                for i in range(dim)\\n            ]\\n        else:\\n            output_size = [\\n                _sym_int(input.size(i + 2) * scale_factors[i]) for i in range(dim)\\n            ]\\n        scale_factors = None\\n\\n    if antialias and not (mode in (\\\"bilinear\\\", \\\"bicubic\\\") and input.ndim == 4):\\n        raise ValueError(\\n            \\\"Anti-alias option is restricted to bilinear and bicubic modes and requires a 4-D tensor as input\\\"\\n        )\\n\\n    if input.dim() == 3 and mode == \\\"nearest\\\":\\n        return torch._C._nn.upsample_nearest1d(input, output_size, scale_factors)\\n    if input.dim() == 4 and mode == \\\"nearest\\\":\\n        return torch._C._nn.upsample_nearest2d(input, output_size, scale_factors)\\n    if input.dim() == 5 and mode == \\\"nearest\\\":\\n        return torch._C._nn.upsample_nearest3d(input, output_size, scale_factors)\\n\\n    if input.dim() == 3 and mode == \\\"nearest-exact\\\":\\n        return torch._C._nn._upsample_nearest_exact1d(input, output_size, scale_factors)\\n    if input.dim() == 4 and mode == \\\"nearest-exact\\\":\\n        return torch._C._nn._upsample_nearest_exact2d(input, output_size, scale_factors)\\n    if input.dim() == 5 and mode == \\\"nearest-exact\\\":\\n        return torch._C._nn._upsample_nearest_exact3d(input, output_size, scale_factors)\\n\\n    if input.dim() == 3 and mode == \\\"area\\\":\\n        assert output_size is not None\\n        return adaptive_avg_pool1d(input, output_size)\\n    if input.dim() == 4 and mode == \\\"area\\\":\\n        assert output_size is not None\\n        return adaptive_avg_pool2d(input, output_size)\\n    if input.dim() == 5 and mode == \\\"area\\\":\\n        assert output_size is not None\\n        return adaptive_avg_pool3d(input, output_size)\\n\\n    if input.dim() == 3 and mode == \\\"linear\\\":\\n        assert align_corners is not None\\n        return torch._C._nn.upsample_linear1d(\\n            input, output_size, align_corners, scale_factors\\n        )\\n    if input.dim() == 4 and mode == \\\"bilinear\\\":\\n        assert align_corners is not None\\n        if antialias:\\n            return torch._C._nn._upsample_bilinear2d_aa(\\n                input, output_size, align_corners, scale_factors\\n            )\\n        # Two levels are necessary to prevent TorchScript from touching\\n        # are_deterministic_algorithms_enabled.\\n        if not torch.jit.is_scripting():\\n            if torch.are_deterministic_algorithms_enabled() and (\\n                input.is_cuda or input.is_xpu\\n            ):\\n                # Use slow decomp whose backward will be in terms of index_put\\n                # importlib is required because the import cannot be top level\\n                # (cycle) and cannot be nested (TS doesn't support)\\n                return importlib.import_module(\\n                    \\\"torch._decomp.decompositions\\\"\\n                )._upsample_linear_vec(input, output_size, align_corners, scale_factors)\\n        return torch._C._nn.upsample_bilinear2d(\\n            input, output_size, align_corners, scale_factors\\n        )\\n    if input.dim() == 5 and mode == \\\"trilinear\\\":\\n        assert align_corners is not None\\n        return torch._C._nn.upsample_trilinear3d(\\n            input, output_size, align_corners, scale_factors\\n        )\\n    if input.dim() == 4 and mode == \\\"bicubic\\\":\\n        assert align_corners is not None\\n        if antialias:\\n            return torch._C._nn._upsample_bicubic2d_aa(\\n                input, output_size, align_corners, scale_factors\\n            )\\n        return torch._C._nn.upsample_bicubic2d(\\n            input, output_size, align_corners, scale_factors\\n        )\\n\\n    if input.dim() == 3 and mode == \\\"bilinear\\\":\\n        raise NotImplementedError(\\\"Got 3D input, but bilinear mode needs 4D input\\\")\\n    if input.dim() == 3 and mode == \\\"trilinear\\\":\\n        raise NotImplementedError(\\\"Got 3D input, but trilinear mode needs 5D input\\\")\\n    if input.dim() == 4 and mode == \\\"linear\\\":\\n        raise NotImplementedError(\\\"Got 4D input, but linear mode needs 3D input\\\")\\n    if input.dim() == 4 and mode == \\\"trilinear\\\":\\n        raise NotImplementedError(\\\"Got 4D input, but trilinear mode needs 5D input\\\")\\n    if input.dim() == 5 and mode == \\\"linear\\\":\\n        raise NotImplementedError(\\\"Got 5D input, but linear mode needs 3D input\\\")\\n    if input.dim() == 5 and mode == \\\"bilinear\\\":\\n        raise NotImplementedError(\\\"Got 5D input, but bilinear mode needs 4D input\\\")\\n\\n    raise NotImplementedError(\\n        \\\"Input Error: Only 3D, 4D and 5D input Tensors supported\\\"\\n        f\\\" (got {input.dim()}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact\\\"\\n        f\\\" (got {mode})\\\"\\n    )\\n\\n\\nif interpolate.__doc__:\\n    interpolate.__doc__ = interpolate.__doc__.format(**reproducibility_notes)\\n\\n\\n@_overload\\ndef upsample_nearest(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[float] = None,\\n) -> Tensor:\\n    pass\\n\\n\\n@_overload\\ndef upsample_nearest(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[List[int]] = None,\\n    scale_factor: Optional[float] = None,\\n) -> Tensor:\\n    pass\\n\\n\\ndef upsample_nearest(input, size=None, scale_factor=None):  # noqa: F811\\n    r\\\"\\\"\\\"Upsamples the input, using nearest neighbours' pixel values.\\n\\n    .. warning::\\n        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.\\n        This is equivalent with ``nn.functional.interpolate(..., mode='nearest')``.\\n\\n    Currently spatial and volumetric upsampling are supported (i.e. expected\\n    inputs are 4 or 5 dimensional).\\n\\n    Args:\\n        input (Tensor): input\\n        size (int or Tuple[int, int] or Tuple[int, int, int]): output spatia\\n            size.\\n        scale_factor (int): multiplier for spatial size. Has to be an integer.\\n\\n    Note:\\n        {backward_reproducibility_note}\\n    \\\"\\\"\\\"\\n    # DeprecationWarning is ignored by default\\n    warnings.warn(\\n        \\\"`nn.functional.upsample_nearest` is deprecated. \\\"\\n        \\\"Use `nn.functional.interpolate` instead.\\\",\\n        stacklevel=2,\\n    )\\n    return interpolate(input, size, scale_factor, mode=\\\"nearest\\\")\\n\\n\\nif upsample_nearest.__doc__:\\n    upsample_nearest.__doc__ = upsample_nearest.__doc__.format(**reproducibility_notes)\\n\\n\\n@_overload\\ndef upsample_bilinear(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[float] = None,\\n) -> Tensor:\\n    pass\\n\\n\\n@_overload\\ndef upsample_bilinear(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[List[int]] = None,\\n    scale_factor: Optional[float] = None,\\n) -> Tensor:\\n    pass\\n\\n\\n@_overload\\ndef upsample_bilinear(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[int] = None,\\n    scale_factor: Optional[List[float]] = None,\\n) -> Tensor:\\n    pass\\n\\n\\n@_overload\\ndef upsample_bilinear(  # noqa: F811\\n    input: Tensor,\\n    size: Optional[List[int]] = None,\\n    scale_factor: Optional[List[float]] = None,\\n) -> Tensor:\\n    pass\\n\\n\\ndef upsample_bilinear(input, size=None, scale_factor=None):  # noqa: F811\\n    r\\\"\\\"\\\"Upsamples the input, using bilinear upsampling.\\n\\n    .. warning::\\n        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.\\n        This is equivalent with\\n        ``nn.functional.interpolate(..., mode='bilinear', align_corners=True)``.\\n\\n    Expected inputs are spatial (4 dimensional). Use `upsample_trilinear` fo\\n    volumetric (5 dimensional) inputs.\\n\\n    Args:\\n        input (Tensor): input\\n        size (int or Tuple[int, int]): output spatial size.\\n        scale_factor (int or Tuple[int, int]): multiplier for spatial size\\n\\n    Note:\\n        {backward_reproducibility_note}\\n    \\\"\\\"\\\"\\n    # DeprecationWarning is ignored by default\\n    warnings.warn(\\n        \\\"`nn.functional.upsample_bilinear` is deprecated. \\\"\\n        \\\"Use `nn.functional.interpolate` instead.\\\",\\n        stacklevel=2,\\n    )\\n    return interpolate(input, size, scale_factor, mode=\\\"bilinear\\\", align_corners=True)\\n\\n\\nif upsample_bilinear.__doc__:\\n    upsample_bilinear.__doc__ = upsample_bilinear.__doc__.format(\\n        **reproducibility_notes\\n    )\\n\\nGRID_SAMPLE_INTERPOLATION_MODES = {\\n    \\\"bilinear\\\": 0,\\n    \\\"nearest\\\": 1,\\n    \\\"bicubic\\\": 2,\\n}\\n\\nGRID_SAMPLE_PADDING_MODES = {\\n    \\\"zeros\\\": 0,\\n    \\\"border\\\": 1,\\n    \\\"reflection\\\": 2,\\n}\\n\\n\\ndef grid_sample(\\n    input: Tensor,\\n    grid: Tensor,\\n    mode: str = \\\"bilinear\\\",\\n    padding_mode: str = \\\"zeros\\\",\\n    align_corners: Optional[bool] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute grid sample.\\n\\n    Given an :attr:`input` and a flow-field :attr:`grid`, computes the\\n    ``output`` using :attr:`input` values and pixel locations from :attr:`grid`.\\n\\n    Currently, only spatial (4-D) and volumetric (5-D) :attr:`input` are\\n    supported.\\n\\n    In the spatial (4-D) case, for :attr:`input` with shape\\n    :math:`(N, C, H_\\\\text{in}, W_\\\\text{in})` and :attr:`grid` with shape\\n    :math:`(N, H_\\\\text{out}, W_\\\\text{out}, 2)`, the output will have shape\\n    :math:`(N, C, H_\\\\text{out}, W_\\\\text{out})`.\\n\\n    For each output location ``output[n, :, h, w]``, the size-2 vector\\n    ``grid[n, h, w]`` specifies :attr:`input` pixel locations ``x`` and ``y``,\\n    which are used to interpolate the output value ``output[n, :, h, w]``.\\n    In the case of 5D inputs, ``grid[n, d, h, w]`` specifies the\\n    ``x``, ``y``, ``z`` pixel locations for interpolating\\n    ``output[n, :, d, h, w]``. :attr:`mode` argument specifies ``nearest`` or\\n    ``bilinear`` interpolation method to sample the input pixels.\\n\\n    :attr:`grid` specifies the sampling pixel locations normalized by the\\n    :attr:`input` spatial dimensions. Therefore, it should have most values in\\n    the range of ``[-1, 1]``. For example, values ``x = -1, y = -1`` is the\\n    left-top pixel of :attr:`input`, and values  ``x = 1, y = 1`` is the\\n    right-bottom pixel of :attr:`input`.\\n\\n    If :attr:`grid` has values outside the range of ``[-1, 1]``, the corresponding\\n    outputs are handled as defined by :attr:`padding_mode`. Options are\\n\\n        * ``padding_mode=\\\"zeros\\\"``: use ``0`` for out-of-bound grid locations,\\n        * ``padding_mode=\\\"border\\\"``: use border values for out-of-bound grid locations,\\n        * ``padding_mode=\\\"reflection\\\"``: use values at locations reflected by\\n          the border for out-of-bound grid locations. For location far away\\n          from the border, it will keep being reflected until becoming in bound,\\n          e.g., (normalized) pixel location ``x = -3.5`` reflects by border ``-1``\\n          and becomes ``x' = 1.5``, then reflects by border ``1`` and becomes\\n          ``x'' = -0.5``.\\n\\n    Note:\\n        This function is often used in conjunction with :func:`affine_grid`\\n        to build `Spatial Transformer Networks`_ .\\n\\n    Note:\\n        When using the CUDA backend, this operation may induce nondeterministic\\n        behaviour in its backward pass that is not easily switched off.\\n        Please see the notes on :doc:`/notes/randomness` for background.\\n\\n    Note:\\n        NaN values in :attr:`grid` would be interpreted as ``-1``.\\n\\n    Args:\\n        input (Tensor): input of shape :math:`(N, C, H_\\\\text{in}, W_\\\\text{in})` (4-D case)\\n                        or :math:`(N, C, D_\\\\text{in}, H_\\\\text{in}, W_\\\\text{in})` (5-D case)\\n        grid (Tensor): flow-field of shape :math:`(N, H_\\\\text{out}, W_\\\\text{out}, 2)` (4-D case)\\n                       or :math:`(N, D_\\\\text{out}, H_\\\\text{out}, W_\\\\text{out}, 3)` (5-D case)\\n        mode (str): interpolation mode to calculate output values\\n            ``'bilinear'`` | ``'nearest'`` | ``'bicubic'``. Default: ``'bilinear'``\\n            Note: ``mode='bicubic'`` supports only 4-D input.\\n            When ``mode='bilinear'`` and the input is 5-D, the interpolation mode\\n            used internally will actually be trilinear. However, when the input is 4-D,\\n            the interpolation mode will legitimately be bilinear.\\n        padding_mode (str): padding mode for outside grid values\\n            ``'zeros'`` | ``'border'`` | ``'reflection'``. Default: ``'zeros'``\\n        align_corners (bool, optional): Geometrically, we consider the pixels of the\\n            input  as squares rather than points.\\n            If set to ``True``, the extrema (``-1`` and ``1``) are considered as referring\\n            to the center points of the input's corner pixels. If set to ``False``, they\\n            are instead considered as referring to the corner points of the input's corner\\n            pixels, making the sampling more resolution agnostic.\\n            This option parallels the ``align_corners`` option in\\n            :func:`interpolate`, and so whichever option is used here\\n            should also be used there to resize the input image before grid sampling.\\n            Default: ``False``\\n\\n    Returns:\\n        output (Tensor): output Tensor\\n\\n    .. _`Spatial Transformer Networks`:\\n        https://arxiv.org/abs/1506.02025\\n\\n    .. warning::\\n        When ``align_corners = True``, the grid positions depend on the pixel\\n        size relative to the input image size, and so the locations sampled by\\n        :func:`grid_sample` will differ for the same input given at different\\n        resolutions (that is, after being upsampled or downsampled).\\n        The default behavior up to version 1.2.0 was ``align_corners = True``.\\n        Since then, the default behavior has been changed to ``align_corners = False``,\\n        in order to bring it in line with the default for :func:`interpolate`.\\n\\n    .. note::\\n        ``mode='bicubic'`` is implemented using the `cubic convolution algorithm`_ with :math:`\\\\alpha=-0.75`.\\n        The constant :math:`\\\\alpha` might be different from packages to packages.\\n        For example, `PIL`_ and `OpenCV`_ use -0.5 and -0.75 respectively.\\n        This algorithm may \\\"overshoot\\\" the range of values it's interpolating.\\n        For example, it may produce negative values or values greater than 255 when interpolating input in [0, 255].\\n        Clamp the results with :func:`torch.clamp` to ensure they are within the valid range.\\n    .. _`cubic convolution algorithm`: https://en.wikipedia.org/wiki/Bicubic_interpolation\\n    .. _`PIL`: https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/src/libImaging/Resample.c#L51\\n    .. _`OpenCV`: https://github.com/opencv/opencv/blob/f345ed564a06178670750bad59526cfa4033be55/modules/imgproc/src/resize.cpp#L908\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, grid):\\n        return handle_torch_function(\\n            grid_sample,\\n            (input, grid),\\n            input,\\n            grid,\\n            mode=mode,\\n            padding_mode=padding_mode,\\n            align_corners=align_corners,\\n        )\\n    if mode != \\\"bilinear\\\" and mode != \\\"nearest\\\" and mode != \\\"bicubic\\\":\\n        raise ValueError(\\n            f\\\"nn.functional.grid_sample(): expected mode to be 'bilinear', 'nearest' or 'bicubic', but got: '{mode}'\\\"\\n        )\\n    if (\\n        padding_mode != \\\"zeros\\\"\\n        and padding_mode != \\\"border\\\"\\n        and padding_mode != \\\"reflection\\\"\\n    ):\\n        raise ValueError(\\n            \\\"nn.functional.grid_sample(): expected padding_mode \\\"\\n            \\\"to be 'zeros', 'border', or 'reflection', \\\"\\n            f\\\"but got: '{padding_mode}'\\\"\\n        )\\n\\n    if mode == \\\"bilinear\\\":\\n        mode_enum = 0\\n    elif mode == \\\"nearest\\\":\\n        mode_enum = 1\\n    else:  # mode == 'bicubic'\\n        mode_enum = 2\\n\\n    if padding_mode == \\\"zeros\\\":\\n        padding_mode_enum = 0\\n    elif padding_mode == \\\"border\\\":\\n        padding_mode_enum = 1\\n    else:  # padding_mode == 'reflection'\\n        padding_mode_enum = 2\\n\\n    if align_corners is None:\\n        warnings.warn(\\n            \\\"Default grid_sample and affine_grid behavior has changed \\\"\\n            \\\"to align_corners=False since 1.3.0. Please specify \\\"\\n            \\\"align_corners=True if the old behavior is desired. \\\"\\n            \\\"See the documentation of grid_sample for details.\\\"\\n        )\\n        align_corners = False\\n\\n    return torch.grid_sampler(input, grid, mode_enum, padding_mode_enum, align_corners)\\n\\n\\ndef affine_grid(\\n    theta: Tensor,\\n    size: List[int],\\n    align_corners: Optional[bool] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Generate 2D or 3D flow field (sampling grid), given a batch of affine matrices :attr:`theta`.\\n\\n    .. note::\\n        This function is often used in conjunction with :func:`grid_sample`\\n        to build `Spatial Transformer Networks`_ .\\n\\n    Args:\\n        theta (Tensor): input batch of affine matrices with shape\\n            (:math:`N \\\\times 2 \\\\times 3`) for 2D or\\n            (:math:`N \\\\times 3 \\\\times 4`) for 3D\\n        size (torch.Size): the target output image size.\\n            (:math:`N \\\\times C \\\\times H \\\\times W` for 2D or\\n            :math:`N \\\\times C \\\\times D \\\\times H \\\\times W` for 3D)\\n            Example: torch.Size((32, 3, 24, 24))\\n        align_corners (bool, optional): if ``True``, consider ``-1`` and ``1``\\n            to refer to the centers of the corner pixels rather than the image corners.\\n            Refer to :func:`grid_sample` for a more complete description.\\n            A grid generated by :func:`affine_grid` should be passed to :func:`grid_sample`\\n            with the same setting for this option.\\n            Default: ``False``\\n\\n    Returns:\\n        output (Tensor): output Tensor of size (:math:`N \\\\times H \\\\times W \\\\times 2`)\\n\\n    .. _`Spatial Transformer Networks`:\\n        https://arxiv.org/abs/1506.02025\\n\\n    .. warning::\\n        When ``align_corners = True``, the grid positions depend on the pixel\\n        size relative to the input image size, and so the locations sampled by\\n        :func:`grid_sample` will differ for the same input given at different\\n        resolutions (that is, after being upsampled or downsampled).\\n        The default behavior up to version 1.2.0 was ``align_corners = True``.\\n        Since then, the default behavior has been changed to ``align_corners = False``,\\n        in order to bring it in line with the default for :func:`interpolate`.\\n    .. warning::\\n        When ``align_corners = True``, 2D affine transforms on 1D data and\\n        3D affine transforms on 2D data (that is, when one of the spatial\\n        dimensions has unit size) are ill-defined, and not an intended use case.\\n        This is not a problem when ``align_corners = False``.\\n        Up to version 1.2.0, all grid points along a unit dimension were\\n        considered arbitrarily to be at ``-1``.\\n        From version 1.3.0, under ``align_corners = True`` all grid points\\n        along a unit dimension are considered to be at ``0``\\n        (the center of the input image).\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(theta):\\n        return handle_torch_function(\\n            affine_grid, (theta,), theta, size, align_corners=align_corners\\n        )\\n    if align_corners is None:\\n        warnings.warn(\\n            \\\"Default grid_sample and affine_grid behavior has changed \\\"\\n            \\\"to align_corners=False since 1.3.0. Please specify \\\"\\n            \\\"align_corners=True if the old behavior is desired. \\\"\\n            \\\"See the documentation of grid_sample for details.\\\"\\n        )\\n        align_corners = False\\n\\n    # enforce floating point dtype on theta\\n    if not theta.is_floating_point():\\n        raise ValueError(\\n            f\\\"Expected theta to have floating point type, but got {theta.dtype}\\\"\\n        )\\n    # check that shapes and sizes match\\n    if len(size) == 4:\\n        if theta.dim() != 3 or theta.shape[-2] != 2 or theta.shape[-1] != 3:\\n            raise ValueError(\\n                f\\\"Expected a batch of 2D affine matrices of shape Nx2x3 for size {size}. Got {theta.shape}.\\\"\\n            )\\n        spatial_size = size[-2:]  # spatial dimension sizes\\n    elif len(size) == 5:\\n        if theta.dim() != 3 or theta.shape[-2] != 3 or theta.shape[-1] != 4:\\n            raise ValueError(\\n                f\\\"Expected a batch of 3D affine matrices of shape Nx3x4 for size {size}. Got {theta.shape}.\\\"\\n            )\\n        spatial_size = size[-3:]  # spatial dimension sizes\\n    else:\\n        raise NotImplementedError(\\n            \\\"affine_grid only supports 4D and 5D sizes, \\\"\\n            \\\"for 2D and 3D affine transforms, respectively. \\\"\\n            f\\\"Got size {size}.\\\"\\n        )\\n    # check for empty span\\n    if align_corners and min(spatial_size) == 1:\\n        warnings.warn(\\n            \\\"Since version 1.3.0, affine_grid behavior has changed \\\"\\n            \\\"for unit-size grids when align_corners=True. \\\"\\n            \\\"This is not an intended use case of affine_grid. \\\"\\n            \\\"See the documentation of affine_grid for details.\\\"\\n        )\\n    elif min(size) <= 0:\\n        raise ValueError(f\\\"Expected non-zero, positive output size. Got {size}\\\")\\n\\n    return torch.affine_grid_generator(theta, size, align_corners)\\n\\n\\ndef pad(\\n    input: Tensor,\\n    pad: List[int],\\n    mode: str = \\\"constant\\\",\\n    value: Optional[float] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"\\n    pad(input, pad, mode=\\\"constant\\\", value=None) -> Tensor\\n\\n    Pads tensor.\\n\\n    Padding size:\\n        The padding size by which to pad some dimensions of :attr:`input`\\n        are described starting from the last dimension and moving forward.\\n        :math:`\\\\left\\\\lfloor\\\\frac{\\\\text{len(pad)}}{2}\\\\right\\\\rfloor` dimensions\\n        of ``input`` will be padded.\\n        For example, to pad only the last dimension of the input tensor, then\\n        :attr:`pad` has the form\\n        :math:`(\\\\text{padding\\\\_left}, \\\\text{padding\\\\_right})`;\\n        to pad the last 2 dimensions of the input tensor, then use\\n        :math:`(\\\\text{padding\\\\_left}, \\\\text{padding\\\\_right},`\\n        :math:`\\\\text{padding\\\\_top}, \\\\text{padding\\\\_bottom})`;\\n        to pad the last 3 dimensions, use\\n        :math:`(\\\\text{padding\\\\_left}, \\\\text{padding\\\\_right},`\\n        :math:`\\\\text{padding\\\\_top}, \\\\text{padding\\\\_bottom}`\\n        :math:`\\\\text{padding\\\\_front}, \\\\text{padding\\\\_back})`.\\n\\n    Padding mode:\\n        See :class:`torch.nn.CircularPad2d`, :class:`torch.nn.ConstantPad2d`,\\n        :class:`torch.nn.ReflectionPad2d`, and :class:`torch.nn.ReplicationPad2d`\\n        for concrete examples on how each of the padding modes works. Constant\\n        padding is implemented for arbitrary dimensions. Circular, replicate and\\n        reflection padding are implemented for padding the last 3 dimensions of a\\n        4D or 5D input tensor, the last 2 dimensions of a 3D or 4D input tensor,\\n        or the last dimension of a 2D or 3D input tensor.\\n\\n    Note:\\n        When using the CUDA backend, this operation may induce nondeterministic\\n        behaviour in its backward pass that is not easily switched off.\\n        Please see the notes on :doc:`/notes/randomness` for background.\\n\\n    Args:\\n        input (Tensor): N-dimensional tensor\\n        pad (tuple): m-elements tuple, where\\n            :math:`\\\\frac{m}{2} \\\\leq` input dimensions and :math:`m` is even.\\n        mode: ``'constant'``, ``'reflect'``, ``'replicate'`` or ``'circular'``.\\n            Default: ``'constant'``\\n        value: fill value for ``'constant'`` padding. Default: ``0``\\n\\n    Examples::\\n\\n        >>> t4d = torch.empty(3, 3, 4, 2)\\n        >>> p1d = (1, 1) # pad last dim by 1 on each side\\n        >>> out = F.pad(t4d, p1d, \\\"constant\\\", 0)  # effectively zero padding\\n        >>> print(out.size())\\n        torch.Size([3, 3, 4, 4])\\n        >>> p2d = (1, 1, 2, 2) # pad last dim by (1, 1) and 2nd to last by (2, 2)\\n        >>> out = F.pad(t4d, p2d, \\\"constant\\\", 0)\\n        >>> print(out.size())\\n        torch.Size([3, 3, 8, 4])\\n        >>> t4d = torch.empty(3, 3, 4, 2)\\n        >>> p3d = (0, 1, 2, 1, 3, 3) # pad by (0, 1), (2, 1), and (3, 3)\\n        >>> out = F.pad(t4d, p3d, \\\"constant\\\", 0)\\n        >>> print(out.size())\\n        torch.Size([3, 9, 7, 3])\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            torch.nn.functional.pad, (input,), input, pad, mode=mode, value=value\\n        )\\n    if not torch.jit.is_scripting():\\n        if torch.are_deterministic_algorithms_enabled() and (\\n            input.is_cuda or input.is_xpu\\n        ):\\n            if mode == \\\"replicate\\\":\\n                # Use slow decomp whose backward will be in terms of index_put.\\n                # importlib is required because the import cannot be top level\\n                # (cycle) and cannot be nested (TS doesn't support)\\n                return importlib.import_module(\\n                    \\\"torch._decomp.decompositions\\\"\\n                )._replication_pad(input, pad)\\n    return torch._C._nn.pad(input, pad, mode, value)\\n\\n\\n# TODO: Fix via https://github.com/pytorch/pytorch/issues/75798\\npad.__module__ = \\\"torch.nn.functional\\\"\\n\\n# distance\\n\\n\\npairwise_distance = _add_docstr(\\n    torch.pairwise_distance,\\n    r\\\"\\\"\\\"\\npairwise_distance(x1, x2, p=2.0, eps=1e-6, keepdim=False) -> Tensor\\n\\nSee :class:`torch.nn.PairwiseDistance` for details\\n\\\"\\\"\\\",\\n)\\n\\n\\npdist = _add_docstr(\\n    torch.pdist,\\n    r\\\"\\\"\\\"\\npdist(input, p=2) -> Tensor\\n\\nComputes the p-norm distance between every pair of row vectors in the input.\\nThis is identical to the upper triangular portion, excluding the diagonal, of\\n`torch.norm(input[:, None] - input, dim=2, p=p)`. This function will be faster\\nif the rows are contiguous.\\n\\nIf input has shape :math:`N \\\\times M` then the output will have shape\\n:math:`\\\\frac{1}{2} N (N - 1)`.\\n\\nThis function is equivalent to ``scipy.spatial.distance.pdist(input,\\n'minkowski', p=p)`` if :math:`p \\\\in (0, \\\\infty)`. When :math:`p = 0` it is\\nequivalent to ``scipy.spatial.distance.pdist(input, 'hamming') * M``.\\nWhen :math:`p = \\\\infty`, the closest scipy function is\\n``scipy.spatial.distance.pdist(xn, lambda x, y: np.abs(x - y).max())``.\\n\\nArgs:\\n    input: input tensor of shape :math:`N \\\\times M`.\\n    p: p value for the p-norm distance to calculate between each vector pair\\n        :math:`\\\\in [0, \\\\infty]`.\\n\\\"\\\"\\\",\\n)\\n\\n\\ncosine_similarity = _add_docstr(\\n    torch.cosine_similarity,\\n    r\\\"\\\"\\\"\\ncosine_similarity(x1, x2, dim=1, eps=1e-8) -> Tensor\\n\\nReturns cosine similarity between ``x1`` and ``x2``, computed along dim. ``x1`` and ``x2`` must be broadcastable\\nto a common shape. ``dim`` refers to the dimension in this common shape. Dimension ``dim`` of the output is\\nsqueezed (see :func:`torch.squeeze`), resulting in the\\noutput tensor having 1 fewer dimension.\\n\\n.. math ::\\n    \\\\text{similarity} = \\\\dfrac{x_1 \\\\cdot x_2}{\\\\max(\\\\Vert x_1 \\\\Vert _2, \\\\epsilon) \\\\cdot \\\\max(\\\\Vert x_2 \\\\Vert _2, \\\\epsilon)}\\n\\nSupports :ref:`type promotion <type-promotion-doc>`.\\n\\nArgs:\\n    x1 (Tensor): First input.\\n    x2 (Tensor): Second input.\\n    dim (int, optional): Dimension along which cosine similarity is computed. Default: 1\\n    eps (float, optional): Small value to avoid division by zero.\\n        Default: 1e-8\\n\\nExample::\\n\\n    >>> input1 = torch.randn(100, 128)\\n    >>> input2 = torch.randn(100, 128)\\n    >>> output = F.cosine_similarity(input1, input2)\\n    >>> print(output)\\n\\\"\\\"\\\",\\n)\\n\\n\\none_hot = _add_docstr(\\n    torch._C._nn.one_hot,\\n    r\\\"\\\"\\\"\\none_hot(tensor, num_classes=-1) -> LongTensor\\n\\nTakes LongTensor with index values of shape ``(*)`` and returns a tensor\\nof shape ``(*, num_classes)`` that have zeros everywhere except where the\\nindex of last dimension matches the corresponding value of the input tensor,\\nin which case it will be 1.\\n\\nSee also `One-hot on Wikipedia`_ .\\n\\n.. _One-hot on Wikipedia:\\n    https://en.wikipedia.org/wiki/One-hot\\n\\nArguments:\\n    tensor (LongTensor): class values of any shape.\\n    num_classes (int):  Total number of classes. If set to -1, the number\\n        of classes will be inferred as one greater than the largest class\\n        value in the input tensor.\\n\\nReturns:\\n    LongTensor that has one more dimension with 1 values at the\\n    index of last dimension indicated by the input, and 0 everywhere\\n    else.\\n\\nExamples:\\n    >>> F.one_hot(torch.arange(0, 5) % 3)\\n    tensor([[1, 0, 0],\\n            [0, 1, 0],\\n            [0, 0, 1],\\n            [1, 0, 0],\\n            [0, 1, 0]])\\n    >>> F.one_hot(torch.arange(0, 5) % 3, num_classes=5)\\n    tensor([[1, 0, 0, 0, 0],\\n            [0, 1, 0, 0, 0],\\n            [0, 0, 1, 0, 0],\\n            [1, 0, 0, 0, 0],\\n            [0, 1, 0, 0, 0]])\\n    >>> F.one_hot(torch.arange(0, 6).view(3,2) % 3)\\n    tensor([[[1, 0, 0],\\n             [0, 1, 0]],\\n            [[0, 0, 1],\\n             [1, 0, 0]],\\n            [[0, 1, 0],\\n             [0, 0, 1]]])\\n\\\"\\\"\\\",\\n)\\n\\n\\ndef triplet_margin_loss(\\n    anchor: Tensor,\\n    positive: Tensor,\\n    negative: Tensor,\\n    margin: float = 1.0,\\n    p: float = 2,\\n    eps: float = 1e-6,\\n    swap: bool = False,\\n    size_average: Optional[bool] = None,\\n    reduce: Optional[bool] = None,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the triplet loss between given input tensors and a margin greater than 0.\\n\\n    See :class:`~torch.nn.TripletMarginLoss` for details.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(anchor, positive, negative):\\n        return handle_torch_function(\\n            triplet_margin_loss,\\n            (anchor, positive, negative),\\n            anchor,\\n            positive,\\n            negative,\\n            margin=margin,\\n            p=p,\\n            eps=eps,\\n            swap=swap,\\n            size_average=size_average,\\n            reduce=reduce,\\n            reduction=reduction,\\n        )\\n    if size_average is not None or reduce is not None:\\n        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)\\n    else:\\n        reduction_enum = _Reduction.get_enum(reduction)\\n    if margin <= 0:\\n        raise ValueError(f\\\"margin must be greater than 0, got {margin}\\\")\\n    return torch.triplet_margin_loss(\\n        anchor, positive, negative, margin, p, eps, swap, reduction_enum\\n    )\\n\\n\\ndef triplet_margin_with_distance_loss(\\n    anchor: Tensor,\\n    positive: Tensor,\\n    negative: Tensor,\\n    *,\\n    distance_function: Optional[Callable[[Tensor, Tensor], Tensor]] = None,\\n    margin: float = 1.0,\\n    swap: bool = False,\\n    reduction: str = \\\"mean\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Compute the triplet margin loss for input tensors using a custom distance function.\\n\\n    See :class:`~torch.nn.TripletMarginWithDistanceLoss` for details.\\n    \\\"\\\"\\\"\\n    if torch.jit.is_scripting():\\n        raise NotImplementedError(\\n            \\\"F.triplet_margin_with_distance_loss does not support JIT scripting: \\\"\\n            \\\"functions requiring Callables cannot be scripted.\\\"\\n        )\\n\\n    if has_torch_function_variadic(anchor, positive, negative):\\n        return handle_torch_function(\\n            triplet_margin_with_distance_loss,\\n            (anchor, positive, negative),\\n            anchor,\\n            positive,\\n            negative,\\n            distance_function=distance_function,\\n            margin=margin,\\n            swap=swap,\\n            reduction=reduction,\\n        )\\n\\n    # Check validity of reduction mode\\n    if reduction not in (\\\"mean\\\", \\\"sum\\\", \\\"none\\\"):\\n        raise ValueError(f\\\"{reduction} is not a valid value for reduction\\\")\\n\\n    # Check validity of margin\\n    if margin <= 0:\\n        raise ValueError(f\\\"margin must be greater than 0, got {margin}\\\")\\n\\n    # Check dimensions\\n    a_dim = anchor.ndim\\n    p_dim = positive.ndim\\n    n_dim = negative.ndim\\n    if not (a_dim == p_dim and p_dim == n_dim):\\n        raise RuntimeError(\\n            f\\\"The anchor, positive, and negative tensors are expected to have \\\"\\n            f\\\"the same number of dimensions, but got: anchor {a_dim}D, \\\"\\n            f\\\"positive {p_dim}D, and negative {n_dim}D inputs\\\"\\n        )\\n\\n    # Calculate loss\\n    if distance_function is None:\\n        distance_function = torch.pairwise_distance\\n\\n    dist_pos = distance_function(anchor, positive)\\n    dist_neg = distance_function(anchor, negative)\\n    # The distance swap is described in the paper \\\"Learning shallow\\n    # convolutional feature descriptors with triplet losses\\\" by V. Balntas, E.\\n    # Riba et al.  If True, and if the positive example is closer to the\\n    # negative example than the anchor is, swaps the positive example and the\\n    # anchor in the loss computation.\\n    if swap:\\n        dist_swap = distance_function(positive, negative)\\n        dist_neg = torch.minimum(dist_neg, dist_swap)\\n    loss = torch.clamp_min(margin + dist_pos - dist_neg, 0)\\n\\n    # Apply reduction\\n    if reduction == \\\"sum\\\":\\n        return torch.sum(loss)\\n    elif reduction == \\\"mean\\\":\\n        return torch.mean(loss)\\n    else:  # reduction == \\\"none\\\"\\n        return loss\\n\\n\\ndef normalize(\\n    input: Tensor,\\n    p: float = 2.0,\\n    dim: int = 1,\\n    eps: float = 1e-12,\\n    out: Optional[Tensor] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Perform :math:`L_p` normalization of inputs over specified dimension.\\n\\n    For a tensor :attr:`input` of sizes :math:`(n_0, ..., n_{dim}, ..., n_k)`, each\\n    :math:`n_{dim}` -element vector :math:`v` along dimension :attr:`dim` is transformed as\\n\\n    .. math::\\n        v = \\\\frac{v}{\\\\max(\\\\lVert v \\\\rVert_p, \\\\epsilon)}.\\n\\n    With the default arguments it uses the Euclidean norm over vectors along dimension :math:`1` for normalization.\\n\\n    Args:\\n        input: input tensor of any shape\\n        p (float): the exponent value in the norm formulation. Default: 2\\n        dim (int or tuple of ints): the dimension to reduce. Default: 1\\n        eps (float): small value to avoid division by zero. Default: 1e-12\\n        out (Tensor, optional): the output tensor. If :attr:`out` is used, this\\n                                operation won't be differentiable.\\n    \\\"\\\"\\\"\\n    if has_torch_function_variadic(input, out):\\n        return handle_torch_function(\\n            normalize, (input, out), input, p=p, dim=dim, eps=eps, out=out\\n        )\\n    if out is None:\\n        denom = input.norm(p, dim, keepdim=True).clamp_min(eps).expand_as(input)\\n        return input / denom\\n    else:\\n        denom = input.norm(p, dim, keepdim=True).clamp_min_(eps).expand_as(input)\\n        return torch.div(input, denom, out=out)\\n\\n\\ndef assert_int_or_pair(arg: List[int], arg_name: str, message: str) -> None:\\n    assert isinstance(arg, int) or len(arg) == 2, message.format(arg_name)\\n\\n\\ndef unfold(\\n    input: Tensor,\\n    kernel_size: BroadcastingList2[int],\\n    dilation: BroadcastingList2[int] = 1,\\n    padding: BroadcastingList2[int] = 0,\\n    stride: BroadcastingList2[int] = 1,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Extract sliding local blocks from a batched input tensor.\\n\\n    .. warning::\\n        Currently, only 4-D input tensors (batched image-like tensors) are\\n        supported.\\n\\n    .. warning::\\n\\n        More than one element of the unfolded tensor may refer to a single\\n        memory location. As a result, in-place operations (especially ones that\\n        are vectorized) may result in incorrect behavior. If you need to write\\n        to the tensor, please clone it first.\\n\\n\\n    See :class:`torch.nn.Unfold` for details\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            unfold,\\n            (input,),\\n            input,\\n            kernel_size,\\n            dilation=dilation,\\n            padding=padding,\\n            stride=stride,\\n        )\\n    return torch._C._nn.im2col(\\n        input, _pair(kernel_size), _pair(dilation), _pair(padding), _pair(stride)\\n    )\\n\\n\\ndef fold(\\n    input: Tensor,\\n    output_size: BroadcastingList2[int],\\n    kernel_size: BroadcastingList2[int],\\n    dilation: BroadcastingList2[int] = 1,\\n    padding: BroadcastingList2[int] = 0,\\n    stride: BroadcastingList2[int] = 1,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Combine an array of sliding local blocks into a large containing tensor.\\n\\n    .. warning::\\n        Currently, only unbatched (3D) or batched (4D) image-like output tensors are supported.\\n\\n    See :class:`torch.nn.Fold` for details\\n    \\\"\\\"\\\"\\n    if has_torch_function_unary(input):\\n        return handle_torch_function(\\n            fold,\\n            (input,),\\n            input,\\n            output_size,\\n            kernel_size,\\n            dilation=dilation,\\n            padding=padding,\\n            stride=stride,\\n        )\\n    return torch._C._nn.col2im(\\n        input,\\n        _pair(output_size),\\n        _pair(kernel_size),\\n        _pair(dilation),\\n        _pair(padding),\\n        _pair(stride),\\n    )\\n\\n\\n#\\n# multihead attention\\n#\\n\\n\\ndef _in_projection_packed(\\n    q: Tensor,\\n    k: Tensor,\\n    v: Tensor,\\n    w: Tensor,\\n    b: Optional[Tensor] = None,\\n) -> List[Tensor]:\\n    r\\\"\\\"\\\"Perform the in-projection step of the attention operation, using packed weights.\\n\\n    Output is a triple containing projection tensors for query, key and value.\\n\\n    Args:\\n        q, k, v: query, key and value tensors to be projected. For self-attention,\\n            these are typically the same tensor; for encoder-decoder attention,\\n            k and v are typically the same tensor. (We take advantage of these\\n            identities for performance if they are present.) Regardless, q, k and v\\n            must share a common embedding dimension; otherwise their shapes may vary.\\n        w: projection weights for q, k and v, packed into a single tensor. Weights\\n            are packed along dimension 0, in q, k, v order.\\n        b: optional projection biases for q, k and v, packed into a single tensor\\n            in q, k, v order.\\n\\n    Shape:\\n        Inputs:\\n        - q: :math:`(..., E)` where E is the embedding dimension\\n        - k: :math:`(..., E)` where E is the embedding dimension\\n        - v: :math:`(..., E)` where E is the embedding dimension\\n        - w: :math:`(E * 3, E)` where E is the embedding dimension\\n        - b: :math:`E * 3` where E is the embedding dimension\\n\\n        Output:\\n        - in output list :math:`[q', k', v']`, each output tensor will have the\\n            same shape as the corresponding input tensor.\\n    \\\"\\\"\\\"\\n    E = q.size(-1)\\n    if k is v:\\n        if q is k:\\n            # self-attention\\n            proj = linear(q, w, b)\\n            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()\\n            proj = (\\n                proj.unflatten(-1, (3, E))\\n                .unsqueeze(0)\\n                .transpose(0, -2)\\n                .squeeze(-2)\\n                .contiguous()\\n            )\\n            return proj[0], proj[1], proj[2]\\n        else:\\n            # encoder-decoder attention\\n            w_q, w_kv = w.split([E, E * 2])\\n            if b is None:\\n                b_q = b_kv = None\\n            else:\\n                b_q, b_kv = b.split([E, E * 2])\\n            q_proj = linear(q, w_q, b_q)\\n            kv_proj = linear(k, w_kv, b_kv)\\n            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()\\n            kv_proj = (\\n                kv_proj.unflatten(-1, (2, E))\\n                .unsqueeze(0)\\n                .transpose(0, -2)\\n                .squeeze(-2)\\n                .contiguous()\\n            )\\n            return (q_proj, kv_proj[0], kv_proj[1])\\n    else:\\n        w_q, w_k, w_v = w.chunk(3)\\n        if b is None:\\n            b_q = b_k = b_v = None\\n        else:\\n            b_q, b_k, b_v = b.chunk(3)\\n        return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)\\n\\n\\ndef _in_projection(\\n    q: Tensor,\\n    k: Tensor,\\n    v: Tensor,\\n    w_q: Tensor,\\n    w_k: Tensor,\\n    w_v: Tensor,\\n    b_q: Optional[Tensor] = None,\\n    b_k: Optional[Tensor] = None,\\n    b_v: Optional[Tensor] = None,\\n) -> Tuple[Tensor, Tensor, Tensor]:\\n    r\\\"\\\"\\\"Perform the in-projection step of the attention operation.\\n\\n    This is simply a triple of linear projections,\\n    with shape constraints on the weights which\\n    ensure embedding dimension uniformity in the projected outputs.\\n    Output is a triple containing projection tensors for query, key and value.\\n\\n    Args:\\n        q, k, v: query, key and value tensors to be projected.\\n        w_q, w_k, w_v: weights for q, k and v, respectively.\\n        b_q, b_k, b_v: optional biases for q, k and v, respectively.\\n\\n    Shape:\\n        Inputs:\\n        - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any\\n            number of leading dimensions.\\n        - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any\\n            number of leading dimensions.\\n        - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any\\n            number of leading dimensions.\\n        - w_q: :math:`(Eq, Eq)`\\n        - w_k: :math:`(Eq, Ek)`\\n        - w_v: :math:`(Eq, Ev)`\\n        - b_q: :math:`(Eq)`\\n        - b_k: :math:`(Eq)`\\n        - b_v: :math:`(Eq)`\\n\\n        Output: in output triple :math:`(q', k', v')`,\\n         - q': :math:`[Qdims..., Eq]`\\n         - k': :math:`[Kdims..., Eq]`\\n         - v': :math:`[Vdims..., Eq]`\\n\\n    \\\"\\\"\\\"\\n    Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)\\n    assert w_q.shape == (\\n        Eq,\\n        Eq,\\n    ), f\\\"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}\\\"\\n    assert w_k.shape == (\\n        Eq,\\n        Ek,\\n    ), f\\\"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}\\\"\\n    assert w_v.shape == (\\n        Eq,\\n        Ev,\\n    ), f\\\"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}\\\"\\n    assert b_q is None or b_q.shape == (\\n        Eq,\\n    ), f\\\"expecting query bias shape of {(Eq,)}, but got {b_q.shape}\\\"\\n    assert b_k is None or b_k.shape == (\\n        Eq,\\n    ), f\\\"expecting key bias shape of {(Eq,)}, but got {b_k.shape}\\\"\\n    assert b_v is None or b_v.shape == (\\n        Eq,\\n    ), f\\\"expecting value bias shape of {(Eq,)}, but got {b_v.shape}\\\"\\n    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)\\n\\n\\nscaled_dot_product_attention = _add_docstr(\\n    torch._C._nn.scaled_dot_product_attention,\\n    r\\\"\\\"\\\"scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,\\n        is_causal=False, scale=None, enable_gqa=False) -> Tensor:\\n\\n    Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed,\\n    and applying dropout if a probability greater than 0.0 is specified. The optional scale argument can only be\\n    specified as a keyword argument.\\n\\n    .. code-block:: python\\n\\n        # Efficient implementation equivalent to the following:\\n        def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,\\n                is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor:\\n            L, S = query.size(-2), key.size(-2)\\n            scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale\\n            attn_bias = torch.zeros(L, S, dtype=query.dtype)\\n            if is_causal:\\n                assert attn_mask is None\\n                temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)\\n                attn_bias.masked_fill_(temp_mask.logical_not(), float(\\\"-inf\\\"))\\n                attn_bias.to(query.dtype)\\n\\n            if attn_mask is not None:\\n                if attn_mask.dtype == torch.bool:\\n                    attn_bias.masked_fill_(attn_mask.logical_not(), float(\\\"-inf\\\"))\\n                else:\\n                    attn_bias += attn_mask\\n\\n            if enable_gqa:\\n                key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)\\n                value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)\\n\\n            attn_weight = query @ key.transpose(-2, -1) * scale_factor\\n            attn_weight += attn_bias\\n            attn_weight = torch.softmax(attn_weight, dim=-1)\\n            attn_weight = torch.dropout(attn_weight, dropout_p, train=True)\\n            return attn_weight @ value\\n\\n    .. warning::\\n        This function is beta and subject to change.\\n\\n    .. warning::\\n        This function always applies dropout according to the specified ``dropout_p`` argument.\\n        To disable dropout during evaluation, be sure to pass a value of ``0.0`` when the module\\n        that makes the function call is not in training mode.\\n\\n        For example:\\n\\n        .. code-block:: python\\n\\n            class MyModel(nn.Module):\\n                def __init__(self, p=0.5):\\n                    super().__init__()\\n                    self.p = p\\n\\n                def forward(self, ...):\\n                    return F.scaled_dot_product_attention(...,\\n                        dropout_p=(self.p if self.training else 0.0))\\n\\n    Note:\\n\\n        There are currently three supported implementations of scaled dot product attention:\\n\\n            - `FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning`_\\n            - `Memory-Efficient Attention`_\\n            - A PyTorch implementation defined in C++ matching the above formulation\\n\\n        The function may call optimized kernels for improved performance when using the CUDA backend.\\n        For all other backends, the PyTorch implementation will be used.\\n\\n        All implementations are enabled by default. Scaled dot product attention attempts to automatically select the\\n        most optimal implementation based on the inputs. In order to provide more fine-grained control over what implementation\\n        is used, the following functions are provided for enabling and disabling implementations.\\n        The context manager is the preferred mechanism:\\n\\n            - :func:`torch.nn.attention.sdpa_kernel`: A context manager used to enable or disable any of the implementations.\\n            - :func:`torch.backends.cuda.enable_flash_sdp`: Globally enables or disables FlashAttention.\\n            - :func:`torch.backends.cuda.enable_mem_efficient_sdp`: Globally enables or disables  Memory-Efficient Attention.\\n            - :func:`torch.backends.cuda.enable_math_sdp`: Globally enables or disables  the PyTorch C++ implementation.\\n\\n        Each of the fused kernels has specific input limitations. If the user requires the use of a specific fused implementation,\\n        disable the PyTorch C++ implementation using :func:`torch.nn.attention.sdpa_kernel`.\\n        In the event that a fused implementation is not available, a warning will be raised with the\\n        reasons why the fused implementation cannot run.\\n\\n        Due to the nature of fusing floating point operations, the output of this function may be different\\n        depending on what backend kernel is chosen.\\n        The c++ implementation supports torch.float64 and can be used when higher precision is required.\\n        For math backend, all intermediates are kept in torch.float if inputs are in torch.half or torch.bfloat16.\\n    For more information please see :doc:`/notes/numerical_accuracy`\\n\\n        Grouped Query Attention (GQA) is an experimental feature. It currently works only for Flash_attention\\n        and math kernel on CUDA tensor, and does not support Nested tensor.\\n        Constraints for GQA:\\n\\n            - number_of_heads_query % number_of_heads_key_value == 0 and,\\n            - number_of_heads_key == number_of_heads_value\\n\\n    Note:\\n\\n        {cudnn_reproducibility_note}\\n    \\\"\\\"\\\".format(\\n        **reproducibility_notes\\n    )\\n    + r\\\"\\\"\\\"\\n    Args:\\n        query (Tensor): Query tensor; shape :math:`(N, ..., Hq, L, E)`.\\n        key (Tensor): Key tensor; shape :math:`(N, ..., H, S, E)`.\\n        value (Tensor): Value tensor; shape :math:`(N, ..., H, S, Ev)`.\\n        attn_mask (optional Tensor): Attention mask; shape must be broadcastable to the shape of attention weights,\\n            which is :math:`(N,..., L, S)`. Two types of masks are supported.\\n            A boolean mask where a value of True indicates that the element *should* take part in attention.\\n            A float mask of the same type as query, key, value that is added to the attention score.\\n        dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied\\n        is_causal (bool): If set to true, the attention masking is a lower triangular matrix when the mask is a\\n            square matrix. The attention masking has the form of the upper left causal bias due to the alignment\\n            (see :class:`torch.nn.attention.bias.CausalBias`) when the mask is a non-square matrix.\\n            An error is thrown if both attn_mask and is_causal are set.\\n        scale (optional float, keyword-only): Scaling factor applied prior to softmax. If None, the default value is set\\n            to :math:`\\\\frac{1}{\\\\sqrt{E}}`.\\n        enable_gqa (bool): If set to True, Grouped Query Attention (GQA) is enabled, by default it is set to False.\\n\\n    Returns:\\n        output (Tensor): Attention output; shape :math:`(N, ..., Hq, L, Ev)`.\\n\\n    Shape legend:\\n        - :math:`N: \\\\text{Batch size} ... : \\\\text{Any number of other batch dimensions (optional)}`\\n        - :math:`S: \\\\text{Source sequence length}`\\n        - :math:`L: \\\\text{Target sequence length}`\\n        - :math:`E: \\\\text{Embedding dimension of the query and key}`\\n        - :math:`Ev: \\\\text{Embedding dimension of the value}`\\n        - :math:`Hq: \\\\text{Number of heads of query}`\\n        - :math:`H: \\\\text{Number of heads of key and value}`\\n\\n    Examples:\\n\\n        >>> # Optionally use the context manager to ensure one of the fused kernels is run\\n        >>> query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):\\n        >>>     F.scaled_dot_product_attention(query,key,value)\\n\\n\\n        >>> # Sample for GQA for llama3\\n        >>> query = torch.rand(32, 32, 128, 64, dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> with sdpa_kernel(backends=[SDPBackend.MATH]):\\n        >>>     F.scaled_dot_product_attention(query,key,value,enable_gqa=True)\\n\\n\\n    .. _FlashAttention-2\\\\: Faster Attention with Better Parallelism and Work Partitioning:\\n        https://arxiv.org/abs/2307.08691\\n    .. _Memory-Efficient Attention:\\n        https://github.com/facebookresearch/xformers\\n    .. _Grouped-Query Attention:\\n        https://arxiv.org/pdf/2305.13245\\n    \\\"\\\"\\\",\\n)\\n\\n\\ndef _mha_shape_check(\\n    query: Tensor,\\n    key: Tensor,\\n    value: Tensor,\\n    key_padding_mask: Optional[Tensor],\\n    attn_mask: Optional[Tensor],\\n    num_heads: int,\\n):\\n    # Verifies the expected shape for `query, `key`, `value`, `key_padding_mask` and `attn_mask`\\n    # and returns if the input is batched or not.\\n    # Raises an error if `query` is not 2-D (unbatched) or 3-D (batched) tensor.\\n\\n    # Shape check.\\n    if query.dim() == 3:\\n        # Batched Inputs\\n        is_batched = True\\n        assert key.dim() == 3 and value.dim() == 3, (\\n            \\\"For batched (3-D) `query`, expected `key` and `value` to be 3-D\\\"\\n            f\\\" but found {key.dim()}-D and {value.dim()}-D tensors respectively\\\"\\n        )\\n        if key_padding_mask is not None:\\n            assert key_padding_mask.dim() == 2, (\\n                \\\"For batched (3-D) `query`, expected `key_padding_mask` to be `None` or 2-D\\\"\\n                f\\\" but found {key_padding_mask.dim()}-D tensor instead\\\"\\n            )\\n        if attn_mask is not None:\\n            assert attn_mask.dim() in (2, 3), (\\n                \\\"For batched (3-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D\\\"\\n                f\\\" but found {attn_mask.dim()}-D tensor instead\\\"\\n            )\\n    elif query.dim() == 2:\\n        # Unbatched Inputs\\n        is_batched = False\\n        assert key.dim() == 2 and value.dim() == 2, (\\n            \\\"For unbatched (2-D) `query`, expected `key` and `value` to be 2-D\\\"\\n            f\\\" but found {key.dim()}-D and {value.dim()}-D tensors respectively\\\"\\n        )\\n\\n        if key_padding_mask is not None:\\n            assert key_padding_mask.dim() == 1, (\\n                \\\"For unbatched (2-D) `query`, expected `key_padding_mask` to be `None` or 1-D\\\"\\n                f\\\" but found {key_padding_mask.dim()}-D tensor instead\\\"\\n            )\\n\\n        if attn_mask is not None:\\n            assert attn_mask.dim() in (2, 3), (\\n                \\\"For unbatched (2-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D\\\"\\n                f\\\" but found {attn_mask.dim()}-D tensor instead\\\"\\n            )\\n            if attn_mask.dim() == 3:\\n                expected_shape = (num_heads, query.shape[0], key.shape[0])\\n                assert (\\n                    attn_mask.shape == expected_shape\\n                ), f\\\"Expected `attn_mask` shape to be {expected_shape} but got {attn_mask.shape}\\\"\\n    else:\\n        raise AssertionError(\\n            f\\\"query should be unbatched 2D or batched 3D tensor but received {query.dim()}-D query tensor\\\"\\n        )\\n\\n    return is_batched\\n\\n\\ndef _canonical_mask(\\n    mask: Optional[Tensor],\\n    mask_name: str,\\n    other_type: Optional[DType],\\n    other_name: str,\\n    target_type: DType,\\n    check_other: bool = True,\\n) -> Optional[Tensor]:\\n    if mask is not None:\\n        _mask_dtype = mask.dtype\\n        _mask_is_float = torch.is_floating_point(mask)\\n        if _mask_dtype != torch.bool and not _mask_is_float:\\n            raise AssertionError(\\n                f\\\"only bool and floating types of {mask_name} are supported\\\"\\n            )\\n        if check_other and other_type is not None:\\n            if _mask_dtype != other_type:\\n                warnings.warn(\\n                    f\\\"Support for mismatched {mask_name} and {other_name} \\\"\\n                    \\\"is deprecated. Use same type for both instead.\\\"\\n                )\\n        if not _mask_is_float:\\n            mask = torch.zeros_like(mask, dtype=target_type).masked_fill_(\\n                mask, float(\\\"-inf\\\")\\n            )\\n    return mask\\n\\n\\ndef _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:\\n    if input is None:\\n        return None\\n    elif isinstance(input, torch.Tensor):\\n        return input.dtype\\n    raise RuntimeError(\\\"input to _none_or_dtype() must be None or torch.Tensor\\\")\\n\\n\\ndef multi_head_attention_forward(\\n    query: Tensor,\\n    key: Tensor,\\n    value: Tensor,\\n    embed_dim_to_check: int,\\n    num_heads: int,\\n    in_proj_weight: Optional[Tensor],\\n    in_proj_bias: Optional[Tensor],\\n    bias_k: Optional[Tensor],\\n    bias_v: Optional[Tensor],\\n    add_zero_attn: bool,\\n    dropout_p: float,\\n    out_proj_weight: Tensor,\\n    out_proj_bias: Optional[Tensor],\\n    training: bool = True,\\n    key_padding_mask: Optional[Tensor] = None,\\n    need_weights: bool = True,\\n    attn_mask: Optional[Tensor] = None,\\n    use_separate_proj_weight: bool = False,\\n    q_proj_weight: Optional[Tensor] = None,\\n    k_proj_weight: Optional[Tensor] = None,\\n    v_proj_weight: Optional[Tensor] = None,\\n    static_k: Optional[Tensor] = None,\\n    static_v: Optional[Tensor] = None,\\n    average_attn_weights: bool = True,\\n    is_causal: bool = False,\\n) -> Tuple[Tensor, Optional[Tensor]]:\\n    r\\\"\\\"\\\"Forward method for MultiHeadAttention.\\n\\n    See :class:`torch.nn.MultiheadAttention` for details.\\n\\n    Args:\\n        query, key, value: map a query and a set of key-value pairs to an output.\\n            See \\\"Attention Is All You Need\\\" for more details.\\n        embed_dim_to_check: total dimension of the model.\\n        num_heads: parallel attention heads.\\n        in_proj_weight, in_proj_bias: input projection weight and bias.\\n        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.\\n        add_zero_attn: add a new batch of zeros to the key and\\n                       value sequences at dim=1.\\n        dropout_p: probability of an element to be zeroed.\\n        out_proj_weight, out_proj_bias: the output projection weight and bias.\\n        training: apply dropout if is ``True``.\\n        key_padding_mask: if provided, specified padding elements in the key will\\n            be ignored by the attention. This is an binary mask. When the value is True,\\n            the corresponding value on the attention layer will be filled with -inf.\\n        need_weights: output attn_output_weights.\\n            Default: `True`\\n            Note: `needs_weight` defaults to `True`, but should be set to `False`\\n            For best performance when attention weights are not needed.\\n            *Setting needs_weights to `True`\\n            leads to a significant performance degradation.*\\n        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all\\n            the batches while a 3D mask allows to specify a different mask for the entries of each batch.\\n        is_causal: If specified, applies a causal mask as attention mask, and ignores\\n            attn_mask for computing scaled dot product attention.\\n            Default: ``False``.\\n            .. warning::\\n                is_causal is provides a hint that the attn_mask is the\\n                causal mask.Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n        use_separate_proj_weight: the function accept the proj. weights for query, key,\\n            and value in different forms. If false, in_proj_weight will be used, which is\\n            a combination of q_proj_weight, k_proj_weight, v_proj_weight.\\n        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.\\n        static_k, static_v: static key and value used for attention operators.\\n        average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across heads.\\n            Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an effect\\n            when ``need_weights=True.``. Default: True\\n\\n\\n    Shape:\\n        Inputs:\\n        - query: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is\\n          the embedding dimension.\\n        - key: :math:`(S, E)` or :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is\\n          the embedding dimension.\\n        - value: :math:`(S, E)` or :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is\\n          the embedding dimension.\\n        - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.\\n          If a FloatTensor is provided, it will be directly added to the value.\\n          If a BoolTensor is provided, the positions with the\\n          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.\\n        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.\\n          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,\\n          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked\\n          positions. If a BoolTensor is provided, positions with ``True``\\n          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor\\n          is provided, it will be added to the attention weight.\\n        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,\\n          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.\\n        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,\\n          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.\\n\\n        Outputs:\\n        - attn_output: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size,\\n          E is the embedding dimension.\\n        - attn_output_weights: Only returned when ``need_weights=True``. If ``average_attn_weights=True``, returns\\n          attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or\\n          :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and\\n          :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per\\n          head of shape :math:`(num_heads, L, S)` when input is unbatched or :math:`(N, num_heads, L, S)`.\\n    \\\"\\\"\\\"\\n    tens_ops = (\\n        query,\\n        key,\\n        value,\\n        in_proj_weight,\\n        in_proj_bias,\\n        bias_k,\\n        bias_v,\\n        out_proj_weight,\\n        out_proj_bias,\\n    )\\n    if has_torch_function(tens_ops):\\n        return handle_torch_function(\\n            multi_head_attention_forward,\\n            tens_ops,\\n            query,\\n            key,\\n            value,\\n            embed_dim_to_check,\\n            num_heads,\\n            in_proj_weight,\\n            in_proj_bias,\\n            bias_k,\\n            bias_v,\\n            add_zero_attn,\\n            dropout_p,\\n            out_proj_weight,\\n            out_proj_bias,\\n            training=training,\\n            key_padding_mask=key_padding_mask,\\n            need_weights=need_weights,\\n            attn_mask=attn_mask,\\n            is_causal=is_causal,\\n            use_separate_proj_weight=use_separate_proj_weight,\\n            q_proj_weight=q_proj_weight,\\n            k_proj_weight=k_proj_weight,\\n            v_proj_weight=v_proj_weight,\\n            static_k=static_k,\\n            static_v=static_v,\\n            average_attn_weights=average_attn_weights,\\n        )\\n\\n    is_batched = _mha_shape_check(\\n        query, key, value, key_padding_mask, attn_mask, num_heads\\n    )\\n\\n    # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input\\n    # is batched, run the computation and before returning squeeze the\\n    # batch dimension so that the output doesn't carry this temporary batch dimension.\\n    if not is_batched:\\n        # unsqueeze if the input is unbatched\\n        query = query.unsqueeze(1)\\n        key = key.unsqueeze(1)\\n        value = value.unsqueeze(1)\\n        if key_padding_mask is not None:\\n            key_padding_mask = key_padding_mask.unsqueeze(0)\\n\\n    # set up shape vars\\n    tgt_len, bsz, embed_dim = query.shape\\n    src_len, _, _ = key.shape\\n\\n    key_padding_mask = _canonical_mask(\\n        mask=key_padding_mask,\\n        mask_name=\\\"key_padding_mask\\\",\\n        other_type=_none_or_dtype(attn_mask),\\n        other_name=\\\"attn_mask\\\",\\n        target_type=query.dtype,\\n    )\\n\\n    if is_causal and attn_mask is None:\\n        raise RuntimeError(\\n            \\\"Need attn_mask if specifying the is_causal hint. \\\"\\n            \\\"You may use the Transformer module method \\\"\\n            \\\"`generate_square_subsequent_mask` to create this mask.\\\"\\n        )\\n\\n    if is_causal and key_padding_mask is None and not need_weights:\\n        # when we have a kpm or need weights, we need attn_mask\\n        # Otherwise, we use the is_causal hint go as is_causal\\n        # indicator to SDPA.\\n        attn_mask = None\\n    else:\\n        attn_mask = _canonical_mask(\\n            mask=attn_mask,\\n            mask_name=\\\"attn_mask\\\",\\n            other_type=None,\\n            other_name=\\\"\\\",\\n            target_type=query.dtype,\\n            check_other=False,\\n        )\\n\\n        if key_padding_mask is not None:\\n            # We have the attn_mask, and use that to merge kpm into it.\\n            # Turn off use of is_causal hint, as the merged mask is no\\n            # longer causal.\\n            is_causal = False\\n\\n    assert (\\n        embed_dim == embed_dim_to_check\\n    ), f\\\"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}\\\"\\n    if isinstance(embed_dim, torch.Tensor):\\n        # embed_dim can be a tensor when JIT tracing\\n        head_dim = embed_dim.div(num_heads, rounding_mode=\\\"trunc\\\")\\n    else:\\n        head_dim = embed_dim // num_heads\\n    assert (\\n        head_dim * num_heads == embed_dim\\n    ), f\\\"embed_dim {embed_dim} not divisible by num_heads {num_heads}\\\"\\n    if use_separate_proj_weight:\\n        # allow MHA to have different embedding dimensions when separate projection weights are used\\n        assert (\\n            key.shape[:2] == value.shape[:2]\\n        ), f\\\"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}\\\"\\n    else:\\n        assert (\\n            key.shape == value.shape\\n        ), f\\\"key shape {key.shape} does not match value shape {value.shape}\\\"\\n\\n    #\\n    # compute in-projection\\n    #\\n    if not use_separate_proj_weight:\\n        assert (\\n            in_proj_weight is not None\\n        ), \\\"use_separate_proj_weight is False but in_proj_weight is None\\\"\\n        q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)\\n    else:\\n        assert (\\n            q_proj_weight is not None\\n        ), \\\"use_separate_proj_weight is True but q_proj_weight is None\\\"\\n        assert (\\n            k_proj_weight is not None\\n        ), \\\"use_separate_proj_weight is True but k_proj_weight is None\\\"\\n        assert (\\n            v_proj_weight is not None\\n        ), \\\"use_separate_proj_weight is True but v_proj_weight is None\\\"\\n        if in_proj_bias is None:\\n            b_q = b_k = b_v = None\\n        else:\\n            b_q, b_k, b_v = in_proj_bias.chunk(3)\\n        q, k, v = _in_projection(\\n            query,\\n            key,\\n            value,\\n            q_proj_weight,\\n            k_proj_weight,\\n            v_proj_weight,\\n            b_q,\\n            b_k,\\n            b_v,\\n        )\\n\\n    # prep attention mask\\n\\n    if attn_mask is not None:\\n        # ensure attn_mask's dim is 3\\n        if attn_mask.dim() == 2:\\n            correct_2d_size = (tgt_len, src_len)\\n            if attn_mask.shape != correct_2d_size:\\n                raise RuntimeError(\\n                    f\\\"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.\\\"\\n                )\\n            attn_mask = attn_mask.unsqueeze(0)\\n        elif attn_mask.dim() == 3:\\n            correct_3d_size = (bsz * num_heads, tgt_len, src_len)\\n            if attn_mask.shape != correct_3d_size:\\n                raise RuntimeError(\\n                    f\\\"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.\\\"\\n                )\\n        else:\\n            raise RuntimeError(\\n                f\\\"attn_mask's dimension {attn_mask.dim()} is not supported\\\"\\n            )\\n\\n    # add bias along batch dimension (currently second)\\n    if bias_k is not None and bias_v is not None:\\n        assert static_k is None, \\\"bias cannot be added to static key.\\\"\\n        assert static_v is None, \\\"bias cannot be added to static value.\\\"\\n        k = torch.cat([k, bias_k.repeat(1, bsz, 1)])\\n        v = torch.cat([v, bias_v.repeat(1, bsz, 1)])\\n        if attn_mask is not None:\\n            attn_mask = pad(attn_mask, (0, 1))\\n        if key_padding_mask is not None:\\n            key_padding_mask = pad(key_padding_mask, (0, 1))\\n    else:\\n        assert bias_k is None\\n        assert bias_v is None\\n\\n    #\\n    # reshape q, k, v for multihead attention and make them batch first\\n    #\\n    q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\\n    if static_k is None:\\n        k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)\\n    else:\\n        # TODO finish disentangling control flow so we don't do in-projections when statics are passed\\n        assert (\\n            static_k.size(0) == bsz * num_heads\\n        ), f\\\"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}\\\"\\n        assert (\\n            static_k.size(2) == head_dim\\n        ), f\\\"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}\\\"\\n        k = static_k\\n    if static_v is None:\\n        v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)\\n    else:\\n        # TODO finish disentangling control flow so we don't do in-projections when statics are passed\\n        assert (\\n            static_v.size(0) == bsz * num_heads\\n        ), f\\\"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}\\\"\\n        assert (\\n            static_v.size(2) == head_dim\\n        ), f\\\"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}\\\"\\n        v = static_v\\n\\n    # add zero attention along batch dimension (now first)\\n    if add_zero_attn:\\n        zero_attn_shape = (bsz * num_heads, 1, head_dim)\\n        k = torch.cat(\\n            [k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1\\n        )\\n        v = torch.cat(\\n            [v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1\\n        )\\n        if attn_mask is not None:\\n            attn_mask = pad(attn_mask, (0, 1))\\n        if key_padding_mask is not None:\\n            key_padding_mask = pad(key_padding_mask, (0, 1))\\n\\n    # update source sequence length after adjustments\\n    src_len = k.size(1)\\n\\n    # merge key padding and attention masks\\n    if key_padding_mask is not None:\\n        assert key_padding_mask.shape == (\\n            bsz,\\n            src_len,\\n        ), f\\\"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}\\\"\\n        key_padding_mask = (\\n            key_padding_mask.view(bsz, 1, 1, src_len)\\n            .expand(-1, num_heads, -1, -1)\\n            .reshape(bsz * num_heads, 1, src_len)\\n        )\\n        if attn_mask is None:\\n            attn_mask = key_padding_mask\\n        else:\\n            attn_mask = attn_mask + key_padding_mask\\n\\n    # adjust dropout probability\\n    if not training:\\n        dropout_p = 0.0\\n\\n    #\\n    # (deep breath) calculate attention and out projection\\n    #\\n\\n    if need_weights:\\n        B, Nt, E = q.shape\\n        q_scaled = q * math.sqrt(1.0 / float(E))\\n\\n        assert not (\\n            is_causal and attn_mask is None\\n        ), \\\"FIXME: is_causal not implemented for need_weights\\\"\\n\\n        if attn_mask is not None:\\n            attn_output_weights = torch.baddbmm(\\n                attn_mask, q_scaled, k.transpose(-2, -1)\\n            )\\n        else:\\n            attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))\\n        attn_output_weights = softmax(attn_output_weights, dim=-1)\\n        if dropout_p > 0.0:\\n            attn_output_weights = dropout(attn_output_weights, p=dropout_p)\\n\\n        attn_output = torch.bmm(attn_output_weights, v)\\n\\n        attn_output = (\\n            attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)\\n        )\\n        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)\\n        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))\\n\\n        # optionally average attention weights over heads\\n        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)\\n        if average_attn_weights:\\n            attn_output_weights = attn_output_weights.mean(dim=1)\\n\\n        if not is_batched:\\n            # squeeze the output if input was unbatched\\n            attn_output = attn_output.squeeze(1)\\n            attn_output_weights = attn_output_weights.squeeze(0)\\n        return attn_output, attn_output_weights\\n    else:\\n        # attn_mask can be either (L,S) or (N*num_heads, L, S)\\n        # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)\\n        # in order to match the input for SDPA of (N, num_heads, L, S)\\n        if attn_mask is not None:\\n            if attn_mask.size(0) == 1 and attn_mask.dim() == 3:\\n                attn_mask = attn_mask.unsqueeze(0)\\n            else:\\n                attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)\\n\\n        q = q.view(bsz, num_heads, tgt_len, head_dim)\\n        k = k.view(bsz, num_heads, src_len, head_dim)\\n        v = v.view(bsz, num_heads, src_len, head_dim)\\n\\n        attn_output = scaled_dot_product_attention(\\n            q, k, v, attn_mask, dropout_p, is_causal\\n        )\\n        attn_output = (\\n            attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)\\n        )\\n\\n        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)\\n        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))\\n        if not is_batched:\\n            # squeeze the output if input was unbatched\\n            attn_output = attn_output.squeeze(1)\\n        return attn_output, None\\n\\n\\nfrom collections import OrderedDict\\n\\nimport torch\\nfrom torch._C import _disabled_torch_function_impl\\n\\n\\n# Metaclass to combine _TensorMeta and the instance check override for Parameter.\\nclass _ParameterMeta(torch._C._TensorMeta):\\n    # Make `isinstance(t, Parameter)` return True for custom tensor instances that have the _is_param flag.\\n    def __instancecheck__(self, instance):\\n        if self is Parameter:\\n            if isinstance(instance, torch.Tensor) and getattr(\\n                instance, \\\"_is_param\\\", False\\n            ):\\n                return True\\n        return super().__instancecheck__(instance)\\n\\n\\nclass Parameter(torch.Tensor, metaclass=_ParameterMeta):\\n    r\\\"\\\"\\\"A kind of Tensor that is to be considered a module parameter.\\n\\n    Parameters are :class:`~torch.Tensor` subclasses, that have a\\n    very special property when used with :class:`Module` s - when they're\\n    assigned as Module attributes they are automatically added to the list of\\n    its parameters, and will appear e.g. in :meth:`~Module.parameters` iterator.\\n    Assigning a Tensor doesn't have such effect. This is because one might\\n    want to cache some temporary state, like last hidden state of the RNN, in\\n    the model. If there was no such class as :class:`Parameter`, these\\n    temporaries would get registered too.\\n\\n    Args:\\n        data (Tensor): parameter tensor.\\n        requires_grad (bool, optional): if the parameter requires gradient. Note that\\n            the torch.no_grad() context does NOT affect the default behavior of\\n            Parameter creation--the Parameter will still have `requires_grad=True` in\\n            :class:`~no_grad` mode. See :ref:`locally-disable-grad-doc` for more\\n            details. Default: `True`\\n    \\\"\\\"\\\"\\n\\n    def __new__(cls, data=None, requires_grad=True):\\n        if data is None:\\n            data = torch.empty(0)\\n        if type(data) is torch.Tensor or type(data) is Parameter:\\n            # For ease of BC maintenance, keep this path for standard Tensor.\\n            # Eventually (tm), we should change the behavior for standard Tensor to match.\\n            return torch.Tensor._make_subclass(cls, data, requires_grad)\\n\\n        # Path for custom tensors: set a flag on the instance to indicate parameter-ness.\\n        t = data.detach().requires_grad_(requires_grad)\\n        if type(t) is not type(data):\\n            raise RuntimeError(\\n                f\\\"Creating a Parameter from an instance of type {type(data).__name__} \\\"\\n                \\\"requires that detach() returns an instance of the same type, but return \\\"\\n                f\\\"type {type(t).__name__} was found instead. To use the type as a \\\"\\n                \\\"Parameter, please correct the detach() semantics defined by \\\"\\n                \\\"its __torch_dispatch__() implementation.\\\"\\n            )\\n        t._is_param = True\\n        return t\\n\\n    # Note: the 3 methods below only apply to standard Tensor. Parameters of custom tensor types\\n    # are still considered that custom tensor type and these methods will not be called for them.\\n    def __deepcopy__(self, memo):\\n        if id(self) in memo:\\n            return memo[id(self)]\\n        else:\\n            result = type(self)(\\n                self.data.clone(memory_format=torch.preserve_format), self.requires_grad\\n            )\\n            memo[id(self)] = result\\n            return result\\n\\n    def __repr__(self):\\n        return \\\"Parameter containing:\\\\n\\\" + super().__repr__()\\n\\n    def __reduce_ex__(self, proto):\\n        state = torch._utils._get_obj_state(self)\\n\\n        # See Note [Don't serialize hooks]\\n        hooks = OrderedDict()\\n        if not state:\\n            return (\\n                torch._utils._rebuild_parameter,\\n                (self.data, self.requires_grad, hooks),\\n            )\\n\\n        return (\\n            torch._utils._rebuild_parameter_with_state,\\n            (self.data, self.requires_grad, hooks, state),\\n        )\\n\\n    __torch_function__ = _disabled_torch_function_impl\\n\\n\\nclass UninitializedTensorMixin:\\n    _allowed_methods = [\\n        torch.Tensor.__hash__,\\n        torch.Tensor.size,\\n        torch.Tensor.copy_,\\n        torch.Tensor.is_complex,\\n        torch.Tensor.is_floating_point,\\n        torch.Tensor.half,\\n        torch.Tensor.float,\\n        torch.Tensor.double,\\n        torch.Tensor.char,\\n        torch.Tensor.short,\\n        torch.Tensor.int,\\n        torch.Tensor.long,\\n        torch.Tensor.cuda,\\n        torch.Tensor.cpu,\\n        torch.Tensor.to,\\n        torch.Tensor.get_device,\\n        torch._has_compatible_shallow_copy_type,\\n    ]\\n\\n    def materialize(self, shape, device=None, dtype=None):\\n        r\\\"\\\"\\\"Create a Parameter or Tensor with the same properties of the uninitialized one.\\n\\n        Given a shape, it materializes a parameter in the same device\\n        and with the same `dtype` as the current one or the specified ones in the\\n        arguments.\\n\\n        Args:\\n            shape : (tuple): the shape for the materialized tensor.\\n            device (:class:`torch.device`): the desired device of the parameters\\n                and buffers in this module. Optional.\\n            dtype (:class:`torch.dtype`): the desired floating point type of\\n                the floating point parameters and buffers in this module. Optional.\\n        \\\"\\\"\\\"\\n        if device is None:\\n            device = self.data.device\\n        if dtype is None:\\n            dtype = self.data.dtype\\n        self.data = torch.empty(shape, device=device, dtype=dtype)\\n        self.__class__ = self.cls_to_become\\n\\n    @property\\n    def shape(self):\\n        raise RuntimeError(\\n            \\\"Can't access the shape of an uninitialized parameter or buffer. \\\"\\n            \\\"This error usually happens in `load_state_dict` when trying to load \\\"\\n            \\\"an uninitialized parameter into an initialized one. \\\"\\n            \\\"Call `forward` to initialize the parameters before accessing their attributes.\\\"\\n        )\\n\\n    def share_memory_(self):\\n        raise RuntimeError(\\n            \\\"Can't share memory on an uninitialized parameter or buffer. \\\"\\n            \\\"Call `forward` to initialize the parameters before calling \\\"\\n            \\\"`module.share_memory()`.\\\"\\n        )\\n\\n    def __repr__(self):\\n        return f\\\"<{self.__class__.__name__}>\\\"\\n\\n    def __reduce_ex__(self, proto):\\n        # See Note [Don't serialize hooks]\\n        return (self.__class__, (self.requires_grad,))\\n\\n    @classmethod\\n    def __torch_function__(cls, func, types, args=(), kwargs=None):\\n        # method-wrapper is to detect access to Tensor properties that are\\n        # wrapped in descriptors\\n        if func in cls._allowed_methods or func.__class__.__name__ == \\\"method-wrapper\\\":\\n            if kwargs is None:\\n                kwargs = {}\\n            return super().__torch_function__(func, types, args, kwargs)\\n        raise ValueError(\\n            f\\\"Attempted to use an uninitialized parameter in {func}. \\\"\\n            \\\"This error happens when you are using a `LazyModule` or \\\"\\n            f\\\"explicitly manipulating `torch.nn.parameter.{cls.__name__}` \\\"\\n            \\\"objects. When using LazyModules Call `forward` with a dummy batch \\\"\\n            \\\"to initialize the parameters before calling torch functions\\\"\\n        )\\n\\n\\ndef is_lazy(param):\\n    return isinstance(param, UninitializedTensorMixin)\\n\\n\\nclass UninitializedParameter(UninitializedTensorMixin, Parameter):\\n    r\\\"\\\"\\\"A parameter that is not initialized.\\n\\n    Uninitialized Parameters are a a special case of :class:`torch.nn.Parameter`\\n    where the shape of the data is still unknown.\\n\\n    Unlike a :class:`torch.nn.Parameter`, uninitialized parameters\\n    hold no data and attempting to access some properties, like their shape,\\n    will throw a runtime error. The only operations that can be performed on a uninitialized\\n    parameter are changing its datatype, moving it to a different device and\\n    converting it to a regular :class:`torch.nn.Parameter`.\\n\\n    The default device or dtype to use when the parameter is materialized can be set\\n    during construction using e.g. ``device='cuda'``.\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = Parameter\\n\\n    def __new__(cls, requires_grad=True, device=None, dtype=None) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        data = torch.empty(0, **factory_kwargs)\\n        return torch.Tensor._make_subclass(cls, data, requires_grad)\\n\\n    def __deepcopy__(self, memo):\\n        if id(self) in memo:\\n            return memo[id(self)]\\n        else:\\n            result = type(self)(self.requires_grad, self.data.device, self.data.dtype)\\n            memo[id(self)] = result\\n            return result\\n\\n\\n# Metaclass to combine _TensorMeta and the instance check override for Buffer.\\nclass _BufferMeta(torch._C._TensorMeta):\\n    # Make `isinstance(t, Buffer)` return True for custom tensor instances that have the _is_buffer flag.\\n    def __instancecheck__(self, instance):\\n        if self is Buffer:\\n            if isinstance(instance, torch.Tensor) and getattr(\\n                instance, \\\"_is_buffer\\\", False\\n            ):\\n                return True\\n        return super().__instancecheck__(instance)\\n\\n\\nclass Buffer(torch.Tensor, metaclass=_BufferMeta):\\n    r\\\"\\\"\\\"A kind of Tensor that should not be considered a model\\n    parameter. For example, BatchNorm's ``running_mean`` is not a parameter, but is part of the module's state.\\n\\n    Buffers are :class:`~torch.Tensor` subclasses, that have a\\n    very special property when used with :class:`Module` s -- when they're\\n    assigned as Module attributes they are automatically added to the list of\\n    its buffers, and will appear e.g. in :meth:`~torch.nn.Module.buffers` iterator.\\n    Assigning a Tensor doesn't have such effect. One can still assign a Tensor as explicitly by using\\n    the :meth:`~torch.nn.Module.register_buffer` function.\\n\\n    Args:\\n        data (Tensor): buffer tensor.\\n        persistent (bool, optional): whether the buffer is part of the module's\\n            :attr:`state_dict`. Default: ``True``\\n    \\\"\\\"\\\"\\n\\n    def __new__(cls, data=None, *, persistent=True):\\n        if data is None:\\n            data = torch.empty(0)\\n\\n        t = data.detach().requires_grad_(data.requires_grad)\\n        t.persistent = persistent\\n        t._is_buffer = True\\n        return t\\n\\n    __torch_function__ = _disabled_torch_function_impl\\n\\n\\nclass UninitializedBuffer(UninitializedTensorMixin, torch.Tensor):\\n    r\\\"\\\"\\\"A buffer that is not initialized.\\n\\n    Uninitialized Buffer is a a special case of :class:`torch.Tensor`\\n    where the shape of the data is still unknown.\\n\\n    Unlike a :class:`torch.Tensor`, uninitialized parameters\\n    hold no data and attempting to access some properties, like their shape,\\n    will throw a runtime error. The only operations that can be performed on a uninitialized\\n    parameter are changing its datatype, moving it to a different device and\\n    converting it to a regular :class:`torch.Tensor`.\\n\\n    The default device or dtype to use when the buffer is materialized can be set\\n    during construction using e.g. ``device='cuda'``.\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = torch.Tensor\\n\\n    def __new__(\\n        cls, requires_grad=False, device=None, dtype=None, persistent=True\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        data = torch.empty(0, **factory_kwargs)\\n        ret = torch.Tensor._make_subclass(cls, data, requires_grad)\\n        ret.persistent = persistent\\n        ret._is_buffer = True\\n        return ret\\n\\n\\nfrom typing import Optional, Tuple, TypeVar, Union\\n\\nfrom torch import Tensor\\n\\n\\n# Create some useful type aliases\\n\\n# Template for arguments which can be supplied as a tuple, or which can be a scalar which PyTorch will internally\\n# broadcast to a tuple.\\n# Comes in several variants: A tuple of unknown size, and a fixed-size tuple for 1d, 2d, or 3d operations.\\nT = TypeVar(\\\"T\\\")\\n_scalar_or_tuple_any_t = Union[T, Tuple[T, ...]]\\n_scalar_or_tuple_1_t = Union[T, Tuple[T]]\\n_scalar_or_tuple_2_t = Union[T, Tuple[T, T]]\\n_scalar_or_tuple_3_t = Union[T, Tuple[T, T, T]]\\n_scalar_or_tuple_4_t = Union[T, Tuple[T, T, T, T]]\\n_scalar_or_tuple_5_t = Union[T, Tuple[T, T, T, T, T]]\\n_scalar_or_tuple_6_t = Union[T, Tuple[T, T, T, T, T, T]]\\n\\n# For arguments which represent size parameters (eg, kernel size, padding)\\n_size_any_t = _scalar_or_tuple_any_t[int]\\n_size_1_t = _scalar_or_tuple_1_t[int]\\n_size_2_t = _scalar_or_tuple_2_t[int]\\n_size_3_t = _scalar_or_tuple_3_t[int]\\n_size_4_t = _scalar_or_tuple_4_t[int]\\n_size_5_t = _scalar_or_tuple_5_t[int]\\n_size_6_t = _scalar_or_tuple_6_t[int]\\n\\n# For arguments which represent optional size parameters (eg, adaptive pool parameters)\\n_size_any_opt_t = _scalar_or_tuple_any_t[Optional[int]]\\n_size_2_opt_t = _scalar_or_tuple_2_t[Optional[int]]\\n_size_3_opt_t = _scalar_or_tuple_3_t[Optional[int]]\\n\\n# For arguments that represent a ratio to adjust each dimension of an input with (eg, upsampling parameters)\\n_ratio_2_t = _scalar_or_tuple_2_t[float]\\n_ratio_3_t = _scalar_or_tuple_3_t[float]\\n_ratio_any_t = _scalar_or_tuple_any_t[float]\\n\\n_tensor_list_t = _scalar_or_tuple_any_t[Tensor]\\n\\n# For the return value of max pooling operations that may or may not return indices.\\n# With the proposed 'Literal' feature to Python typing, it might be possible to\\n# eventually eliminate this.\\n_maybe_indices_t = _scalar_or_tuple_2_t[Tensor]\\n\\n\\n# mypy: allow-untyped-defs\\nfrom torch.nn.parameter import (  # usort: skip\\n    Buffer as Buffer,\\n    Parameter as Parameter,\\n    UninitializedBuffer as UninitializedBuffer,\\n    UninitializedParameter as UninitializedParameter,\\n)\\nfrom torch.nn.modules import *  # usort: skip # noqa: F403\\nfrom torch.nn import (\\n    attention as attention,\\n    functional as functional,\\n    init as init,\\n    modules as modules,\\n    parallel as parallel,\\n    parameter as parameter,\\n    utils as utils,\\n)\\nfrom torch.nn.parallel import DataParallel as DataParallel\\n\\n\\ndef factory_kwargs(kwargs):\\n    r\\\"\\\"\\\"Return a canonicalized dict of factory kwargs.\\n\\n    Given kwargs, returns a canonicalized dict of factory kwargs that can be directly passed\\n    to factory functions like torch.empty, or errors if unrecognized kwargs are present.\\n\\n    This function makes it simple to write code like this::\\n\\n        class MyModule(nn.Module):\\n            def __init__(self, **kwargs):\\n                factory_kwargs = torch.nn.factory_kwargs(kwargs)\\n                self.weight = Parameter(torch.empty(10, **factory_kwargs))\\n\\n    Why should you use this function instead of just passing `kwargs` along directly?\\n\\n    1. This function does error validation, so if there are unexpected kwargs we will\\n    immediately report an error, instead of deferring it to the factory call\\n    2. This function supports a special `factory_kwargs` argument, which can be used to\\n    explicitly specify a kwarg to be used for factory functions, in the event one of the\\n    factory kwargs conflicts with an already existing argument in the signature (e.g.\\n    in the signature ``def f(dtype, **kwargs)``, you can specify ``dtype`` for factory\\n    functions, as distinct from the dtype argument, by saying\\n    ``f(dtype1, factory_kwargs={\\\"dtype\\\": dtype2})``)\\n    \\\"\\\"\\\"\\n    if kwargs is None:\\n        return {}\\n    simple_keys = {\\\"device\\\", \\\"dtype\\\", \\\"memory_format\\\"}\\n    expected_keys = simple_keys | {\\\"factory_kwargs\\\"}\\n    if not kwargs.keys() <= expected_keys:\\n        raise TypeError(f\\\"unexpected kwargs {kwargs.keys() - expected_keys}\\\")\\n\\n    # guarantee no input kwargs is untouched\\n    r = dict(kwargs.get(\\\"factory_kwargs\\\", {}))\\n    for k in simple_keys:\\n        if k in kwargs:\\n            if k in r:\\n                raise TypeError(\\n                    f\\\"{k} specified twice, in **kwargs and in factory_kwargs\\\"\\n                )\\n            r[k] = kwargs[k]\\n\\n    return r\\n\\n\\n# mypy: allow-untyped-defs\\n\\\"\\\"\\\"Gradient interface.\\\"\\\"\\\"\\n\\nimport torch\\nfrom torch.nn.modules.utils import _pair, _single, _triple\\n\\n\\ndef conv1d_input(\\n    input_size,\\n    weight,\\n    grad_output,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    r\\\"\\\"\\\"Compute the gradient of conv1d with respect to the input of the convolution.\\n\\n    This is same as the 1D transposed convolution operator under the hood but requires\\n    the shape of the gradient w.r.t. input to be specified explicitly.\\n\\n    Args:\\n        input_size : Shape of the input gradient tensor\\n        weight: weight tensor (out_channels x in_channels/groups x kW)\\n        grad_output : output gradient tensor (minibatch x out_channels x oW)\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n\\n    Examples::\\n\\n        >>> input = torch.randn(1, 1, 3, requires_grad=True)\\n        >>> weight = torch.randn(1, 1, 1, requires_grad=True)\\n        >>> output = F.conv1d(input, weight)\\n        >>> grad_output = torch.randn(output.shape)\\n        >>> grad_input = torch.autograd.grad(output, input, grad_output)\\n        >>> F.grad.conv1d_input(input.shape, weight, grad_output)\\n\\n    \\\"\\\"\\\"\\n    input = grad_output.new_empty(1).expand(input_size)\\n\\n    return torch.ops.aten.convolution_backward(\\n        grad_output,\\n        input,\\n        weight,\\n        None,\\n        _single(stride),\\n        _single(padding),\\n        _single(dilation),\\n        False,\\n        [0],\\n        groups,\\n        (True, False, False),\\n    )[0]\\n\\n\\ndef conv1d_weight(\\n    input,\\n    weight_size,\\n    grad_output,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    r\\\"\\\"\\\"Compute the gradient of conv1d with respect to the weight of the convolution.\\n\\n    Args:\\n        input: input tensor of shape (minibatch x in_channels x iW)\\n        weight_size : Shape of the weight gradient tensor\\n        grad_output : output gradient tensor (minibatch x out_channels x oW)\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n\\n    Examples::\\n\\n        >>> input = torch.randn(1, 1, 3, requires_grad=True)\\n        >>> weight = torch.randn(1, 1, 1, requires_grad=True)\\n        >>> output = F.conv1d(input, weight)\\n        >>> grad_output = torch.randn(output.shape)\\n        >>> # xdoctest: +SKIP\\n        >>> grad_weight = torch.autograd.grad(output, filter, grad_output)\\n        >>> F.grad.conv1d_weight(input, weight.shape, grad_output)\\n\\n    \\\"\\\"\\\"\\n    weight = grad_output.new_empty(1).expand(weight_size)\\n\\n    return torch.ops.aten.convolution_backward(\\n        grad_output,\\n        input,\\n        weight,\\n        None,\\n        _single(stride),\\n        _single(padding),\\n        _single(dilation),\\n        False,\\n        [0],\\n        groups,\\n        (False, True, False),\\n    )[1]\\n\\n\\ndef conv2d_input(\\n    input_size,\\n    weight,\\n    grad_output,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    r\\\"\\\"\\\"Compute the gradient of conv2d with respect to the input of the convolution.\\n\\n    This is same as the 2D transposed convolution operator under the hood but requires\\n    the shape of the gradient w.r.t. input to be specified explicitly.\\n\\n    Args:\\n        input_size : Shape of the input gradient tensor\\n        weight: weight tensor (out_channels x in_channels/groups x kH x kW)\\n        grad_output : output gradient tensor (minibatch x out_channels x oH x oW)\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n\\n    Examples::\\n\\n        >>> input = torch.randn(1, 1, 3, 3, requires_grad=True)\\n        >>> weight = torch.randn(1, 1, 1, 2, requires_grad=True)\\n        >>> output = F.conv2d(input, weight)\\n        >>> grad_output = torch.randn(output.shape)\\n        >>> grad_input = torch.autograd.grad(output, input, grad_output)\\n        >>> F.grad.conv2d_input(input.shape, weight, grad_output)\\n\\n    \\\"\\\"\\\"\\n    input = grad_output.new_empty(1).expand(input_size)\\n\\n    return torch.ops.aten.convolution_backward(\\n        grad_output,\\n        input,\\n        weight,\\n        None,\\n        _pair(stride),\\n        _pair(padding),\\n        _pair(dilation),\\n        False,\\n        [0],\\n        groups,\\n        (True, False, False),\\n    )[0]\\n\\n\\ndef conv2d_weight(\\n    input,\\n    weight_size,\\n    grad_output,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    r\\\"\\\"\\\"Compute the gradient of conv2d with respect to the weight of the convolution.\\n\\n    Args:\\n        input: input tensor of shape (minibatch x in_channels x iH x iW)\\n        weight_size : Shape of the weight gradient tensor\\n        grad_output : output gradient tensor (minibatch x out_channels x oH x oW)\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n\\n    Examples::\\n\\n        >>> input = torch.randn(1, 1, 3, 3, requires_grad=True)\\n        >>> weight = torch.randn(1, 1, 1, 2, requires_grad=True)\\n        >>> output = F.conv2d(input, weight)\\n        >>> grad_output = torch.randn(output.shape)\\n        >>> # xdoctest: +SKIP\\n        >>> grad_weight = torch.autograd.grad(output, filter, grad_output)\\n        >>> F.grad.conv2d_weight(input, weight.shape, grad_output)\\n\\n    \\\"\\\"\\\"\\n    weight = grad_output.new_empty(1).expand(weight_size)\\n\\n    return torch.ops.aten.convolution_backward(\\n        grad_output,\\n        input,\\n        weight,\\n        None,\\n        _pair(stride),\\n        _pair(padding),\\n        _pair(dilation),\\n        False,\\n        [0],\\n        groups,\\n        (False, True, False),\\n    )[1]\\n\\n\\ndef conv3d_input(\\n    input_size,\\n    weight,\\n    grad_output,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    r\\\"\\\"\\\"Compute the gradient of conv3d with respect to the input of the convolution.\\n\\n    This is same as the 3D transposed convolution operator under the hood but requires\\n    the shape of the gradient w.r.t. input to be specified explicitly.\\n\\n    Args:\\n        input_size : Shape of the input gradient tensor\\n        weight: weights tensor (out_channels x in_channels/groups x kT x kH x kW)\\n        grad_output : output gradient tensor (minibatch x out_channels x oT x oH x oW)\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n\\n    Examples::\\n\\n        >>> input = torch.randn(2, 8, 10, 10, 20, requires_grad=True)\\n        >>> weight = torch.randn(4, 8, 2, 3, 3, requires_grad=True)\\n        >>> output = F.conv3d(input, weight)\\n        >>> grad_output = torch.randn(output.shape)\\n        >>> grad_input = torch.autograd.grad(output, input, grad_output)\\n        >>> F.grad.conv3d_input(input.shape, weight, grad_output)\\n\\n    \\\"\\\"\\\"\\n    input = grad_output.new_empty(1).expand(input_size)\\n\\n    return torch.ops.aten.convolution_backward(\\n        grad_output,\\n        input,\\n        weight,\\n        None,\\n        _triple(stride),\\n        _triple(padding),\\n        _triple(dilation),\\n        False,\\n        [0],\\n        groups,\\n        (True, False, False),\\n    )[0]\\n\\n\\ndef conv3d_weight(\\n    input,\\n    weight_size,\\n    grad_output,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    r\\\"\\\"\\\"Compute the gradient of conv3d with respect to the weight of the convolution.\\n\\n    Args:\\n        input: input tensor of shape (minibatch x in_channels x iT x iH x iW)\\n        weight_size : Shape of the weight gradient tensor\\n        grad_output : output gradient tensor (minibatch x out_channels x oT x oH x oW)\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n\\n    Examples::\\n\\n        >>> input = torch.randn(2, 8, 10, 10, 20, requires_grad=True)\\n        >>> weight = torch.randn(4, 8, 2, 3, 3, requires_grad=True)\\n        >>> output = F.conv3d(input, weight)\\n        >>> grad_output = torch.randn(output.shape)\\n        >>> grad_weight = torch.autograd.grad(output, weight, grad_output)\\n        >>> F.grad.conv3d_weight(input, weight.shape, grad_output)\\n\\n    \\\"\\\"\\\"\\n    weight = grad_output.new_empty(1).expand(weight_size)\\n\\n    return torch.ops.aten.convolution_backward(\\n        grad_output,\\n        input,\\n        weight,\\n        None,\\n        _triple(stride),\\n        _triple(padding),\\n        _triple(dilation),\\n        False,\\n        [0],\\n        groups,\\n        (False, True, False),\\n    )[1]\\n\\n\\nimport warnings\\nfrom typing import Optional\\n\\n\\n# NB: Keep this file in sync with enums in aten/src/ATen/core/Reduction.h\\n\\n\\ndef get_enum(reduction: str) -> int:\\n    if reduction == \\\"none\\\":\\n        ret = 0\\n    elif reduction == \\\"mean\\\":\\n        ret = 1\\n    elif reduction == \\\"elementwise_mean\\\":\\n        warnings.warn(\\n            \\\"reduction='elementwise_mean' is deprecated. \\\"\\n            \\\"Please use reduction='mean' instead.\\\"\\n        )\\n        ret = 1\\n    elif reduction == \\\"sum\\\":\\n        ret = 2\\n    else:\\n        ret = -1  # TODO: remove once JIT exceptions support control flow\\n        raise ValueError(f\\\"{reduction} is not a valid value for reduction\\\")\\n    return ret\\n\\n\\n# In order to support previous versions, accept boolean size_average and reduce\\n# and convert them into the new constants for now\\n\\n\\n# We use these functions in torch/legacy as well, in which case we'll silence the warning\\ndef legacy_get_string(\\n    size_average: Optional[bool],\\n    reduce: Optional[bool],\\n    emit_warning: bool = True,\\n) -> str:\\n    warning = \\\"size_average and reduce args will be deprecated, please use reduction='{}' instead.\\\"\\n\\n    if size_average is None:\\n        size_average = True\\n    if reduce is None:\\n        reduce = True\\n\\n    if size_average and reduce:\\n        ret = \\\"mean\\\"\\n    elif reduce:\\n        ret = \\\"sum\\\"\\n    else:\\n        ret = \\\"none\\\"\\n    if emit_warning:\\n        warnings.warn(warning.format(ret))\\n    return ret\\n\\n\\ndef legacy_get_enum(\\n    size_average: Optional[bool],\\n    reduce: Optional[bool],\\n    emit_warning: bool = True,\\n) -> int:\\n    return get_enum(legacy_get_string(size_average, reduce, emit_warning))\\n\\n\\n# mypy: allow-untyped-defs\\n\\\"\\\"\\\"This file contains utilities for initializing neural network parameters.\\\"\\\"\\\"\\nimport math\\nimport warnings\\nfrom typing import Optional as _Optional\\n\\nimport torch\\nfrom torch import Tensor\\n\\n\\n# These no_grad_* functions are necessary as wrappers around the parts of these\\n# functions that use `with torch.no_grad()`. The JIT doesn't support context\\n# managers, so these need to be implemented as builtins. Using these wrappers\\n# lets us keep those builtins small and re-usable.\\ndef _no_grad_uniform_(tensor, a, b, generator=None):\\n    with torch.no_grad():\\n        return tensor.uniform_(a, b, generator=generator)\\n\\n\\ndef _no_grad_normal_(tensor, mean, std, generator=None):\\n    with torch.no_grad():\\n        return tensor.normal_(mean, std, generator=generator)\\n\\n\\ndef _no_grad_trunc_normal_(tensor, mean, std, a, b, generator=None):\\n    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf\\n    def norm_cdf(x):\\n        # Computes standard normal cumulative distribution function\\n        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0\\n\\n    if (mean < a - 2 * std) or (mean > b + 2 * std):\\n        warnings.warn(\\n            \\\"mean is more than 2 std from [a, b] in nn.init.trunc_normal_. \\\"\\n            \\\"The distribution of values may be incorrect.\\\",\\n            stacklevel=2,\\n        )\\n\\n    with torch.no_grad():\\n        # Values are generated by using a truncated uniform distribution and\\n        # then using the inverse CDF for the normal distribution.\\n        # Get upper and lower cdf values\\n        l = norm_cdf((a - mean) / std)\\n        u = norm_cdf((b - mean) / std)\\n\\n        # Uniformly fill tensor with values from [l, u], then translate to\\n        # [2l-1, 2u-1].\\n        tensor.uniform_(2 * l - 1, 2 * u - 1, generator=generator)\\n\\n        # Use inverse cdf transform for normal distribution to get truncated\\n        # standard normal\\n        tensor.erfinv_()\\n\\n        # Transform to proper mean, std\\n        tensor.mul_(std * math.sqrt(2.0))\\n        tensor.add_(mean)\\n\\n        # Clamp to ensure it's in the proper range\\n        tensor.clamp_(min=a, max=b)\\n        return tensor\\n\\n\\ndef _no_grad_fill_(tensor, val):\\n    with torch.no_grad():\\n        return tensor.fill_(val)\\n\\n\\ndef _no_grad_zero_(tensor):\\n    with torch.no_grad():\\n        return tensor.zero_()\\n\\n\\ndef calculate_gain(nonlinearity, param=None):\\n    r\\\"\\\"\\\"Return the recommended gain value for the given nonlinearity function.\\n\\n    The values are as follows:\\n\\n    ================= ====================================================\\n    nonlinearity      gain\\n    ================= ====================================================\\n    Linear / Identity :math:`1`\\n    Conv{1,2,3}D      :math:`1`\\n    Sigmoid           :math:`1`\\n    Tanh              :math:`\\\\frac{5}{3}`\\n    ReLU              :math:`\\\\sqrt{2}`\\n    Leaky Relu        :math:`\\\\sqrt{\\\\frac{2}{1 + \\\\text{negative\\\\_slope}^2}}`\\n    SELU              :math:`\\\\frac{3}{4}`\\n    ================= ====================================================\\n\\n    .. warning::\\n        In order to implement `Self-Normalizing Neural Networks`_ ,\\n        you should use ``nonlinearity='linear'`` instead of ``nonlinearity='selu'``.\\n        This gives the initial weights a variance of ``1 / N``,\\n        which is necessary to induce a stable fixed point in the forward pass.\\n        In contrast, the default gain for ``SELU`` sacrifices the normalization\\n        effect for more stable gradient flow in rectangular layers.\\n\\n    Args:\\n        nonlinearity: the non-linear function (`nn.functional` name)\\n        param: optional parameter for the non-linear function\\n\\n    Examples:\\n        >>> gain = nn.init.calculate_gain('leaky_relu', 0.2)  # leaky_relu with negative_slope=0.2\\n\\n    .. _Self-Normalizing Neural Networks: https://papers.nips.cc/paper/2017/hash/5d44ee6f2c3f71b73125876103c8f6c4-Abstract.html\\n    \\\"\\\"\\\"\\n    linear_fns = [\\n        \\\"linear\\\",\\n        \\\"conv1d\\\",\\n        \\\"conv2d\\\",\\n        \\\"conv3d\\\",\\n        \\\"conv_transpose1d\\\",\\n        \\\"conv_transpose2d\\\",\\n        \\\"conv_transpose3d\\\",\\n    ]\\n    if nonlinearity in linear_fns or nonlinearity == \\\"sigmoid\\\":\\n        return 1\\n    elif nonlinearity == \\\"tanh\\\":\\n        return 5.0 / 3\\n    elif nonlinearity == \\\"relu\\\":\\n        return math.sqrt(2.0)\\n    elif nonlinearity == \\\"leaky_relu\\\":\\n        if param is None:\\n            negative_slope = 0.01\\n        elif (\\n            not isinstance(param, bool)\\n            and isinstance(param, int)\\n            or isinstance(param, float)\\n        ):\\n            # True/False are instances of int, hence check above\\n            negative_slope = param\\n        else:\\n            raise ValueError(f\\\"negative_slope {param} not a valid number\\\")\\n        return math.sqrt(2.0 / (1 + negative_slope**2))\\n    elif nonlinearity == \\\"selu\\\":\\n        return (\\n            3.0 / 4\\n        )  # Value found empirically (https://github.com/pytorch/pytorch/pull/50664)\\n    else:\\n        raise ValueError(f\\\"Unsupported nonlinearity {nonlinearity}\\\")\\n\\n\\ndef uniform_(\\n    tensor: Tensor,\\n    a: float = 0.0,\\n    b: float = 1.0,\\n    generator: _Optional[torch.Generator] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input Tensor with values drawn from the uniform distribution.\\n\\n    :math:`\\\\mathcal{U}(a, b)`.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        a: the lower bound of the uniform distribution\\n        b: the upper bound of the uniform distribution\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.uniform_(w)\\n    \\\"\\\"\\\"\\n    if torch.overrides.has_torch_function_variadic(tensor):\\n        return torch.overrides.handle_torch_function(\\n            uniform_, (tensor,), tensor=tensor, a=a, b=b, generator=generator\\n        )\\n    return _no_grad_uniform_(tensor, a, b, generator)\\n\\n\\ndef normal_(\\n    tensor: Tensor,\\n    mean: float = 0.0,\\n    std: float = 1.0,\\n    generator: _Optional[torch.Generator] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input Tensor with values drawn from the normal distribution.\\n\\n    :math:`\\\\mathcal{N}(\\\\text{mean}, \\\\text{std}^2)`.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        mean: the mean of the normal distribution\\n        std: the standard deviation of the normal distribution\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.normal_(w)\\n    \\\"\\\"\\\"\\n    if torch.overrides.has_torch_function_variadic(tensor):\\n        return torch.overrides.handle_torch_function(\\n            normal_, (tensor,), tensor=tensor, mean=mean, std=std, generator=generator\\n        )\\n    return _no_grad_normal_(tensor, mean, std, generator)\\n\\n\\ndef trunc_normal_(\\n    tensor: Tensor,\\n    mean: float = 0.0,\\n    std: float = 1.0,\\n    a: float = -2.0,\\n    b: float = 2.0,\\n    generator: _Optional[torch.Generator] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input Tensor with values drawn from a truncated normal distribution.\\n\\n    The values are effectively drawn from the\\n    normal distribution :math:`\\\\mathcal{N}(\\\\text{mean}, \\\\text{std}^2)`\\n    with values outside :math:`[a, b]` redrawn until they are within\\n    the bounds. The method used for generating the random values works\\n    best when :math:`a \\\\leq \\\\text{mean} \\\\leq b`.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        mean: the mean of the normal distribution\\n        std: the standard deviation of the normal distribution\\n        a: the minimum cutoff value\\n        b: the maximum cutoff value\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.trunc_normal_(w)\\n    \\\"\\\"\\\"\\n    return _no_grad_trunc_normal_(tensor, mean, std, a, b, generator=generator)\\n\\n\\ndef constant_(tensor: Tensor, val: float) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input Tensor with the value :math:`\\\\text{val}`.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        val: the value to fill the tensor with\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.constant_(w, 0.3)\\n    \\\"\\\"\\\"\\n    if torch.overrides.has_torch_function_variadic(tensor):\\n        return torch.overrides.handle_torch_function(\\n            constant_, (tensor,), tensor=tensor, val=val\\n        )\\n    return _no_grad_fill_(tensor, val)\\n\\n\\ndef ones_(tensor: Tensor) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input Tensor with the scalar value `1`.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.ones_(w)\\n    \\\"\\\"\\\"\\n    return _no_grad_fill_(tensor, 1.0)\\n\\n\\ndef zeros_(tensor: Tensor) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input Tensor with the scalar value `0`.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.zeros_(w)\\n    \\\"\\\"\\\"\\n    return _no_grad_zero_(tensor)\\n\\n\\ndef eye_(tensor):\\n    r\\\"\\\"\\\"Fill the 2-dimensional input `Tensor` with the identity matrix.\\n\\n    Preserves the identity of the inputs in `Linear` layers, where as\\n    many inputs are preserved as possible.\\n\\n    Args:\\n        tensor: a 2-dimensional `torch.Tensor`\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.eye_(w)\\n    \\\"\\\"\\\"\\n    if tensor.ndimension() != 2:\\n        raise ValueError(\\\"Only tensors with 2 dimensions are supported\\\")\\n\\n    with torch.no_grad():\\n        torch.eye(*tensor.shape, out=tensor, requires_grad=tensor.requires_grad)\\n    return tensor\\n\\n\\ndef dirac_(tensor, groups=1):\\n    r\\\"\\\"\\\"Fill the {3, 4, 5}-dimensional input `Tensor` with the Dirac delta function.\\n\\n    Preserves the identity of the inputs in `Convolutional`\\n    layers, where as many input channels are preserved as possible. In case\\n    of groups>1, each group of channels preserves identity\\n\\n    Args:\\n        tensor: a {3, 4, 5}-dimensional `torch.Tensor`\\n        groups (int, optional): number of groups in the conv layer (default: 1)\\n    Examples:\\n        >>> w = torch.empty(3, 16, 5, 5)\\n        >>> nn.init.dirac_(w)\\n        >>> w = torch.empty(3, 24, 5, 5)\\n        >>> nn.init.dirac_(w, 3)\\n    \\\"\\\"\\\"\\n    dimensions = tensor.ndimension()\\n    if dimensions not in [3, 4, 5]:\\n        raise ValueError(\\\"Only tensors with 3, 4, or 5 dimensions are supported\\\")\\n\\n    sizes = tensor.size()\\n\\n    if sizes[0] % groups != 0:\\n        raise ValueError(\\\"dim 0 must be divisible by groups\\\")\\n\\n    out_chans_per_grp = sizes[0] // groups\\n    min_dim = min(out_chans_per_grp, sizes[1])\\n\\n    with torch.no_grad():\\n        tensor.zero_()\\n\\n        for g in range(groups):\\n            for d in range(min_dim):\\n                if dimensions == 3:  # Temporal convolution\\n                    tensor[g * out_chans_per_grp + d, d, tensor.size(2) // 2] = 1\\n                elif dimensions == 4:  # Spatial convolution\\n                    tensor[\\n                        g * out_chans_per_grp + d,\\n                        d,\\n                        tensor.size(2) // 2,\\n                        tensor.size(3) // 2,\\n                    ] = 1\\n                else:  # Volumetric convolution\\n                    tensor[\\n                        g * out_chans_per_grp + d,\\n                        d,\\n                        tensor.size(2) // 2,\\n                        tensor.size(3) // 2,\\n                        tensor.size(4) // 2,\\n                    ] = 1\\n    return tensor\\n\\n\\ndef _calculate_fan_in_and_fan_out(tensor):\\n    dimensions = tensor.dim()\\n    if dimensions < 2:\\n        raise ValueError(\\n            \\\"Fan in and fan out can not be computed for tensor with fewer than 2 dimensions\\\"\\n        )\\n\\n    num_input_fmaps = tensor.size(1)\\n    num_output_fmaps = tensor.size(0)\\n    receptive_field_size = 1\\n    if tensor.dim() > 2:\\n        # math.prod is not always available, accumulate the product manually\\n        # we could use functools.reduce but that is not supported by TorchScript\\n        for s in tensor.shape[2:]:\\n            receptive_field_size *= s\\n    fan_in = num_input_fmaps * receptive_field_size\\n    fan_out = num_output_fmaps * receptive_field_size\\n\\n    return fan_in, fan_out\\n\\n\\ndef xavier_uniform_(\\n    tensor: Tensor,\\n    gain: float = 1.0,\\n    generator: _Optional[torch.Generator] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input `Tensor` with values using a Xavier uniform distribution.\\n\\n    The method is described in `Understanding the difficulty of training\\n    deep feedforward neural networks` - Glorot, X. & Bengio, Y. (2010).\\n    The resulting tensor will have values sampled from\\n    :math:`\\\\mathcal{U}(-a, a)` where\\n\\n    .. math::\\n        a = \\\\text{gain} \\\\times \\\\sqrt{\\\\frac{6}{\\\\text{fan\\\\_in} + \\\\text{fan\\\\_out}}}\\n\\n    Also known as Glorot initialization.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        gain: an optional scaling factor\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain('relu'))\\n\\n    Note:\\n        Be aware that ``fan_in`` and ``fan_out`` are calculated assuming\\n        that the weight matrix is used in a transposed manner,\\n        (i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).\\n        This is important for correct initialization.\\n        If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,\\n        pass in a transposed weight matrix, i.e. ``nn.init.xavier_uniform_(w.T, ...)``.\\n    \\\"\\\"\\\"\\n    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)\\n    std = gain * math.sqrt(2.0 / float(fan_in + fan_out))\\n    a = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation\\n\\n    return _no_grad_uniform_(tensor, -a, a, generator)\\n\\n\\ndef xavier_normal_(\\n    tensor: Tensor,\\n    gain: float = 1.0,\\n    generator: _Optional[torch.Generator] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Fill the input `Tensor` with values using a Xavier normal distribution.\\n\\n    The method is described in `Understanding the difficulty of training deep feedforward\\n    neural networks` - Glorot, X. & Bengio, Y. (2010). The resulting tensor\\n    will have values sampled from :math:`\\\\mathcal{N}(0, \\\\text{std}^2)` where\\n\\n    .. math::\\n        \\\\text{std} = \\\\text{gain} \\\\times \\\\sqrt{\\\\frac{2}{\\\\text{fan\\\\_in} + \\\\text{fan\\\\_out}}}\\n\\n    Also known as Glorot initialization.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        gain: an optional scaling factor\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.xavier_normal_(w)\\n\\n    Note:\\n        Be aware that ``fan_in`` and ``fan_out`` are calculated assuming\\n        that the weight matrix is used in a transposed manner,\\n        (i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).\\n        This is important for correct initialization.\\n        If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,\\n        pass in a transposed weight matrix, i.e. ``nn.init.xavier_normal_(w.T, ...)``.\\n    \\\"\\\"\\\"\\n    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)\\n    std = gain * math.sqrt(2.0 / float(fan_in + fan_out))\\n\\n    return _no_grad_normal_(tensor, 0.0, std, generator)\\n\\n\\ndef _calculate_correct_fan(tensor, mode):\\n    mode = mode.lower()\\n    valid_modes = [\\\"fan_in\\\", \\\"fan_out\\\"]\\n    if mode not in valid_modes:\\n        raise ValueError(f\\\"Mode {mode} not supported, please use one of {valid_modes}\\\")\\n\\n    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)\\n    return fan_in if mode == \\\"fan_in\\\" else fan_out\\n\\n\\ndef kaiming_uniform_(\\n    tensor: Tensor,\\n    a: float = 0,\\n    mode: str = \\\"fan_in\\\",\\n    nonlinearity: str = \\\"leaky_relu\\\",\\n    generator: _Optional[torch.Generator] = None,\\n):\\n    r\\\"\\\"\\\"Fill the input `Tensor` with values using a Kaiming uniform distribution.\\n\\n    The method is described in `Delving deep into rectifiers: Surpassing\\n    human-level performance on ImageNet classification` - He, K. et al. (2015).\\n    The resulting tensor will have values sampled from\\n    :math:`\\\\mathcal{U}(-\\\\text{bound}, \\\\text{bound})` where\\n\\n    .. math::\\n        \\\\text{bound} = \\\\text{gain} \\\\times \\\\sqrt{\\\\frac{3}{\\\\text{fan\\\\_mode}}}\\n\\n    Also known as He initialization.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        a: the negative slope of the rectifier used after this layer (only\\n            used with ``'leaky_relu'``)\\n        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``\\n            preserves the magnitude of the variance of the weights in the\\n            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the\\n            backwards pass.\\n        nonlinearity: the non-linear function (`nn.functional` name),\\n            recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default).\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.kaiming_uniform_(w, mode='fan_in', nonlinearity='relu')\\n\\n    Note:\\n        Be aware that ``fan_in`` and ``fan_out`` are calculated assuming\\n        that the weight matrix is used in a transposed manner,\\n        (i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).\\n        This is important for correct initialization.\\n        If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,\\n        pass in a transposed weight matrix, i.e. ``nn.init.kaiming_uniform_(w.T, ...)``.\\n    \\\"\\\"\\\"\\n    if torch.overrides.has_torch_function_variadic(tensor):\\n        return torch.overrides.handle_torch_function(\\n            kaiming_uniform_,\\n            (tensor,),\\n            tensor=tensor,\\n            a=a,\\n            mode=mode,\\n            nonlinearity=nonlinearity,\\n            generator=generator,\\n        )\\n\\n    if 0 in tensor.shape:\\n        warnings.warn(\\\"Initializing zero-element tensors is a no-op\\\")\\n        return tensor\\n    fan = _calculate_correct_fan(tensor, mode)\\n    gain = calculate_gain(nonlinearity, a)\\n    std = gain / math.sqrt(fan)\\n    bound = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation\\n    with torch.no_grad():\\n        return tensor.uniform_(-bound, bound, generator=generator)\\n\\n\\ndef kaiming_normal_(\\n    tensor: Tensor,\\n    a: float = 0,\\n    mode: str = \\\"fan_in\\\",\\n    nonlinearity: str = \\\"leaky_relu\\\",\\n    generator: _Optional[torch.Generator] = None,\\n):\\n    r\\\"\\\"\\\"Fill the input `Tensor` with values using a Kaiming normal distribution.\\n\\n    The method is described in `Delving deep into rectifiers: Surpassing\\n    human-level performance on ImageNet classification` - He, K. et al. (2015).\\n    The resulting tensor will have values sampled from\\n    :math:`\\\\mathcal{N}(0, \\\\text{std}^2)` where\\n\\n    .. math::\\n        \\\\text{std} = \\\\frac{\\\\text{gain}}{\\\\sqrt{\\\\text{fan\\\\_mode}}}\\n\\n    Also known as He initialization.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        a: the negative slope of the rectifier used after this layer (only\\n            used with ``'leaky_relu'``)\\n        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``\\n            preserves the magnitude of the variance of the weights in the\\n            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the\\n            backwards pass.\\n        nonlinearity: the non-linear function (`nn.functional` name),\\n            recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default).\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.kaiming_normal_(w, mode='fan_out', nonlinearity='relu')\\n\\n    Note:\\n        Be aware that ``fan_in`` and ``fan_out`` are calculated assuming\\n        that the weight matrix is used in a transposed manner,\\n        (i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).\\n        This is important for correct initialization.\\n        If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,\\n        pass in a transposed weight matrix, i.e. ``nn.init.kaiming_normal_(w.T, ...)``.\\n    \\\"\\\"\\\"\\n    if 0 in tensor.shape:\\n        warnings.warn(\\\"Initializing zero-element tensors is a no-op\\\")\\n        return tensor\\n    fan = _calculate_correct_fan(tensor, mode)\\n    gain = calculate_gain(nonlinearity, a)\\n    std = gain / math.sqrt(fan)\\n    with torch.no_grad():\\n        return tensor.normal_(0, std, generator=generator)\\n\\n\\ndef orthogonal_(\\n    tensor,\\n    gain=1,\\n    generator: _Optional[torch.Generator] = None,\\n):\\n    r\\\"\\\"\\\"Fill the input `Tensor` with a (semi) orthogonal matrix.\\n\\n    Described in `Exact solutions to the nonlinear dynamics of learning in deep\\n    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have\\n    at least 2 dimensions, and for tensors with more than 2 dimensions the\\n    trailing dimensions are flattened.\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`, where :math:`n \\\\geq 2`\\n        gain: optional scaling factor\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.orthogonal_(w)\\n    \\\"\\\"\\\"\\n    if tensor.ndimension() < 2:\\n        raise ValueError(\\\"Only tensors with 2 or more dimensions are supported\\\")\\n\\n    if tensor.numel() == 0:\\n        # no-op\\n        return tensor\\n    rows = tensor.size(0)\\n    cols = tensor.numel() // rows\\n    flattened = tensor.new_empty((rows, cols)).normal_(0, 1, generator=generator)\\n\\n    if rows < cols:\\n        flattened.t_()\\n\\n    # Compute the qr factorization\\n    q, r = torch.linalg.qr(flattened)\\n    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf\\n    d = torch.diag(r, 0)\\n    ph = d.sign()\\n    q *= ph\\n\\n    if rows < cols:\\n        q.t_()\\n\\n    with torch.no_grad():\\n        tensor.view_as(q).copy_(q)\\n        tensor.mul_(gain)\\n    return tensor\\n\\n\\ndef sparse_(\\n    tensor,\\n    sparsity,\\n    std=0.01,\\n    generator: _Optional[torch.Generator] = None,\\n):\\n    r\\\"\\\"\\\"Fill the 2D input `Tensor` as a sparse matrix.\\n\\n    The non-zero elements will be drawn from the normal distribution\\n    :math:`\\\\mathcal{N}(0, 0.01)`, as described in `Deep learning via\\n    Hessian-free optimization` - Martens, J. (2010).\\n\\n    Args:\\n        tensor: an n-dimensional `torch.Tensor`\\n        sparsity: The fraction of elements in each column to be set to zero\\n        std: the standard deviation of the normal distribution used to generate\\n            the non-zero values\\n        generator: the torch Generator to sample from (default: None)\\n\\n    Examples:\\n        >>> w = torch.empty(3, 5)\\n        >>> nn.init.sparse_(w, sparsity=0.1)\\n    \\\"\\\"\\\"\\n    if tensor.ndimension() != 2:\\n        raise ValueError(\\\"Only tensors with 2 dimensions are supported\\\")\\n\\n    rows, cols = tensor.shape\\n    num_zeros = int(math.ceil(sparsity * rows))\\n\\n    with torch.no_grad():\\n        tensor.normal_(0, std, generator=generator)\\n        for col_idx in range(cols):\\n            row_indices = torch.randperm(rows)\\n            zero_indices = row_indices[:num_zeros]\\n            tensor[zero_indices, col_idx] = 0\\n    return tensor\\n\\n\\n# for backward compatibility\\ndef _make_deprecate(meth):\\n    new_name = meth.__name__\\n    old_name = new_name[:-1]\\n\\n    def deprecated_init(*args, **kwargs):\\n        warnings.warn(\\n            f\\\"`nn.init.{old_name}` is now deprecated in favor of `nn.init.{new_name}`.\\\",\\n            FutureWarning,\\n            stacklevel=2,\\n        )\\n        return meth(*args, **kwargs)\\n\\n    deprecated_init.__doc__ = rf\\\"\\\"\\\"\\n    {old_name}(...)\\n\\n    .. warning::\\n        This method is now deprecated in favor of :func:`torch.nn.init.{new_name}`.\\n\\n    See :func:`~torch.nn.init.{new_name}` for details.\\\"\\\"\\\"\\n    deprecated_init.__name__ = old_name\\n    return deprecated_init\\n\\n\\nuniform = _make_deprecate(uniform_)\\nnormal = _make_deprecate(normal_)\\nconstant = _make_deprecate(constant_)\\neye = _make_deprecate(eye_)\\ndirac = _make_deprecate(dirac_)\\nxavier_uniform = _make_deprecate(xavier_uniform_)\\nxavier_normal = _make_deprecate(xavier_normal_)\\nkaiming_uniform = _make_deprecate(kaiming_uniform_)\\nkaiming_normal = _make_deprecate(kaiming_normal_)\\northogonal = _make_deprecate(orthogonal_)\\nsparse = _make_deprecate(sparse_)\\n\\n\\nr\\\"\\\"\\\"nn.quantized.functional.\\n\\nQuantized equivalents of the `nn.functional`.\\n\\nNote::\\n    This location is in the process of being deprecated.\\n    Please, use the `torch.ao.nn.quantized.functional` instead.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.functional import *  # noqa: F401,F403\\n\\n\\nfrom torch.nn.quantized import dynamic, functional, modules  # noqa: F403\\nfrom torch.nn.quantized.modules import *  # noqa: F403\\nfrom torch.nn.quantized.modules import MaxPool2d\\n\\n\\n__all__ = [\\n    \\\"BatchNorm2d\\\",\\n    \\\"BatchNorm3d\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n    \\\"DeQuantize\\\",\\n    \\\"Dropout\\\",\\n    \\\"ELU\\\",\\n    \\\"Embedding\\\",\\n    \\\"EmbeddingBag\\\",\\n    \\\"GroupNorm\\\",\\n    \\\"Hardswish\\\",\\n    \\\"InstanceNorm1d\\\",\\n    \\\"InstanceNorm2d\\\",\\n    \\\"InstanceNorm3d\\\",\\n    \\\"LayerNorm\\\",\\n    \\\"LeakyReLU\\\",\\n    \\\"Linear\\\",\\n    \\\"LSTM\\\",\\n    \\\"MultiheadAttention\\\",\\n    \\\"PReLU\\\",\\n    \\\"Quantize\\\",\\n    \\\"ReLU6\\\",\\n    \\\"Sigmoid\\\",\\n    \\\"Softmax\\\",\\n    # Wrapper modules\\n    \\\"FloatFunctional\\\",\\n    \\\"FXFloatFunctional\\\",\\n    \\\"QFunctional\\\",\\n]\\n\\n\\nfrom torch.nn.quantized._reference.modules import *  # noqa: F403\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Reference Modules.\\n\\nThis module is in the process of migration to\\n`torch/ao/nn/quantized/reference`, and is kept here for\\ncompatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/reference`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.reference.modules.linear import Linear\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Reference Modules.\\n\\nThis module is in the process of migration to\\n`torch/ao/nn/quantized/reference`, and is kept here for\\ncompatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/reference`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.reference.modules.rnn import (\\n    GRUCell,\\n    LSTM,\\n    LSTMCell,\\n    RNNBase,\\n    RNNCell,\\n    RNNCellBase,\\n)\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Reference Modules.\\n\\nThis module is in the process of migration to\\n`torch/ao/nn/quantized/reference`, and is kept here for\\ncompatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/reference`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.reference.modules.utils import (\\n    _get_weight_qparam_keys,\\n    _quantize_and_dequantize_weight,\\n    _quantize_weight,\\n    _save_weight_qparams,\\n    ReferenceQuantizedModule,\\n)\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Reference Modules.\\n\\nThis module is in the process of migration to\\n`torch/ao/nn/quantized/reference`, and is kept here for\\ncompatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/reference`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.reference.modules.sparse import Embedding, EmbeddingBag\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Reference Modules.\\n\\nThis module is in the process of migration to\\n`torch/ao/nn/quantized/reference`, and is kept here for\\ncompatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/reference`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.reference.modules.conv import (\\n    _ConvNd,\\n    _ConvTransposeNd,\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n)\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Reference Modules.\\n\\nThis module is in the process of migration to\\n`torch/ao/nn/quantized/reference`, and is kept here for\\ncompatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/reference`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.reference.modules.conv import (\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n)\\nfrom torch.ao.nn.quantized.reference.modules.linear import Linear\\nfrom torch.ao.nn.quantized.reference.modules.rnn import GRUCell, LSTM, LSTMCell, RNNCell\\nfrom torch.ao.nn.quantized.reference.modules.sparse import Embedding, EmbeddingBag\\n\\n\\n__all__ = [\\n    \\\"Linear\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n    \\\"RNNCell\\\",\\n    \\\"LSTMCell\\\",\\n    \\\"GRUCell\\\",\\n    \\\"LSTM\\\",\\n    \\\"Embedding\\\",\\n    \\\"EmbeddingBag\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.linear import Linear, LinearPackedParams\\n\\n\\n__all__ = [\\\"LinearPackedParams\\\", \\\"Linear\\\"]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.rnn import LSTM\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.normalization import (\\n    GroupNorm,\\n    InstanceNorm1d,\\n    InstanceNorm2d,\\n    InstanceNorm3d,\\n    LayerNorm,\\n)\\n\\n\\n__all__ = [\\n    \\\"LayerNorm\\\",\\n    \\\"GroupNorm\\\",\\n    \\\"InstanceNorm1d\\\",\\n    \\\"InstanceNorm2d\\\",\\n    \\\"InstanceNorm3d\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.utils import (\\n    _hide_packed_params_repr,\\n    _ntuple_from_first,\\n    _pair_from_first,\\n    _quantize_weight,\\n    WeightedQuantizedModule,\\n)\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.batchnorm import BatchNorm2d, BatchNorm3d\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.activation import (\\n    ELU,\\n    Hardswish,\\n    LeakyReLU,\\n    MultiheadAttention,\\n    PReLU,\\n    ReLU6,\\n    Sigmoid,\\n    Softmax,\\n)\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.functional_modules import (\\n    FloatFunctional,\\n    FXFloatFunctional,\\n    QFunctional,\\n)\\n\\n\\n__all__ = [\\\"FloatFunctional\\\", \\\"FXFloatFunctional\\\", \\\"QFunctional\\\"]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.dropout import Dropout\\n\\n\\n__all__ = [\\\"Dropout\\\"]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.conv import (\\n    _reverse_repeat_padding,\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n)\\n\\n\\n__all__ = [\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n]\\n\\n\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nNote::\\n    The `torch.nn.quantized` namespace is in the process of being deprecated.\\n    Please, use `torch.ao.nn.quantized` instead.\\n\\\"\\\"\\\"\\n\\n# The following imports are needed in case the user decides\\n# to import the files directly,\\n# s.a. `from torch.nn.quantized.modules.conv import ...`.\\n# No need to add them to the `__all__`.\\nfrom torch.ao.nn.quantized.modules import (\\n    activation,\\n    batchnorm,\\n    conv,\\n    DeQuantize,\\n    dropout,\\n    embedding_ops,\\n    functional_modules,\\n    linear,\\n    MaxPool2d,\\n    normalization,\\n    Quantize,\\n    rnn,\\n    utils,\\n)\\nfrom torch.ao.nn.quantized.modules.activation import (\\n    ELU,\\n    Hardswish,\\n    LeakyReLU,\\n    MultiheadAttention,\\n    PReLU,\\n    ReLU6,\\n    Sigmoid,\\n    Softmax,\\n)\\nfrom torch.ao.nn.quantized.modules.batchnorm import BatchNorm2d, BatchNorm3d\\nfrom torch.ao.nn.quantized.modules.conv import (\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n)\\nfrom torch.ao.nn.quantized.modules.dropout import Dropout\\nfrom torch.ao.nn.quantized.modules.embedding_ops import Embedding, EmbeddingBag\\nfrom torch.ao.nn.quantized.modules.functional_modules import (\\n    FloatFunctional,\\n    FXFloatFunctional,\\n    QFunctional,\\n)\\nfrom torch.ao.nn.quantized.modules.linear import Linear\\nfrom torch.ao.nn.quantized.modules.normalization import (\\n    GroupNorm,\\n    InstanceNorm1d,\\n    InstanceNorm2d,\\n    InstanceNorm3d,\\n    LayerNorm,\\n)\\nfrom torch.ao.nn.quantized.modules.rnn import LSTM\\n\\n\\n__all__ = [\\n    \\\"BatchNorm2d\\\",\\n    \\\"BatchNorm3d\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n    \\\"DeQuantize\\\",\\n    \\\"ELU\\\",\\n    \\\"Embedding\\\",\\n    \\\"EmbeddingBag\\\",\\n    \\\"GroupNorm\\\",\\n    \\\"Hardswish\\\",\\n    \\\"InstanceNorm1d\\\",\\n    \\\"InstanceNorm2d\\\",\\n    \\\"InstanceNorm3d\\\",\\n    \\\"LayerNorm\\\",\\n    \\\"LeakyReLU\\\",\\n    \\\"Linear\\\",\\n    \\\"LSTM\\\",\\n    \\\"MultiheadAttention\\\",\\n    \\\"Quantize\\\",\\n    \\\"ReLU6\\\",\\n    \\\"Sigmoid\\\",\\n    \\\"Softmax\\\",\\n    \\\"Dropout\\\",\\n    \\\"PReLU\\\",\\n    # Wrapper modules\\n    \\\"FloatFunctional\\\",\\n    \\\"FXFloatFunctional\\\",\\n    \\\"QFunctional\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.modules.embedding_ops import (\\n    Embedding,\\n    EmbeddingBag,\\n    EmbeddingPackedParams,\\n)\\n\\n\\n__all__ = [\\\"EmbeddingPackedParams\\\", \\\"Embedding\\\", \\\"EmbeddingBag\\\"]\\n\\n\\nfrom torch.ao.nn.quantized.dynamic import *  # noqa: F403\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Dynamic Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized/dynamic`,\\nand is kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/dynamic/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\nfrom torch.ao.nn.quantized.dynamic.modules.linear import Linear\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Dynamic Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized/dynamic`,\\nand is kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/dynamic/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.dynamic.modules.rnn import (\\n    GRU,\\n    GRUCell,\\n    LSTM,\\n    LSTMCell,\\n    pack_weight_bias,\\n    PackedParameter,\\n    RNNBase,\\n    RNNCell,\\n    RNNCellBase,\\n)\\n\\n\\n__all__ = [\\n    \\\"pack_weight_bias\\\",\\n    \\\"PackedParameter\\\",\\n    \\\"RNNBase\\\",\\n    \\\"LSTM\\\",\\n    \\\"GRU\\\",\\n    \\\"RNNCellBase\\\",\\n    \\\"RNNCell\\\",\\n    \\\"LSTMCell\\\",\\n    \\\"GRUCell\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Dynamic Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized/dynamic`,\\nand is kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/dynamic/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.dynamic.modules.conv import (\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n)\\n\\n\\n__all__ = [\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantized Dynamic Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantized/dynamic`,\\nand is kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantized/dynamic`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantized.dynamic.modules import conv, linear, rnn\\nfrom torch.ao.nn.quantized.dynamic.modules.conv import (\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n)\\nfrom torch.ao.nn.quantized.dynamic.modules.linear import Linear\\nfrom torch.ao.nn.quantized.dynamic.modules.rnn import (\\n    GRU,\\n    GRUCell,\\n    LSTM,\\n    LSTMCell,\\n    RNNCell,\\n)\\n\\n\\n__all__ = [\\n    \\\"Linear\\\",\\n    \\\"LSTM\\\",\\n    \\\"GRU\\\",\\n    \\\"LSTMCell\\\",\\n    \\\"RNNCell\\\",\\n    \\\"GRUCell\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n]\\n\\n\\n# mypy: allow-untyped-defs\\nimport math\\nfrom typing import Any\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch.nn import functional as F, init\\nfrom torch.nn.parameter import Parameter, UninitializedParameter\\n\\nfrom .lazy import LazyModuleMixin\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"Bilinear\\\",\\n    \\\"Identity\\\",\\n    \\\"LazyLinear\\\",\\n    \\\"Linear\\\",\\n]\\n\\n\\nclass Identity(Module):\\n    r\\\"\\\"\\\"A placeholder identity operator that is argument-insensitive.\\n\\n    Args:\\n        args: any argument (unused)\\n        kwargs: any keyword argument (unused)\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    Examples::\\n\\n        >>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False)\\n        >>> input = torch.randn(128, 20)\\n        >>> output = m(input)\\n        >>> print(output.size())\\n        torch.Size([128, 20])\\n\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, *args: Any, **kwargs: Any) -> None:\\n        super().__init__()\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return input\\n\\n\\nclass Linear(Module):\\n    r\\\"\\\"\\\"Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`.\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    Args:\\n        in_features: size of each input sample\\n        out_features: size of each output sample\\n        bias: If set to ``False``, the layer will not learn an additive bias.\\n            Default: ``True``\\n\\n    Shape:\\n        - Input: :math:`(*, H_{in})` where :math:`*` means any number of\\n          dimensions including none and :math:`H_{in} = \\\\text{in\\\\_features}`.\\n        - Output: :math:`(*, H_{out})` where all but the last dimension\\n          are the same shape as the input and :math:`H_{out} = \\\\text{out\\\\_features}`.\\n\\n    Attributes:\\n        weight: the learnable weights of the module of shape\\n            :math:`(\\\\text{out\\\\_features}, \\\\text{in\\\\_features})`. The values are\\n            initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`, where\\n            :math:`k = \\\\frac{1}{\\\\text{in\\\\_features}}`\\n        bias:   the learnable bias of the module of shape :math:`(\\\\text{out\\\\_features})`.\\n                If :attr:`bias` is ``True``, the values are initialized from\\n                :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                :math:`k = \\\\frac{1}{\\\\text{in\\\\_features}}`\\n\\n    Examples::\\n\\n        >>> m = nn.Linear(20, 30)\\n        >>> input = torch.randn(128, 20)\\n        >>> output = m(input)\\n        >>> print(output.size())\\n        torch.Size([128, 30])\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"in_features\\\", \\\"out_features\\\"]\\n    in_features: int\\n    out_features: int\\n    weight: Tensor\\n\\n    def __init__(\\n        self,\\n        in_features: int,\\n        out_features: int,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.in_features = in_features\\n        self.out_features = out_features\\n        self.weight = Parameter(\\n            torch.empty((out_features, in_features), **factory_kwargs)\\n        )\\n        if bias:\\n            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))\\n        else:\\n            self.register_parameter(\\\"bias\\\", None)\\n        self.reset_parameters()\\n\\n    def reset_parameters(self) -> None:\\n        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with\\n        # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see\\n        # https://github.com/pytorch/pytorch/issues/57109\\n        init.kaiming_uniform_(self.weight, a=math.sqrt(5))\\n        if self.bias is not None:\\n            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)\\n            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\\n            init.uniform_(self.bias, -bound, bound)\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.linear(input, self.weight, self.bias)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}\\\"\\n\\n\\n# This class exists solely to avoid triggering an obscure error when scripting\\n# an improperly quantized attention layer. See this issue for details:\\n# https://github.com/pytorch/pytorch/issues/58969\\n# TODO: fail fast on quantization API usage error, then remove this class\\n# and replace uses of it with plain Linear\\nclass NonDynamicallyQuantizableLinear(Linear):\\n    def __init__(\\n        self,\\n        in_features: int,\\n        out_features: int,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        super().__init__(\\n            in_features, out_features, bias=bias, device=device, dtype=dtype\\n        )\\n\\n\\nclass Bilinear(Module):\\n    r\\\"\\\"\\\"Applies a bilinear transformation to the incoming data: :math:`y = x_1^T A x_2 + b`.\\n\\n    Args:\\n        in1_features: size of each first input sample\\n        in2_features: size of each second input sample\\n        out_features: size of each output sample\\n        bias: If set to False, the layer will not learn an additive bias.\\n            Default: ``True``\\n\\n    Shape:\\n        - Input1: :math:`(*, H_{in1})` where :math:`H_{in1}=\\\\text{in1\\\\_features}` and\\n          :math:`*` means any number of additional dimensions including none. All but the last dimension\\n          of the inputs should be the same.\\n        - Input2: :math:`(*, H_{in2})` where :math:`H_{in2}=\\\\text{in2\\\\_features}`.\\n        - Output: :math:`(*, H_{out})` where :math:`H_{out}=\\\\text{out\\\\_features}`\\n          and all but the last dimension are the same shape as the input.\\n\\n    Attributes:\\n        weight: the learnable weights of the module of shape\\n            :math:`(\\\\text{out\\\\_features}, \\\\text{in1\\\\_features}, \\\\text{in2\\\\_features})`.\\n            The values are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`, where\\n            :math:`k = \\\\frac{1}{\\\\text{in1\\\\_features}}`\\n        bias:   the learnable bias of the module of shape :math:`(\\\\text{out\\\\_features})`.\\n                If :attr:`bias` is ``True``, the values are initialized from\\n                :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`, where\\n                :math:`k = \\\\frac{1}{\\\\text{in1\\\\_features}}`\\n\\n    Examples::\\n\\n        >>> m = nn.Bilinear(20, 30, 40)\\n        >>> input1 = torch.randn(128, 20)\\n        >>> input2 = torch.randn(128, 30)\\n        >>> output = m(input1, input2)\\n        >>> print(output.size())\\n        torch.Size([128, 40])\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"in1_features\\\", \\\"in2_features\\\", \\\"out_features\\\"]\\n    in1_features: int\\n    in2_features: int\\n    out_features: int\\n    weight: Tensor\\n\\n    def __init__(\\n        self,\\n        in1_features: int,\\n        in2_features: int,\\n        out_features: int,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.in1_features = in1_features\\n        self.in2_features = in2_features\\n        self.out_features = out_features\\n        self.weight = Parameter(\\n            torch.empty((out_features, in1_features, in2_features), **factory_kwargs)\\n        )\\n\\n        if bias:\\n            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))\\n        else:\\n            self.register_parameter(\\\"bias\\\", None)\\n        self.reset_parameters()\\n\\n    def reset_parameters(self) -> None:\\n        bound = 1 / math.sqrt(self.weight.size(1))\\n        init.uniform_(self.weight, -bound, bound)\\n        if self.bias is not None:\\n            init.uniform_(self.bias, -bound, bound)\\n\\n    def forward(self, input1: Tensor, input2: Tensor) -> Tensor:\\n        return F.bilinear(input1, input2, self.weight, self.bias)\\n\\n    def extra_repr(self) -> str:\\n        return (\\n            f\\\"in1_features={self.in1_features}, in2_features={self.in2_features}, \\\"\\n            f\\\"out_features={self.out_features}, bias={self.bias is not None}\\\"\\n        )\\n\\n\\nclass LazyLinear(LazyModuleMixin, Linear):\\n    r\\\"\\\"\\\"A :class:`torch.nn.Linear` module where `in_features` is inferred.\\n\\n    In this module, the `weight` and `bias` are of :class:`torch.nn.UninitializedParameter`\\n    class. They will be initialized after the first call to ``forward`` is done and the\\n    module will become a regular :class:`torch.nn.Linear` module. The ``in_features`` argument\\n    of the :class:`Linear` is inferred from the ``input.shape[-1]``.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_features: size of each output sample\\n        bias: If set to ``False``, the layer will not learn an additive bias.\\n            Default: ``True``\\n\\n    Attributes:\\n        weight: the learnable weights of the module of shape\\n            :math:`(\\\\text{out\\\\_features}, \\\\text{in\\\\_features})`. The values are\\n            initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`, where\\n            :math:`k = \\\\frac{1}{\\\\text{in\\\\_features}}`\\n        bias:   the learnable bias of the module of shape :math:`(\\\\text{out\\\\_features})`.\\n                If :attr:`bias` is ``True``, the values are initialized from\\n                :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                :math:`k = \\\\frac{1}{\\\\text{in\\\\_features}}`\\n\\n\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = Linear  # type: ignore[assignment]\\n    weight: UninitializedParameter\\n    bias: UninitializedParameter  # type: ignore[assignment]\\n\\n    def __init__(\\n        self, out_features: int, bias: bool = True, device=None, dtype=None\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        # bias is hardcoded to False to avoid creating tensor\\n        # that will soon be overwritten.\\n        super().__init__(0, 0, False)\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_features = out_features\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def reset_parameters(self) -> None:\\n        if not self.has_uninitialized_params() and self.in_features != 0:\\n            super().reset_parameters()\\n\\n    def initialize_parameters(self, input) -> None:  # type: ignore[override]\\n        if self.has_uninitialized_params():\\n            with torch.no_grad():\\n                self.in_features = input.shape[-1]\\n                self.weight.materialize((self.out_features, self.in_features))\\n                if self.bias is not None:\\n                    self.bias.materialize((self.out_features,))\\n                self.reset_parameters()\\n\\n\\n# TODO: PartialLinear - maybe in sparse?\\n\\n\\n# mypy: allow-untyped-decorators\\n# mypy: allow-untyped-defs\\nimport math\\nimport numbers\\nimport warnings\\nimport weakref\\nfrom typing import List, Optional, overload, Tuple\\nfrom typing_extensions import deprecated\\n\\nimport torch\\nfrom torch import _VF, Tensor\\nfrom torch.nn import init\\nfrom torch.nn.parameter import Parameter\\nfrom torch.nn.utils.rnn import PackedSequence\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"RNNBase\\\",\\n    \\\"RNN\\\",\\n    \\\"LSTM\\\",\\n    \\\"GRU\\\",\\n    \\\"RNNCellBase\\\",\\n    \\\"RNNCell\\\",\\n    \\\"LSTMCell\\\",\\n    \\\"GRUCell\\\",\\n]\\n\\n_rnn_impls = {\\n    \\\"RNN_TANH\\\": _VF.rnn_tanh,\\n    \\\"RNN_RELU\\\": _VF.rnn_relu,\\n}\\n\\n\\ndef _apply_permutation(tensor: Tensor, permutation: Tensor, dim: int = 1) -> Tensor:\\n    return tensor.index_select(dim, permutation)\\n\\n\\n@deprecated(\\n    \\\"`apply_permutation` is deprecated, please use `tensor.index_select(dim, permutation)` instead\\\",\\n    category=FutureWarning,\\n)\\ndef apply_permutation(tensor: Tensor, permutation: Tensor, dim: int = 1) -> Tensor:\\n    return _apply_permutation(tensor, permutation, dim)\\n\\n\\nclass RNNBase(Module):\\n    r\\\"\\\"\\\"Base class for RNN modules (RNN, LSTM, GRU).\\n\\n    Implements aspects of RNNs shared by the RNN, LSTM, and GRU classes, such as module initialization\\n    and utility methods for parameter storage management.\\n\\n    .. note::\\n        The forward method is not implemented by the RNNBase class.\\n\\n    .. note::\\n        LSTM and GRU classes override some methods implemented by RNNBase.\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\n        \\\"mode\\\",\\n        \\\"input_size\\\",\\n        \\\"hidden_size\\\",\\n        \\\"num_layers\\\",\\n        \\\"bias\\\",\\n        \\\"batch_first\\\",\\n        \\\"dropout\\\",\\n        \\\"bidirectional\\\",\\n        \\\"proj_size\\\",\\n    ]\\n    __jit_unused_properties__ = [\\\"all_weights\\\"]\\n\\n    mode: str\\n    input_size: int\\n    hidden_size: int\\n    num_layers: int\\n    bias: bool\\n    batch_first: bool\\n    dropout: float\\n    bidirectional: bool\\n    proj_size: int\\n\\n    def __init__(\\n        self,\\n        mode: str,\\n        input_size: int,\\n        hidden_size: int,\\n        num_layers: int = 1,\\n        bias: bool = True,\\n        batch_first: bool = False,\\n        dropout: float = 0.0,\\n        bidirectional: bool = False,\\n        proj_size: int = 0,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.mode = mode\\n        self.input_size = input_size\\n        self.hidden_size = hidden_size\\n        self.num_layers = num_layers\\n        self.bias = bias\\n        self.batch_first = batch_first\\n        self.dropout = float(dropout)\\n        self.bidirectional = bidirectional\\n        self.proj_size = proj_size\\n        self._flat_weight_refs: List[Optional[weakref.ReferenceType[Parameter]]] = []\\n        num_directions = 2 if bidirectional else 1\\n\\n        if (\\n            not isinstance(dropout, numbers.Number)\\n            or not 0 <= dropout <= 1\\n            or isinstance(dropout, bool)\\n        ):\\n            raise ValueError(\\n                \\\"dropout should be a number in range [0, 1] \\\"\\n                \\\"representing the probability of an element being \\\"\\n                \\\"zeroed\\\"\\n            )\\n        if dropout > 0 and num_layers == 1:\\n            warnings.warn(\\n                \\\"dropout option adds dropout after all but last \\\"\\n                \\\"recurrent layer, so non-zero dropout expects \\\"\\n                f\\\"num_layers greater than 1, but got dropout={dropout} and \\\"\\n                f\\\"num_layers={num_layers}\\\"\\n            )\\n\\n        if not isinstance(hidden_size, int):\\n            raise TypeError(\\n                f\\\"hidden_size should be of type int, got: {type(hidden_size).__name__}\\\"\\n            )\\n        if hidden_size <= 0:\\n            raise ValueError(\\\"hidden_size must be greater than zero\\\")\\n        if num_layers <= 0:\\n            raise ValueError(\\\"num_layers must be greater than zero\\\")\\n        if proj_size < 0:\\n            raise ValueError(\\n                \\\"proj_size should be a positive integer or zero to disable projections\\\"\\n            )\\n        if proj_size >= hidden_size:\\n            raise ValueError(\\\"proj_size has to be smaller than hidden_size\\\")\\n\\n        if mode == \\\"LSTM\\\":\\n            gate_size = 4 * hidden_size\\n        elif mode == \\\"GRU\\\":\\n            gate_size = 3 * hidden_size\\n        elif mode == \\\"RNN_TANH\\\":\\n            gate_size = hidden_size\\n        elif mode == \\\"RNN_RELU\\\":\\n            gate_size = hidden_size\\n        else:\\n            raise ValueError(\\\"Unrecognized RNN mode: \\\" + mode)\\n\\n        self._flat_weights_names = []\\n        self._all_weights = []\\n        for layer in range(num_layers):\\n            for direction in range(num_directions):\\n                real_hidden_size = proj_size if proj_size > 0 else hidden_size\\n                layer_input_size = (\\n                    input_size if layer == 0 else real_hidden_size * num_directions\\n                )\\n\\n                w_ih = Parameter(\\n                    torch.empty((gate_size, layer_input_size), **factory_kwargs)\\n                )\\n                w_hh = Parameter(\\n                    torch.empty((gate_size, real_hidden_size), **factory_kwargs)\\n                )\\n                b_ih = Parameter(torch.empty(gate_size, **factory_kwargs))\\n                # Second bias vector included for CuDNN compatibility. Only one\\n                # bias vector is needed in standard definition.\\n                b_hh = Parameter(torch.empty(gate_size, **factory_kwargs))\\n                layer_params: Tuple[Tensor, ...] = ()\\n                if self.proj_size == 0:\\n                    if bias:\\n                        layer_params = (w_ih, w_hh, b_ih, b_hh)\\n                    else:\\n                        layer_params = (w_ih, w_hh)\\n                else:\\n                    w_hr = Parameter(\\n                        torch.empty((proj_size, hidden_size), **factory_kwargs)\\n                    )\\n                    if bias:\\n                        layer_params = (w_ih, w_hh, b_ih, b_hh, w_hr)\\n                    else:\\n                        layer_params = (w_ih, w_hh, w_hr)\\n\\n                suffix = \\\"_reverse\\\" if direction == 1 else \\\"\\\"\\n                param_names = [\\\"weight_ih_l{}{}\\\", \\\"weight_hh_l{}{}\\\"]\\n                if bias:\\n                    param_names += [\\\"bias_ih_l{}{}\\\", \\\"bias_hh_l{}{}\\\"]\\n                if self.proj_size > 0:\\n                    param_names += [\\\"weight_hr_l{}{}\\\"]\\n                param_names = [x.format(layer, suffix) for x in param_names]\\n\\n                for name, param in zip(param_names, layer_params):\\n                    setattr(self, name, param)\\n                self._flat_weights_names.extend(param_names)\\n                self._all_weights.append(param_names)\\n\\n        self._init_flat_weights()\\n\\n        self.reset_parameters()\\n\\n    def _init_flat_weights(self):\\n        self._flat_weights = [\\n            getattr(self, wn) if hasattr(self, wn) else None\\n            for wn in self._flat_weights_names\\n        ]\\n        self._flat_weight_refs = [\\n            weakref.ref(w) if w is not None else None for w in self._flat_weights\\n        ]\\n        self.flatten_parameters()\\n\\n    def __setattr__(self, attr, value):\\n        if hasattr(self, \\\"_flat_weights_names\\\") and attr in self._flat_weights_names:\\n            # keep self._flat_weights up to date if you do self.weight = ...\\n            idx = self._flat_weights_names.index(attr)\\n            self._flat_weights[idx] = value\\n        super().__setattr__(attr, value)\\n\\n    def flatten_parameters(self) -> None:\\n        \\\"\\\"\\\"Reset parameter data pointer so that they can use faster code paths.\\n\\n        Right now, this works only if the module is on the GPU and cuDNN is enabled.\\n        Otherwise, it's a no-op.\\n        \\\"\\\"\\\"\\n        # Short-circuits if _flat_weights is only partially instantiated\\n        if len(self._flat_weights) != len(self._flat_weights_names):\\n            return\\n\\n        for w in self._flat_weights:\\n            if not isinstance(w, Tensor):\\n                return\\n        # Short-circuits if any tensor in self._flat_weights is not acceptable to cuDNN\\n        # or the tensors in _flat_weights are of different dtypes\\n\\n        first_fw = self._flat_weights[0]\\n        dtype = first_fw.dtype\\n        for fw in self._flat_weights:\\n            if (\\n                not isinstance(fw, Tensor)\\n                or not (fw.dtype == dtype)\\n                or not fw.is_cuda\\n                or not torch.backends.cudnn.is_acceptable(fw)\\n            ):\\n                return\\n\\n        # If any parameters alias, we fall back to the slower, copying code path. This is\\n        # a sufficient check, because overlapping parameter buffers that don't completely\\n        # alias would break the assumptions of the uniqueness check in\\n        # Module.named_parameters().\\n        unique_data_ptrs = {p.data_ptr() for p in self._flat_weights}\\n        if len(unique_data_ptrs) != len(self._flat_weights):\\n            return\\n\\n        with torch.cuda.device_of(first_fw):\\n            import torch.backends.cudnn.rnn as rnn\\n\\n            # Note: no_grad() is necessary since _cudnn_rnn_flatten_weight is\\n            # an inplace operation on self._flat_weights\\n            with torch.no_grad():\\n                if torch._use_cudnn_rnn_flatten_weight():\\n                    num_weights = 4 if self.bias else 2\\n                    if self.proj_size > 0:\\n                        num_weights += 1\\n                    torch._cudnn_rnn_flatten_weight(\\n                        self._flat_weights,\\n                        num_weights,\\n                        self.input_size,\\n                        rnn.get_cudnn_mode(self.mode),\\n                        self.hidden_size,\\n                        self.proj_size,\\n                        self.num_layers,\\n                        self.batch_first,\\n                        bool(self.bidirectional),\\n                    )\\n\\n    def _apply(self, fn, recurse=True):\\n        self._flat_weight_refs = []\\n        ret = super()._apply(fn, recurse)\\n\\n        # Resets _flat_weights\\n        # Note: be v. careful before removing this, as 3rd party device types\\n        # likely rely on this behavior to properly .to() modules like LSTM.\\n        self._init_flat_weights()\\n\\n        return ret\\n\\n    def reset_parameters(self) -> None:\\n        stdv = 1.0 / math.sqrt(self.hidden_size) if self.hidden_size > 0 else 0\\n        for weight in self.parameters():\\n            init.uniform_(weight, -stdv, stdv)\\n\\n    def check_input(self, input: Tensor, batch_sizes: Optional[Tensor]) -> None:\\n        if not torch.jit.is_scripting():\\n            if (\\n                input.dtype != self._flat_weights[0].dtype\\n                and not torch._C._is_any_autocast_enabled()\\n            ):\\n                raise ValueError(\\n                    f\\\"input must have the type {self._flat_weights[0].dtype}, got type {input.dtype}\\\"\\n                )\\n        expected_input_dim = 2 if batch_sizes is not None else 3\\n        if input.dim() != expected_input_dim:\\n            raise RuntimeError(\\n                f\\\"input must have {expected_input_dim} dimensions, got {input.dim()}\\\"\\n            )\\n        if self.input_size != input.size(-1):\\n            raise RuntimeError(\\n                f\\\"input.size(-1) must be equal to input_size. Expected {self.input_size}, got {input.size(-1)}\\\"\\n            )\\n\\n    def get_expected_hidden_size(\\n        self, input: Tensor, batch_sizes: Optional[Tensor]\\n    ) -> Tuple[int, int, int]:\\n        if batch_sizes is not None:\\n            mini_batch = int(batch_sizes[0])\\n        else:\\n            mini_batch = input.size(0) if self.batch_first else input.size(1)\\n        num_directions = 2 if self.bidirectional else 1\\n        if self.proj_size > 0:\\n            expected_hidden_size = (\\n                self.num_layers * num_directions,\\n                mini_batch,\\n                self.proj_size,\\n            )\\n        else:\\n            expected_hidden_size = (\\n                self.num_layers * num_directions,\\n                mini_batch,\\n                self.hidden_size,\\n            )\\n        return expected_hidden_size\\n\\n    def check_hidden_size(\\n        self,\\n        hx: Tensor,\\n        expected_hidden_size: Tuple[int, int, int],\\n        msg: str = \\\"Expected hidden size {}, got {}\\\",\\n    ) -> None:\\n        if hx.size() != expected_hidden_size:\\n            raise RuntimeError(msg.format(expected_hidden_size, list(hx.size())))\\n\\n    def _weights_have_changed(self):\\n        # Returns True if the weight tensors have changed since the last forward pass.\\n        # This is the case when used with torch.func.functional_call(), for example.\\n        weights_changed = False\\n        for ref, name in zip(self._flat_weight_refs, self._flat_weights_names):\\n            weight = getattr(self, name) if hasattr(self, name) else None\\n            if weight is not None and ref is not None and ref() is not weight:\\n                weights_changed = True\\n                break\\n        return weights_changed\\n\\n    def check_forward_args(\\n        self, input: Tensor, hidden: Tensor, batch_sizes: Optional[Tensor]\\n    ):\\n        self.check_input(input, batch_sizes)\\n        expected_hidden_size = self.get_expected_hidden_size(input, batch_sizes)\\n\\n        self.check_hidden_size(hidden, expected_hidden_size)\\n\\n    def permute_hidden(self, hx: Tensor, permutation: Optional[Tensor]):\\n        if permutation is None:\\n            return hx\\n        return _apply_permutation(hx, permutation)\\n\\n    def extra_repr(self) -> str:\\n        s = \\\"{input_size}, {hidden_size}\\\"\\n        if self.proj_size != 0:\\n            s += \\\", proj_size={proj_size}\\\"\\n        if self.num_layers != 1:\\n            s += \\\", num_layers={num_layers}\\\"\\n        if self.bias is not True:\\n            s += \\\", bias={bias}\\\"\\n        if self.batch_first is not False:\\n            s += \\\", batch_first={batch_first}\\\"\\n        if self.dropout != 0:\\n            s += \\\", dropout={dropout}\\\"\\n        if self.bidirectional is not False:\\n            s += \\\", bidirectional={bidirectional}\\\"\\n        return s.format(**self.__dict__)\\n\\n    def _update_flat_weights(self):\\n        if not torch.jit.is_scripting():\\n            if self._weights_have_changed():\\n                self._init_flat_weights()\\n\\n    def __getstate__(self):\\n        # If weights have been changed, update the _flat_weights in __getstate__ here.\\n        self._update_flat_weights()\\n        # Don't serialize the weight references.\\n        state = self.__dict__.copy()\\n        del state[\\\"_flat_weight_refs\\\"]\\n        return state\\n\\n    def __setstate__(self, d):\\n        super().__setstate__(d)\\n        if \\\"all_weights\\\" in d:\\n            self._all_weights = d[\\\"all_weights\\\"]\\n        # In PyTorch 1.8 we added a proj_size member variable to LSTM.\\n        # LSTMs that were serialized via torch.save(module) before PyTorch 1.8\\n        # don't have it, so to preserve compatibility we set proj_size here.\\n        if \\\"proj_size\\\" not in d:\\n            self.proj_size = 0\\n\\n        if not isinstance(self._all_weights[0][0], str):\\n            num_layers = self.num_layers\\n            num_directions = 2 if self.bidirectional else 1\\n            self._flat_weights_names = []\\n            self._all_weights = []\\n            for layer in range(num_layers):\\n                for direction in range(num_directions):\\n                    suffix = \\\"_reverse\\\" if direction == 1 else \\\"\\\"\\n                    weights = [\\n                        \\\"weight_ih_l{}{}\\\",\\n                        \\\"weight_hh_l{}{}\\\",\\n                        \\\"bias_ih_l{}{}\\\",\\n                        \\\"bias_hh_l{}{}\\\",\\n                        \\\"weight_hr_l{}{}\\\",\\n                    ]\\n                    weights = [x.format(layer, suffix) for x in weights]\\n                    if self.bias:\\n                        if self.proj_size > 0:\\n                            self._all_weights += [weights]\\n                            self._flat_weights_names.extend(weights)\\n                        else:\\n                            self._all_weights += [weights[:4]]\\n                            self._flat_weights_names.extend(weights[:4])\\n                    else:\\n                        if self.proj_size > 0:\\n                            self._all_weights += [weights[:2]] + [weights[-1:]]\\n                            self._flat_weights_names.extend(\\n                                weights[:2] + [weights[-1:]]\\n                            )\\n                        else:\\n                            self._all_weights += [weights[:2]]\\n                            self._flat_weights_names.extend(weights[:2])\\n            self._flat_weights = [\\n                getattr(self, wn) if hasattr(self, wn) else None\\n                for wn in self._flat_weights_names\\n            ]\\n\\n        self._flat_weight_refs = [\\n            weakref.ref(w) if w is not None else None for w in self._flat_weights\\n        ]\\n\\n    @property\\n    def all_weights(self) -> List[List[Parameter]]:\\n        return [\\n            [getattr(self, weight) for weight in weights]\\n            for weights in self._all_weights\\n        ]\\n\\n    def _replicate_for_data_parallel(self):\\n        replica = super()._replicate_for_data_parallel()\\n        # Need to copy these caches, otherwise the replica will share the same\\n        # flat weights list.\\n        replica._flat_weights = replica._flat_weights[:]\\n        replica._flat_weights_names = replica._flat_weights_names[:]\\n        return replica\\n\\n\\nclass RNN(RNNBase):\\n    r\\\"\\\"\\\"__init__(input_size,hidden_size,num_layers=1,nonlinearity='tanh',bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)\\n\\n    Apply a multi-layer Elman RNN with :math:`\\\\tanh` or :math:`\\\\text{ReLU}`\\n    non-linearity to an input sequence. For each element in the input sequence,\\n    each layer computes the following function:\\n\\n    .. math::\\n        h_t = \\\\tanh(x_t W_{ih}^T + b_{ih} + h_{t-1}W_{hh}^T + b_{hh})\\n\\n    where :math:`h_t` is the hidden state at time `t`, :math:`x_t` is\\n    the input at time `t`, and :math:`h_{(t-1)}` is the hidden state of the\\n    previous layer at time `t-1` or the initial hidden state at time `0`.\\n    If :attr:`nonlinearity` is ``'relu'``, then :math:`\\\\text{ReLU}` is used instead of :math:`\\\\tanh`.\\n\\n    .. code-block:: python\\n\\n        # Efficient implementation equivalent to the following with bidirectional=False\\n        def forward(x, h_0=None):\\n            if batch_first:\\n                x = x.transpose(0, 1)\\n            seq_len, batch_size, _ = x.size()\\n            if h_0 is None:\\n                h_0 = torch.zeros(num_layers, batch_size, hidden_size)\\n            h_t_minus_1 = h_0\\n            h_t = h_0\\n            output = []\\n            for t in range(seq_len):\\n                for layer in range(num_layers):\\n                    h_t[layer] = torch.tanh(\\n                        x[t] @ weight_ih[layer].T\\n                        + bias_ih[layer]\\n                        + h_t_minus_1[layer] @ weight_hh[layer].T\\n                        + bias_hh[layer]\\n                    )\\n                output.append(h_t[-1])\\n                h_t_minus_1 = h_t\\n            output = torch.stack(output)\\n            if batch_first:\\n                output = output.transpose(0, 1)\\n            return output, h_t\\n\\n    Args:\\n        input_size: The number of expected features in the input `x`\\n        hidden_size: The number of features in the hidden state `h`\\n        num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``\\n            would mean stacking two RNNs together to form a `stacked RNN`,\\n            with the second RNN taking in outputs of the first RNN and\\n            computing the final results. Default: 1\\n        nonlinearity: The non-linearity to use. Can be either ``'tanh'`` or ``'relu'``. Default: ``'tanh'``\\n        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.\\n            Default: ``True``\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as `(batch, seq, feature)` instead of `(seq, batch, feature)`.\\n            Note that this does not apply to hidden or cell states. See the\\n            Inputs/Outputs sections below for details.  Default: ``False``\\n        dropout: If non-zero, introduces a `Dropout` layer on the outputs of each\\n            RNN layer except the last layer, with dropout probability equal to\\n            :attr:`dropout`. Default: 0\\n        bidirectional: If ``True``, becomes a bidirectional RNN. Default: ``False``\\n\\n    Inputs: input, h_0\\n        * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,\\n          :math:`(L, N, H_{in})` when ``batch_first=False`` or\\n          :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of\\n          the input sequence.  The input can also be a packed variable length sequence.\\n          See :func:`torch.nn.utils.rnn.pack_padded_sequence` or\\n          :func:`torch.nn.utils.rnn.pack_sequence` for details.\\n        * **h_0**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{out})` for unbatched input or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{out})` containing the initial hidden\\n          state for the input sequence batch. Defaults to zeros if not provided.\\n\\n        where:\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                N ={} & \\\\text{batch size} \\\\\\\\\\n                L ={} & \\\\text{sequence length} \\\\\\\\\\n                D ={} & 2 \\\\text{ if bidirectional=True otherwise } 1 \\\\\\\\\\n                H_{in} ={} & \\\\text{input\\\\_size} \\\\\\\\\\n                H_{out} ={} & \\\\text{hidden\\\\_size}\\n            \\\\end{aligned}\\n\\n    Outputs: output, h_n\\n        * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,\\n          :math:`(L, N, D * H_{out})` when ``batch_first=False`` or\\n          :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features\\n          `(h_t)` from the last layer of the RNN, for each `t`. If a\\n          :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output\\n          will also be a packed sequence.\\n        * **h_n**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{out})` for unbatched input or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{out})` containing the final hidden state\\n          for each element in the batch.\\n\\n    Attributes:\\n        weight_ih_l[k]: the learnable input-hidden weights of the k-th layer,\\n            of shape `(hidden_size, input_size)` for `k = 0`. Otherwise, the shape is\\n            `(hidden_size, num_directions * hidden_size)`\\n        weight_hh_l[k]: the learnable hidden-hidden weights of the k-th layer,\\n            of shape `(hidden_size, hidden_size)`\\n        bias_ih_l[k]: the learnable input-hidden bias of the k-th layer,\\n            of shape `(hidden_size)`\\n        bias_hh_l[k]: the learnable hidden-hidden bias of the k-th layer,\\n            of shape `(hidden_size)`\\n\\n    .. note::\\n        All the weights and biases are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`\\n        where :math:`k = \\\\frac{1}{\\\\text{hidden\\\\_size}}`\\n\\n    .. note::\\n        For bidirectional RNNs, forward and backward are directions 0 and 1 respectively.\\n        Example of splitting the output layers when ``batch_first=False``:\\n        ``output.view(seq_len, batch, num_directions, hidden_size)``.\\n\\n    .. note::\\n        ``batch_first`` argument is ignored for unbatched inputs.\\n\\n    .. include:: ../cudnn_rnn_determinism.rst\\n\\n    .. include:: ../cudnn_persistent_rnn.rst\\n\\n    Examples::\\n\\n        >>> rnn = nn.RNN(10, 20, 2)\\n        >>> input = torch.randn(5, 3, 10)\\n        >>> h0 = torch.randn(2, 3, 20)\\n        >>> output, hn = rnn(input, h0)\\n    \\\"\\\"\\\"\\n\\n    @overload\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        num_layers: int = 1,\\n        nonlinearity: str = \\\"tanh\\\",\\n        bias: bool = True,\\n        batch_first: bool = False,\\n        dropout: float = 0.0,\\n        bidirectional: bool = False,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        ...\\n\\n    @overload\\n    def __init__(self, *args, **kwargs):\\n        ...\\n\\n    def __init__(self, *args, **kwargs):\\n        if \\\"proj_size\\\" in kwargs:\\n            raise ValueError(\\n                \\\"proj_size argument is only supported for LSTM, not RNN or GRU\\\"\\n            )\\n        if len(args) > 3:\\n            self.nonlinearity = args[3]\\n            args = args[:3] + args[4:]\\n        else:\\n            self.nonlinearity = kwargs.pop(\\\"nonlinearity\\\", \\\"tanh\\\")\\n        if self.nonlinearity == \\\"tanh\\\":\\n            mode = \\\"RNN_TANH\\\"\\n        elif self.nonlinearity == \\\"relu\\\":\\n            mode = \\\"RNN_RELU\\\"\\n        else:\\n            raise ValueError(\\n                f\\\"Unknown nonlinearity '{self.nonlinearity}'. Select from 'tanh' or 'relu'.\\\"\\n            )\\n        super().__init__(mode, *args, **kwargs)\\n\\n    @overload\\n    @torch._jit_internal._overload_method  # noqa: F811\\n    def forward(\\n        self, input: Tensor, hx: Optional[Tensor] = None\\n    ) -> Tuple[Tensor, Tensor]:\\n        pass\\n\\n    @overload\\n    @torch._jit_internal._overload_method  # noqa: F811\\n    def forward(\\n        self, input: PackedSequence, hx: Optional[Tensor] = None\\n    ) -> Tuple[PackedSequence, Tensor]:\\n        pass\\n\\n    def forward(self, input, hx=None):  # noqa: F811\\n        self._update_flat_weights()\\n\\n        num_directions = 2 if self.bidirectional else 1\\n        orig_input = input\\n\\n        if isinstance(orig_input, PackedSequence):\\n            input, batch_sizes, sorted_indices, unsorted_indices = input\\n            max_batch_size = batch_sizes[0]\\n            # script() is unhappy when max_batch_size is different type in cond branches, so we duplicate\\n            if hx is None:\\n                hx = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    self.hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n            else:\\n                # Each batch of the hidden state should match the input sequence that\\n                # the user believes he/she is passing in.\\n                hx = self.permute_hidden(hx, sorted_indices)\\n        else:\\n            batch_sizes = None\\n            if input.dim() not in (2, 3):\\n                raise ValueError(\\n                    f\\\"RNN: Expected input to be 2D or 3D, got {input.dim()}D tensor instead\\\"\\n                )\\n            is_batched = input.dim() == 3\\n            batch_dim = 0 if self.batch_first else 1\\n            if not is_batched:\\n                input = input.unsqueeze(batch_dim)\\n                if hx is not None:\\n                    if hx.dim() != 2:\\n                        raise RuntimeError(\\n                            f\\\"For unbatched 2-D input, hx should also be 2-D but got {hx.dim()}-D tensor\\\"\\n                        )\\n                    hx = hx.unsqueeze(1)\\n            else:\\n                if hx is not None and hx.dim() != 3:\\n                    raise RuntimeError(\\n                        f\\\"For batched 3-D input, hx should also be 3-D but got {hx.dim()}-D tensor\\\"\\n                    )\\n            max_batch_size = input.size(0) if self.batch_first else input.size(1)\\n            sorted_indices = None\\n            unsorted_indices = None\\n            if hx is None:\\n                hx = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    self.hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n            else:\\n                # Each batch of the hidden state should match the input sequence that\\n                # the user believes he/she is passing in.\\n                hx = self.permute_hidden(hx, sorted_indices)\\n\\n        assert hx is not None\\n        self.check_forward_args(input, hx, batch_sizes)\\n        assert self.mode == \\\"RNN_TANH\\\" or self.mode == \\\"RNN_RELU\\\"\\n        if batch_sizes is None:\\n            if self.mode == \\\"RNN_TANH\\\":\\n                result = _VF.rnn_tanh(\\n                    input,\\n                    hx,\\n                    self._flat_weights,\\n                    self.bias,\\n                    self.num_layers,\\n                    self.dropout,\\n                    self.training,\\n                    self.bidirectional,\\n                    self.batch_first,\\n                )\\n            else:\\n                result = _VF.rnn_relu(\\n                    input,\\n                    hx,\\n                    self._flat_weights,\\n                    self.bias,\\n                    self.num_layers,\\n                    self.dropout,\\n                    self.training,\\n                    self.bidirectional,\\n                    self.batch_first,\\n                )\\n        else:\\n            if self.mode == \\\"RNN_TANH\\\":\\n                result = _VF.rnn_tanh(\\n                    input,\\n                    batch_sizes,\\n                    hx,\\n                    self._flat_weights,\\n                    self.bias,\\n                    self.num_layers,\\n                    self.dropout,\\n                    self.training,\\n                    self.bidirectional,\\n                )\\n            else:\\n                result = _VF.rnn_relu(\\n                    input,\\n                    batch_sizes,\\n                    hx,\\n                    self._flat_weights,\\n                    self.bias,\\n                    self.num_layers,\\n                    self.dropout,\\n                    self.training,\\n                    self.bidirectional,\\n                )\\n\\n        output = result[0]\\n        hidden = result[1]\\n\\n        if isinstance(orig_input, PackedSequence):\\n            output_packed = PackedSequence(\\n                output, batch_sizes, sorted_indices, unsorted_indices\\n            )\\n            return output_packed, self.permute_hidden(hidden, unsorted_indices)\\n\\n        if not is_batched:  # type: ignore[possibly-undefined]\\n            output = output.squeeze(batch_dim)  # type: ignore[possibly-undefined]\\n            hidden = hidden.squeeze(1)\\n\\n        return output, self.permute_hidden(hidden, unsorted_indices)\\n\\n\\n# XXX: LSTM and GRU implementation is different from RNNBase, this is because:\\n# 1. we want to support nn.LSTM and nn.GRU in TorchScript and TorchScript in\\n#    its current state could not support the python Union Type or Any Type\\n# 2. TorchScript static typing does not allow a Function or Callable type in\\n#    Dict values, so we have to separately call _VF instead of using _rnn_impls\\n# 3. This is temporary only and in the transition state that we want to make it\\n#    on time for the release\\n#\\n# More discussion details in https://github.com/pytorch/pytorch/pull/23266\\n#\\n# TODO: remove the overriding implementations for LSTM and GRU when TorchScript\\n# support expressing these two modules generally.\\n\\n\\nclass LSTM(RNNBase):\\n    r\\\"\\\"\\\"__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,proj_size=0,device=None,dtype=None)\\n\\n    Apply a multi-layer long short-term memory (LSTM) RNN to an input sequence.\\n    For each element in the input sequence, each layer computes the following\\n    function:\\n\\n    .. math::\\n        \\\\begin{array}{ll} \\\\\\\\\\n            i_t = \\\\sigma(W_{ii} x_t + b_{ii} + W_{hi} h_{t-1} + b_{hi}) \\\\\\\\\\n            f_t = \\\\sigma(W_{if} x_t + b_{if} + W_{hf} h_{t-1} + b_{hf}) \\\\\\\\\\n            g_t = \\\\tanh(W_{ig} x_t + b_{ig} + W_{hg} h_{t-1} + b_{hg}) \\\\\\\\\\n            o_t = \\\\sigma(W_{io} x_t + b_{io} + W_{ho} h_{t-1} + b_{ho}) \\\\\\\\\\n            c_t = f_t \\\\odot c_{t-1} + i_t \\\\odot g_t \\\\\\\\\\n            h_t = o_t \\\\odot \\\\tanh(c_t) \\\\\\\\\\n        \\\\end{array}\\n\\n    where :math:`h_t` is the hidden state at time `t`, :math:`c_t` is the cell\\n    state at time `t`, :math:`x_t` is the input at time `t`, :math:`h_{t-1}`\\n    is the hidden state of the layer at time `t-1` or the initial hidden\\n    state at time `0`, and :math:`i_t`, :math:`f_t`, :math:`g_t`,\\n    :math:`o_t` are the input, forget, cell, and output gates, respectively.\\n    :math:`\\\\sigma` is the sigmoid function, and :math:`\\\\odot` is the Hadamard product.\\n\\n    In a multilayer LSTM, the input :math:`x^{(l)}_t` of the :math:`l` -th layer\\n    (:math:`l \\\\ge 2`) is the hidden state :math:`h^{(l-1)}_t` of the previous layer multiplied by\\n    dropout :math:`\\\\delta^{(l-1)}_t` where each :math:`\\\\delta^{(l-1)}_t` is a Bernoulli random\\n    variable which is :math:`0` with probability :attr:`dropout`.\\n\\n    If ``proj_size > 0`` is specified, LSTM with projections will be used. This changes\\n    the LSTM cell in the following way. First, the dimension of :math:`h_t` will be changed from\\n    ``hidden_size`` to ``proj_size`` (dimensions of :math:`W_{hi}` will be changed accordingly).\\n    Second, the output hidden state of each layer will be multiplied by a learnable projection\\n    matrix: :math:`h_t = W_{hr}h_t`. Note that as a consequence of this, the output\\n    of LSTM network will be of different shape as well. See Inputs/Outputs sections below for exact\\n    dimensions of all variables. You can find more details in https://arxiv.org/abs/1402.1128.\\n\\n    Args:\\n        input_size: The number of expected features in the input `x`\\n        hidden_size: The number of features in the hidden state `h`\\n        num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``\\n            would mean stacking two LSTMs together to form a `stacked LSTM`,\\n            with the second LSTM taking in outputs of the first LSTM and\\n            computing the final results. Default: 1\\n        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.\\n            Default: ``True``\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as `(batch, seq, feature)` instead of `(seq, batch, feature)`.\\n            Note that this does not apply to hidden or cell states. See the\\n            Inputs/Outputs sections below for details.  Default: ``False``\\n        dropout: If non-zero, introduces a `Dropout` layer on the outputs of each\\n            LSTM layer except the last layer, with dropout probability equal to\\n            :attr:`dropout`. Default: 0\\n        bidirectional: If ``True``, becomes a bidirectional LSTM. Default: ``False``\\n        proj_size: If ``> 0``, will use LSTM with projections of corresponding size. Default: 0\\n\\n    Inputs: input, (h_0, c_0)\\n        * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,\\n          :math:`(L, N, H_{in})` when ``batch_first=False`` or\\n          :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of\\n          the input sequence.  The input can also be a packed variable length sequence.\\n          See :func:`torch.nn.utils.rnn.pack_padded_sequence` or\\n          :func:`torch.nn.utils.rnn.pack_sequence` for details.\\n        * **h_0**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{out})` for unbatched input or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{out})` containing the\\n          initial hidden state for each element in the input sequence.\\n          Defaults to zeros if (h_0, c_0) is not provided.\\n        * **c_0**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{cell})` for unbatched input or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{cell})` containing the\\n          initial cell state for each element in the input sequence.\\n          Defaults to zeros if (h_0, c_0) is not provided.\\n\\n        where:\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                N ={} & \\\\text{batch size} \\\\\\\\\\n                L ={} & \\\\text{sequence length} \\\\\\\\\\n                D ={} & 2 \\\\text{ if bidirectional=True otherwise } 1 \\\\\\\\\\n                H_{in} ={} & \\\\text{input\\\\_size} \\\\\\\\\\n                H_{cell} ={} & \\\\text{hidden\\\\_size} \\\\\\\\\\n                H_{out} ={} & \\\\text{proj\\\\_size if } \\\\text{proj\\\\_size}>0 \\\\text{ otherwise hidden\\\\_size} \\\\\\\\\\n            \\\\end{aligned}\\n\\n    Outputs: output, (h_n, c_n)\\n        * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,\\n          :math:`(L, N, D * H_{out})` when ``batch_first=False`` or\\n          :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features\\n          `(h_t)` from the last layer of the LSTM, for each `t`. If a\\n          :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output\\n          will also be a packed sequence. When ``bidirectional=True``, `output` will contain\\n          a concatenation of the forward and reverse hidden states at each time step in the sequence.\\n        * **h_n**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{out})` for unbatched input or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{out})` containing the\\n          final hidden state for each element in the sequence. When ``bidirectional=True``,\\n          `h_n` will contain a concatenation of the final forward and reverse hidden states, respectively.\\n        * **c_n**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{cell})` for unbatched input or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{cell})` containing the\\n          final cell state for each element in the sequence. When ``bidirectional=True``,\\n          `c_n` will contain a concatenation of the final forward and reverse cell states, respectively.\\n\\n    Attributes:\\n        weight_ih_l[k] : the learnable input-hidden weights of the :math:`\\\\text{k}^{th}` layer\\n            `(W_ii|W_if|W_ig|W_io)`, of shape `(4*hidden_size, input_size)` for `k = 0`.\\n            Otherwise, the shape is `(4*hidden_size, num_directions * hidden_size)`. If\\n            ``proj_size > 0`` was specified, the shape will be\\n            `(4*hidden_size, num_directions * proj_size)` for `k > 0`\\n        weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\\\\text{k}^{th}` layer\\n            `(W_hi|W_hf|W_hg|W_ho)`, of shape `(4*hidden_size, hidden_size)`. If ``proj_size > 0``\\n            was specified, the shape will be `(4*hidden_size, proj_size)`.\\n        bias_ih_l[k] : the learnable input-hidden bias of the :math:`\\\\text{k}^{th}` layer\\n            `(b_ii|b_if|b_ig|b_io)`, of shape `(4*hidden_size)`\\n        bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\\\\text{k}^{th}` layer\\n            `(b_hi|b_hf|b_hg|b_ho)`, of shape `(4*hidden_size)`\\n        weight_hr_l[k] : the learnable projection weights of the :math:`\\\\text{k}^{th}` layer\\n            of shape `(proj_size, hidden_size)`. Only present when ``proj_size > 0`` was\\n            specified.\\n        weight_ih_l[k]_reverse: Analogous to `weight_ih_l[k]` for the reverse direction.\\n            Only present when ``bidirectional=True``.\\n        weight_hh_l[k]_reverse:  Analogous to `weight_hh_l[k]` for the reverse direction.\\n            Only present when ``bidirectional=True``.\\n        bias_ih_l[k]_reverse:  Analogous to `bias_ih_l[k]` for the reverse direction.\\n            Only present when ``bidirectional=True``.\\n        bias_hh_l[k]_reverse:  Analogous to `bias_hh_l[k]` for the reverse direction.\\n            Only present when ``bidirectional=True``.\\n        weight_hr_l[k]_reverse:  Analogous to `weight_hr_l[k]` for the reverse direction.\\n            Only present when ``bidirectional=True`` and ``proj_size > 0`` was specified.\\n\\n    .. note::\\n        All the weights and biases are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`\\n        where :math:`k = \\\\frac{1}{\\\\text{hidden\\\\_size}}`\\n\\n    .. note::\\n        For bidirectional LSTMs, forward and backward are directions 0 and 1 respectively.\\n        Example of splitting the output layers when ``batch_first=False``:\\n        ``output.view(seq_len, batch, num_directions, hidden_size)``.\\n\\n    .. note::\\n        For bidirectional LSTMs, `h_n` is not equivalent to the last element of `output`; the\\n        former contains the final forward and reverse hidden states, while the latter contains the\\n        final forward hidden state and the initial reverse hidden state.\\n\\n    .. note::\\n        ``batch_first`` argument is ignored for unbatched inputs.\\n\\n    .. note::\\n        ``proj_size`` should be smaller than ``hidden_size``.\\n\\n    .. include:: ../cudnn_rnn_determinism.rst\\n\\n    .. include:: ../cudnn_persistent_rnn.rst\\n\\n    Examples::\\n\\n        >>> rnn = nn.LSTM(10, 20, 2)\\n        >>> input = torch.randn(5, 3, 10)\\n        >>> h0 = torch.randn(2, 3, 20)\\n        >>> c0 = torch.randn(2, 3, 20)\\n        >>> output, (hn, cn) = rnn(input, (h0, c0))\\n    \\\"\\\"\\\"\\n\\n    @overload\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        num_layers: int = 1,\\n        bias: bool = True,\\n        batch_first: bool = False,\\n        dropout: float = 0.0,\\n        bidirectional: bool = False,\\n        proj_size: int = 0,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        ...\\n\\n    @overload\\n    def __init__(self, *args, **kwargs):\\n        ...\\n\\n    def __init__(self, *args, **kwargs):\\n        super().__init__(\\\"LSTM\\\", *args, **kwargs)\\n\\n    def get_expected_cell_size(\\n        self, input: Tensor, batch_sizes: Optional[Tensor]\\n    ) -> Tuple[int, int, int]:\\n        if batch_sizes is not None:\\n            mini_batch = int(batch_sizes[0])\\n        else:\\n            mini_batch = input.size(0) if self.batch_first else input.size(1)\\n        num_directions = 2 if self.bidirectional else 1\\n        expected_hidden_size = (\\n            self.num_layers * num_directions,\\n            mini_batch,\\n            self.hidden_size,\\n        )\\n        return expected_hidden_size\\n\\n    # In the future, we should prevent mypy from applying contravariance rules here.\\n    # See torch/nn/modules/module.py::_forward_unimplemented\\n    def check_forward_args(\\n        self,\\n        input: Tensor,\\n        hidden: Tuple[Tensor, Tensor],  # type: ignore[override]\\n        batch_sizes: Optional[Tensor],\\n    ):\\n        self.check_input(input, batch_sizes)\\n        self.check_hidden_size(\\n            hidden[0],\\n            self.get_expected_hidden_size(input, batch_sizes),\\n            \\\"Expected hidden[0] size {}, got {}\\\",\\n        )\\n        self.check_hidden_size(\\n            hidden[1],\\n            self.get_expected_cell_size(input, batch_sizes),\\n            \\\"Expected hidden[1] size {}, got {}\\\",\\n        )\\n\\n    # Same as above, see torch/nn/modules/module.py::_forward_unimplemented\\n    def permute_hidden(  # type: ignore[override]\\n        self,\\n        hx: Tuple[Tensor, Tensor],\\n        permutation: Optional[Tensor],\\n    ) -> Tuple[Tensor, Tensor]:\\n        if permutation is None:\\n            return hx\\n        return _apply_permutation(hx[0], permutation), _apply_permutation(\\n            hx[1], permutation\\n        )\\n\\n    # Same as above, see torch/nn/modules/module.py::_forward_unimplemented\\n    @overload  # type: ignore[override]\\n    @torch._jit_internal._overload_method  # noqa: F811\\n    def forward(\\n        self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = None\\n    ) -> Tuple[Tensor, Tuple[Tensor, Tensor]]:  # noqa: F811\\n        pass\\n\\n    # Same as above, see torch/nn/modules/module.py::_forward_unimplemented\\n    @overload\\n    @torch._jit_internal._overload_method  # noqa: F811\\n    def forward(\\n        self, input: PackedSequence, hx: Optional[Tuple[Tensor, Tensor]] = None\\n    ) -> Tuple[PackedSequence, Tuple[Tensor, Tensor]]:  # noqa: F811\\n        pass\\n\\n    def forward(self, input, hx=None):  # noqa: F811\\n        self._update_flat_weights()\\n\\n        orig_input = input\\n        # xxx: isinstance check needs to be in conditional for TorchScript to compile\\n        batch_sizes = None\\n        do_permute = False\\n        num_directions = 2 if self.bidirectional else 1\\n        real_hidden_size = self.proj_size if self.proj_size > 0 else self.hidden_size\\n        if isinstance(orig_input, PackedSequence):\\n            input, batch_sizes, sorted_indices, unsorted_indices = input\\n            max_batch_size = batch_sizes[0]\\n            if hx is None:\\n                h_zeros = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    real_hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n                c_zeros = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    self.hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n                hx = (h_zeros, c_zeros)\\n            else:\\n                # Each batch of the hidden state should match the input sequence that\\n                # the user believes he/she is passing in.\\n                hx = self.permute_hidden(hx, sorted_indices)\\n        else:\\n            if input.dim() not in (2, 3):\\n                raise ValueError(\\n                    f\\\"LSTM: Expected input to be 2D or 3D, got {input.dim()}D instead\\\"\\n                )\\n            is_batched = input.dim() == 3\\n            batch_dim = 0 if self.batch_first else 1\\n            if not is_batched:\\n                input = input.unsqueeze(batch_dim)\\n            max_batch_size = input.size(0) if self.batch_first else input.size(1)\\n            sorted_indices = None\\n            unsorted_indices = None\\n            if hx is None:\\n                h_zeros = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    real_hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n                c_zeros = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    self.hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n                hx = (h_zeros, c_zeros)\\n                self.check_forward_args(input, hx, batch_sizes)\\n            else:\\n                if is_batched:\\n                    if hx[0].dim() != 3 or hx[1].dim() != 3:\\n                        msg = (\\n                            \\\"For batched 3-D input, hx and cx should \\\"\\n                            f\\\"also be 3-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors\\\"\\n                        )\\n                        raise RuntimeError(msg)\\n                else:\\n                    if hx[0].dim() != 2 or hx[1].dim() != 2:\\n                        msg = (\\n                            \\\"For unbatched 2-D input, hx and cx should \\\"\\n                            f\\\"also be 2-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors\\\"\\n                        )\\n                        raise RuntimeError(msg)\\n                    hx = (hx[0].unsqueeze(1), hx[1].unsqueeze(1))\\n                # Each batch of the hidden state should match the input sequence that\\n                # the user believes he/she is passing in.\\n                self.check_forward_args(input, hx, batch_sizes)\\n                hx = self.permute_hidden(hx, sorted_indices)\\n\\n        if batch_sizes is None:\\n            result = _VF.lstm(\\n                input,\\n                hx,\\n                self._flat_weights,\\n                self.bias,\\n                self.num_layers,\\n                self.dropout,\\n                self.training,\\n                self.bidirectional,\\n                self.batch_first,\\n            )\\n        else:\\n            result = _VF.lstm(\\n                input,\\n                batch_sizes,\\n                hx,\\n                self._flat_weights,\\n                self.bias,\\n                self.num_layers,\\n                self.dropout,\\n                self.training,\\n                self.bidirectional,\\n            )\\n        output = result[0]\\n        hidden = result[1:]\\n        # xxx: isinstance check needs to be in conditional for TorchScript to compile\\n        if isinstance(orig_input, PackedSequence):\\n            output_packed = PackedSequence(\\n                output, batch_sizes, sorted_indices, unsorted_indices\\n            )\\n            return output_packed, self.permute_hidden(hidden, unsorted_indices)\\n        else:\\n            if not is_batched:  # type: ignore[possibly-undefined]\\n                output = output.squeeze(batch_dim)  # type: ignore[possibly-undefined]\\n                hidden = (hidden[0].squeeze(1), hidden[1].squeeze(1))\\n            return output, self.permute_hidden(hidden, unsorted_indices)\\n\\n\\nclass GRU(RNNBase):\\n    r\\\"\\\"\\\"__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)\\n\\n    Apply a multi-layer gated recurrent unit (GRU) RNN to an input sequence.\\n    For each element in the input sequence, each layer computes the following\\n    function:\\n\\n    .. math::\\n        \\\\begin{array}{ll}\\n            r_t = \\\\sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{(t-1)} + b_{hr}) \\\\\\\\\\n            z_t = \\\\sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{(t-1)} + b_{hz}) \\\\\\\\\\n            n_t = \\\\tanh(W_{in} x_t + b_{in} + r_t \\\\odot (W_{hn} h_{(t-1)}+ b_{hn})) \\\\\\\\\\n            h_t = (1 - z_t) \\\\odot n_t + z_t \\\\odot h_{(t-1)}\\n        \\\\end{array}\\n\\n    where :math:`h_t` is the hidden state at time `t`, :math:`x_t` is the input\\n    at time `t`, :math:`h_{(t-1)}` is the hidden state of the layer\\n    at time `t-1` or the initial hidden state at time `0`, and :math:`r_t`,\\n    :math:`z_t`, :math:`n_t` are the reset, update, and new gates, respectively.\\n    :math:`\\\\sigma` is the sigmoid function, and :math:`\\\\odot` is the Hadamard product.\\n\\n    In a multilayer GRU, the input :math:`x^{(l)}_t` of the :math:`l` -th layer\\n    (:math:`l \\\\ge 2`) is the hidden state :math:`h^{(l-1)}_t` of the previous layer multiplied by\\n    dropout :math:`\\\\delta^{(l-1)}_t` where each :math:`\\\\delta^{(l-1)}_t` is a Bernoulli random\\n    variable which is :math:`0` with probability :attr:`dropout`.\\n\\n    Args:\\n        input_size: The number of expected features in the input `x`\\n        hidden_size: The number of features in the hidden state `h`\\n        num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``\\n            would mean stacking two GRUs together to form a `stacked GRU`,\\n            with the second GRU taking in outputs of the first GRU and\\n            computing the final results. Default: 1\\n        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.\\n            Default: ``True``\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as `(batch, seq, feature)` instead of `(seq, batch, feature)`.\\n            Note that this does not apply to hidden or cell states. See the\\n            Inputs/Outputs sections below for details.  Default: ``False``\\n        dropout: If non-zero, introduces a `Dropout` layer on the outputs of each\\n            GRU layer except the last layer, with dropout probability equal to\\n            :attr:`dropout`. Default: 0\\n        bidirectional: If ``True``, becomes a bidirectional GRU. Default: ``False``\\n\\n    Inputs: input, h_0\\n        * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,\\n          :math:`(L, N, H_{in})` when ``batch_first=False`` or\\n          :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of\\n          the input sequence.  The input can also be a packed variable length sequence.\\n          See :func:`torch.nn.utils.rnn.pack_padded_sequence` or\\n          :func:`torch.nn.utils.rnn.pack_sequence` for details.\\n        * **h_0**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{out})` or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{out})`\\n          containing the initial hidden state for the input sequence. Defaults to zeros if not provided.\\n\\n        where:\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                N ={} & \\\\text{batch size} \\\\\\\\\\n                L ={} & \\\\text{sequence length} \\\\\\\\\\n                D ={} & 2 \\\\text{ if bidirectional=True otherwise } 1 \\\\\\\\\\n                H_{in} ={} & \\\\text{input\\\\_size} \\\\\\\\\\n                H_{out} ={} & \\\\text{hidden\\\\_size}\\n            \\\\end{aligned}\\n\\n    Outputs: output, h_n\\n        * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,\\n          :math:`(L, N, D * H_{out})` when ``batch_first=False`` or\\n          :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features\\n          `(h_t)` from the last layer of the GRU, for each `t`. If a\\n          :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output\\n          will also be a packed sequence.\\n        * **h_n**: tensor of shape :math:`(D * \\\\text{num\\\\_layers}, H_{out})` or\\n          :math:`(D * \\\\text{num\\\\_layers}, N, H_{out})` containing the final hidden state\\n          for the input sequence.\\n\\n    Attributes:\\n        weight_ih_l[k] : the learnable input-hidden weights of the :math:`\\\\text{k}^{th}` layer\\n            (W_ir|W_iz|W_in), of shape `(3*hidden_size, input_size)` for `k = 0`.\\n            Otherwise, the shape is `(3*hidden_size, num_directions * hidden_size)`\\n        weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\\\\text{k}^{th}` layer\\n            (W_hr|W_hz|W_hn), of shape `(3*hidden_size, hidden_size)`\\n        bias_ih_l[k] : the learnable input-hidden bias of the :math:`\\\\text{k}^{th}` layer\\n            (b_ir|b_iz|b_in), of shape `(3*hidden_size)`\\n        bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\\\\text{k}^{th}` layer\\n            (b_hr|b_hz|b_hn), of shape `(3*hidden_size)`\\n\\n    .. note::\\n        All the weights and biases are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`\\n        where :math:`k = \\\\frac{1}{\\\\text{hidden\\\\_size}}`\\n\\n    .. note::\\n        For bidirectional GRUs, forward and backward are directions 0 and 1 respectively.\\n        Example of splitting the output layers when ``batch_first=False``:\\n        ``output.view(seq_len, batch, num_directions, hidden_size)``.\\n\\n    .. note::\\n        ``batch_first`` argument is ignored for unbatched inputs.\\n\\n    .. note::\\n        The calculation of new gate :math:`n_t` subtly differs from the original paper and other frameworks.\\n        In the original implementation, the Hadamard product :math:`(\\\\odot)` between :math:`r_t` and the\\n        previous hidden state :math:`h_{(t-1)}` is done before the multiplication with the weight matrix\\n        `W` and addition of bias:\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                n_t = \\\\tanh(W_{in} x_t + b_{in} + W_{hn} ( r_t \\\\odot h_{(t-1)} ) + b_{hn})\\n            \\\\end{aligned}\\n\\n        This is in contrast to PyTorch implementation, which is done after :math:`W_{hn} h_{(t-1)}`\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                n_t = \\\\tanh(W_{in} x_t + b_{in} + r_t \\\\odot (W_{hn} h_{(t-1)}+ b_{hn}))\\n            \\\\end{aligned}\\n\\n        This implementation differs on purpose for efficiency.\\n\\n    .. include:: ../cudnn_persistent_rnn.rst\\n\\n    Examples::\\n\\n        >>> rnn = nn.GRU(10, 20, 2)\\n        >>> input = torch.randn(5, 3, 10)\\n        >>> h0 = torch.randn(2, 3, 20)\\n        >>> output, hn = rnn(input, h0)\\n    \\\"\\\"\\\"\\n\\n    @overload\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        num_layers: int = 1,\\n        bias: bool = True,\\n        batch_first: bool = False,\\n        dropout: float = 0.0,\\n        bidirectional: bool = False,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        ...\\n\\n    @overload\\n    def __init__(self, *args, **kwargs):\\n        ...\\n\\n    def __init__(self, *args, **kwargs):\\n        if \\\"proj_size\\\" in kwargs:\\n            raise ValueError(\\n                \\\"proj_size argument is only supported for LSTM, not RNN or GRU\\\"\\n            )\\n        super().__init__(\\\"GRU\\\", *args, **kwargs)\\n\\n    @overload  # type: ignore[override]\\n    @torch._jit_internal._overload_method  # noqa: F811\\n    def forward(\\n        self, input: Tensor, hx: Optional[Tensor] = None\\n    ) -> Tuple[Tensor, Tensor]:  # noqa: F811\\n        pass\\n\\n    @overload\\n    @torch._jit_internal._overload_method  # noqa: F811\\n    def forward(\\n        self, input: PackedSequence, hx: Optional[Tensor] = None\\n    ) -> Tuple[PackedSequence, Tensor]:  # noqa: F811\\n        pass\\n\\n    def forward(self, input, hx=None):  # noqa: F811\\n        self._update_flat_weights()\\n\\n        orig_input = input\\n        # xxx: isinstance check needs to be in conditional for TorchScript to compile\\n        if isinstance(orig_input, PackedSequence):\\n            input, batch_sizes, sorted_indices, unsorted_indices = input\\n            max_batch_size = batch_sizes[0]\\n            if hx is None:\\n                num_directions = 2 if self.bidirectional else 1\\n                hx = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    self.hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n            else:\\n                # Each batch of the hidden state should match the input sequence that\\n                # the user believes he/she is passing in.\\n                hx = self.permute_hidden(hx, sorted_indices)\\n        else:\\n            batch_sizes = None\\n            if input.dim() not in (2, 3):\\n                raise ValueError(\\n                    f\\\"GRU: Expected input to be 2D or 3D, got {input.dim()}D instead\\\"\\n                )\\n            is_batched = input.dim() == 3\\n            batch_dim = 0 if self.batch_first else 1\\n            if not is_batched:\\n                input = input.unsqueeze(batch_dim)\\n                if hx is not None:\\n                    if hx.dim() != 2:\\n                        raise RuntimeError(\\n                            f\\\"For unbatched 2-D input, hx should also be 2-D but got {hx.dim()}-D tensor\\\"\\n                        )\\n                    hx = hx.unsqueeze(1)\\n            else:\\n                if hx is not None and hx.dim() != 3:\\n                    raise RuntimeError(\\n                        f\\\"For batched 3-D input, hx should also be 3-D but got {hx.dim()}-D tensor\\\"\\n                    )\\n            max_batch_size = input.size(0) if self.batch_first else input.size(1)\\n            sorted_indices = None\\n            unsorted_indices = None\\n            if hx is None:\\n                num_directions = 2 if self.bidirectional else 1\\n                hx = torch.zeros(\\n                    self.num_layers * num_directions,\\n                    max_batch_size,\\n                    self.hidden_size,\\n                    dtype=input.dtype,\\n                    device=input.device,\\n                )\\n            else:\\n                # Each batch of the hidden state should match the input sequence that\\n                # the user believes he/she is passing in.\\n                hx = self.permute_hidden(hx, sorted_indices)\\n\\n        self.check_forward_args(input, hx, batch_sizes)\\n        if batch_sizes is None:\\n            result = _VF.gru(\\n                input,\\n                hx,\\n                self._flat_weights,\\n                self.bias,\\n                self.num_layers,\\n                self.dropout,\\n                self.training,\\n                self.bidirectional,\\n                self.batch_first,\\n            )\\n        else:\\n            result = _VF.gru(\\n                input,\\n                batch_sizes,\\n                hx,\\n                self._flat_weights,\\n                self.bias,\\n                self.num_layers,\\n                self.dropout,\\n                self.training,\\n                self.bidirectional,\\n            )\\n        output = result[0]\\n        hidden = result[1]\\n\\n        # xxx: isinstance check needs to be in conditional for TorchScript to compile\\n        if isinstance(orig_input, PackedSequence):\\n            output_packed = PackedSequence(\\n                output, batch_sizes, sorted_indices, unsorted_indices\\n            )\\n            return output_packed, self.permute_hidden(hidden, unsorted_indices)\\n        else:\\n            if not is_batched:  # type: ignore[possibly-undefined]\\n                output = output.squeeze(batch_dim)  # type: ignore[possibly-undefined]\\n                hidden = hidden.squeeze(1)\\n\\n            return output, self.permute_hidden(hidden, unsorted_indices)\\n\\n\\nclass RNNCellBase(Module):\\n    __constants__ = [\\\"input_size\\\", \\\"hidden_size\\\", \\\"bias\\\"]\\n\\n    input_size: int\\n    hidden_size: int\\n    bias: bool\\n    weight_ih: Tensor\\n    weight_hh: Tensor\\n    # WARNING: bias_ih and bias_hh purposely not defined here.\\n    # See https://github.com/pytorch/pytorch/issues/39670\\n\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        bias: bool,\\n        num_chunks: int,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.input_size = input_size\\n        self.hidden_size = hidden_size\\n        self.bias = bias\\n        self.weight_ih = Parameter(\\n            torch.empty((num_chunks * hidden_size, input_size), **factory_kwargs)\\n        )\\n        self.weight_hh = Parameter(\\n            torch.empty((num_chunks * hidden_size, hidden_size), **factory_kwargs)\\n        )\\n        if bias:\\n            self.bias_ih = Parameter(\\n                torch.empty(num_chunks * hidden_size, **factory_kwargs)\\n            )\\n            self.bias_hh = Parameter(\\n                torch.empty(num_chunks * hidden_size, **factory_kwargs)\\n            )\\n        else:\\n            self.register_parameter(\\\"bias_ih\\\", None)\\n            self.register_parameter(\\\"bias_hh\\\", None)\\n\\n        self.reset_parameters()\\n\\n    def extra_repr(self) -> str:\\n        s = \\\"{input_size}, {hidden_size}\\\"\\n        if \\\"bias\\\" in self.__dict__ and self.bias is not True:\\n            s += \\\", bias={bias}\\\"\\n        if \\\"nonlinearity\\\" in self.__dict__ and self.nonlinearity != \\\"tanh\\\":\\n            s += \\\", nonlinearity={nonlinearity}\\\"\\n        return s.format(**self.__dict__)\\n\\n    def reset_parameters(self) -> None:\\n        stdv = 1.0 / math.sqrt(self.hidden_size) if self.hidden_size > 0 else 0\\n        for weight in self.parameters():\\n            init.uniform_(weight, -stdv, stdv)\\n\\n\\nclass RNNCell(RNNCellBase):\\n    r\\\"\\\"\\\"An Elman RNN cell with tanh or ReLU non-linearity.\\n\\n    .. math::\\n\\n        h' = \\\\tanh(W_{ih} x + b_{ih}  +  W_{hh} h + b_{hh})\\n\\n    If :attr:`nonlinearity` is `'relu'`, then ReLU is used in place of tanh.\\n\\n    Args:\\n        input_size: The number of expected features in the input `x`\\n        hidden_size: The number of features in the hidden state `h`\\n        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.\\n            Default: ``True``\\n        nonlinearity: The non-linearity to use. Can be either ``'tanh'`` or ``'relu'``. Default: ``'tanh'``\\n\\n    Inputs: input, hidden\\n        - **input**: tensor containing input features\\n        - **hidden**: tensor containing the initial hidden state\\n          Defaults to zero if not provided.\\n\\n    Outputs: h'\\n        - **h'** of shape `(batch, hidden_size)`: tensor containing the next hidden state\\n          for each element in the batch\\n\\n    Shape:\\n        - input: :math:`(N, H_{in})` or :math:`(H_{in})` tensor containing input features where\\n          :math:`H_{in}` = `input_size`.\\n        - hidden: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the initial hidden\\n          state where :math:`H_{out}` = `hidden_size`. Defaults to zero if not provided.\\n        - output: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the next hidden state.\\n\\n    Attributes:\\n        weight_ih: the learnable input-hidden weights, of shape\\n            `(hidden_size, input_size)`\\n        weight_hh: the learnable hidden-hidden weights, of shape\\n            `(hidden_size, hidden_size)`\\n        bias_ih: the learnable input-hidden bias, of shape `(hidden_size)`\\n        bias_hh: the learnable hidden-hidden bias, of shape `(hidden_size)`\\n\\n    .. note::\\n        All the weights and biases are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`\\n        where :math:`k = \\\\frac{1}{\\\\text{hidden\\\\_size}}`\\n\\n    Examples::\\n\\n        >>> rnn = nn.RNNCell(10, 20)\\n        >>> input = torch.randn(6, 3, 10)\\n        >>> hx = torch.randn(3, 20)\\n        >>> output = []\\n        >>> for i in range(6):\\n        ...     hx = rnn(input[i], hx)\\n        ...     output.append(hx)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"input_size\\\", \\\"hidden_size\\\", \\\"bias\\\", \\\"nonlinearity\\\"]\\n    nonlinearity: str\\n\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        bias: bool = True,\\n        nonlinearity: str = \\\"tanh\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(input_size, hidden_size, bias, num_chunks=1, **factory_kwargs)\\n        self.nonlinearity = nonlinearity\\n\\n    def forward(self, input: Tensor, hx: Optional[Tensor] = None) -> Tensor:\\n        if input.dim() not in (1, 2):\\n            raise ValueError(\\n                f\\\"RNNCell: Expected input to be 1D or 2D, got {input.dim()}D instead\\\"\\n            )\\n        if hx is not None and hx.dim() not in (1, 2):\\n            raise ValueError(\\n                f\\\"RNNCell: Expected hidden to be 1D or 2D, got {hx.dim()}D instead\\\"\\n            )\\n        is_batched = input.dim() == 2\\n        if not is_batched:\\n            input = input.unsqueeze(0)\\n\\n        if hx is None:\\n            hx = torch.zeros(\\n                input.size(0), self.hidden_size, dtype=input.dtype, device=input.device\\n            )\\n        else:\\n            hx = hx.unsqueeze(0) if not is_batched else hx\\n\\n        if self.nonlinearity == \\\"tanh\\\":\\n            ret = _VF.rnn_tanh_cell(\\n                input,\\n                hx,\\n                self.weight_ih,\\n                self.weight_hh,\\n                self.bias_ih,\\n                self.bias_hh,\\n            )\\n        elif self.nonlinearity == \\\"relu\\\":\\n            ret = _VF.rnn_relu_cell(\\n                input,\\n                hx,\\n                self.weight_ih,\\n                self.weight_hh,\\n                self.bias_ih,\\n                self.bias_hh,\\n            )\\n        else:\\n            ret = input  # TODO: remove when jit supports exception flow\\n            raise RuntimeError(f\\\"Unknown nonlinearity: {self.nonlinearity}\\\")\\n\\n        if not is_batched:\\n            ret = ret.squeeze(0)\\n\\n        return ret\\n\\n\\nclass LSTMCell(RNNCellBase):\\n    r\\\"\\\"\\\"A long short-term memory (LSTM) cell.\\n\\n    .. math::\\n\\n        \\\\begin{array}{ll}\\n        i = \\\\sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\\\\\\\\\n        f = \\\\sigma(W_{if} x + b_{if} + W_{hf} h + b_{hf}) \\\\\\\\\\n        g = \\\\tanh(W_{ig} x + b_{ig} + W_{hg} h + b_{hg}) \\\\\\\\\\n        o = \\\\sigma(W_{io} x + b_{io} + W_{ho} h + b_{ho}) \\\\\\\\\\n        c' = f \\\\odot c + i \\\\odot g \\\\\\\\\\n        h' = o \\\\odot \\\\tanh(c') \\\\\\\\\\n        \\\\end{array}\\n\\n    where :math:`\\\\sigma` is the sigmoid function, and :math:`\\\\odot` is the Hadamard product.\\n\\n    Args:\\n        input_size: The number of expected features in the input `x`\\n        hidden_size: The number of features in the hidden state `h`\\n        bias: If ``False``, then the layer does not use bias weights `b_ih` and\\n            `b_hh`. Default: ``True``\\n\\n    Inputs: input, (h_0, c_0)\\n        - **input** of shape `(batch, input_size)` or `(input_size)`: tensor containing input features\\n        - **h_0** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the initial hidden state\\n        - **c_0** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the initial cell state\\n\\n          If `(h_0, c_0)` is not provided, both **h_0** and **c_0** default to zero.\\n\\n    Outputs: (h_1, c_1)\\n        - **h_1** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the next hidden state\\n        - **c_1** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the next cell state\\n\\n    Attributes:\\n        weight_ih: the learnable input-hidden weights, of shape\\n            `(4*hidden_size, input_size)`\\n        weight_hh: the learnable hidden-hidden weights, of shape\\n            `(4*hidden_size, hidden_size)`\\n        bias_ih: the learnable input-hidden bias, of shape `(4*hidden_size)`\\n        bias_hh: the learnable hidden-hidden bias, of shape `(4*hidden_size)`\\n\\n    .. note::\\n        All the weights and biases are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`\\n        where :math:`k = \\\\frac{1}{\\\\text{hidden\\\\_size}}`\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    Examples::\\n\\n        >>> rnn = nn.LSTMCell(10, 20)  # (input_size, hidden_size)\\n        >>> input = torch.randn(2, 3, 10)  # (time_steps, batch, input_size)\\n        >>> hx = torch.randn(3, 20)  # (batch, hidden_size)\\n        >>> cx = torch.randn(3, 20)\\n        >>> output = []\\n        >>> for i in range(input.size()[0]):\\n        ...     hx, cx = rnn(input[i], (hx, cx))\\n        ...     output.append(hx)\\n        >>> output = torch.stack(output, dim=0)\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(input_size, hidden_size, bias, num_chunks=4, **factory_kwargs)\\n\\n    def forward(\\n        self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = None\\n    ) -> Tuple[Tensor, Tensor]:\\n        if input.dim() not in (1, 2):\\n            raise ValueError(\\n                f\\\"LSTMCell: Expected input to be 1D or 2D, got {input.dim()}D instead\\\"\\n            )\\n        if hx is not None:\\n            for idx, value in enumerate(hx):\\n                if value.dim() not in (1, 2):\\n                    raise ValueError(\\n                        f\\\"LSTMCell: Expected hx[{idx}] to be 1D or 2D, got {value.dim()}D instead\\\"\\n                    )\\n        is_batched = input.dim() == 2\\n        if not is_batched:\\n            input = input.unsqueeze(0)\\n\\n        if hx is None:\\n            zeros = torch.zeros(\\n                input.size(0), self.hidden_size, dtype=input.dtype, device=input.device\\n            )\\n            hx = (zeros, zeros)\\n        else:\\n            hx = (hx[0].unsqueeze(0), hx[1].unsqueeze(0)) if not is_batched else hx\\n\\n        ret = _VF.lstm_cell(\\n            input,\\n            hx,\\n            self.weight_ih,\\n            self.weight_hh,\\n            self.bias_ih,\\n            self.bias_hh,\\n        )\\n\\n        if not is_batched:\\n            ret = (ret[0].squeeze(0), ret[1].squeeze(0))\\n        return ret\\n\\n\\nclass GRUCell(RNNCellBase):\\n    r\\\"\\\"\\\"A gated recurrent unit (GRU) cell.\\n\\n    .. math::\\n\\n        \\\\begin{array}{ll}\\n        r = \\\\sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\\\\\\\\\n        z = \\\\sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\\\\\\\\\n        n = \\\\tanh(W_{in} x + b_{in} + r \\\\odot (W_{hn} h + b_{hn})) \\\\\\\\\\n        h' = (1 - z) \\\\odot n + z \\\\odot h\\n        \\\\end{array}\\n\\n    where :math:`\\\\sigma` is the sigmoid function, and :math:`\\\\odot` is the Hadamard product.\\n\\n    Args:\\n        input_size: The number of expected features in the input `x`\\n        hidden_size: The number of features in the hidden state `h`\\n        bias: If ``False``, then the layer does not use bias weights `b_ih` and\\n            `b_hh`. Default: ``True``\\n\\n    Inputs: input, hidden\\n        - **input** : tensor containing input features\\n        - **hidden** : tensor containing the initial hidden\\n          state for each element in the batch.\\n          Defaults to zero if not provided.\\n\\n    Outputs: h'\\n        - **h'** : tensor containing the next hidden state\\n          for each element in the batch\\n\\n    Shape:\\n        - input: :math:`(N, H_{in})` or :math:`(H_{in})` tensor containing input features where\\n          :math:`H_{in}` = `input_size`.\\n        - hidden: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the initial hidden\\n          state where :math:`H_{out}` = `hidden_size`. Defaults to zero if not provided.\\n        - output: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the next hidden state.\\n\\n    Attributes:\\n        weight_ih: the learnable input-hidden weights, of shape\\n            `(3*hidden_size, input_size)`\\n        weight_hh: the learnable hidden-hidden weights, of shape\\n            `(3*hidden_size, hidden_size)`\\n        bias_ih: the learnable input-hidden bias, of shape `(3*hidden_size)`\\n        bias_hh: the learnable hidden-hidden bias, of shape `(3*hidden_size)`\\n\\n    .. note::\\n        All the weights and biases are initialized from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})`\\n        where :math:`k = \\\\frac{1}{\\\\text{hidden\\\\_size}}`\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    Examples::\\n\\n        >>> rnn = nn.GRUCell(10, 20)\\n        >>> input = torch.randn(6, 3, 10)\\n        >>> hx = torch.randn(3, 20)\\n        >>> output = []\\n        >>> for i in range(6):\\n        ...     hx = rnn(input[i], hx)\\n        ...     output.append(hx)\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        input_size: int,\\n        hidden_size: int,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(input_size, hidden_size, bias, num_chunks=3, **factory_kwargs)\\n\\n    def forward(self, input: Tensor, hx: Optional[Tensor] = None) -> Tensor:\\n        if input.dim() not in (1, 2):\\n            raise ValueError(\\n                f\\\"GRUCell: Expected input to be 1D or 2D, got {input.dim()}D instead\\\"\\n            )\\n        if hx is not None and hx.dim() not in (1, 2):\\n            raise ValueError(\\n                f\\\"GRUCell: Expected hidden to be 1D or 2D, got {hx.dim()}D instead\\\"\\n            )\\n        is_batched = input.dim() == 2\\n        if not is_batched:\\n            input = input.unsqueeze(0)\\n\\n        if hx is None:\\n            hx = torch.zeros(\\n                input.size(0), self.hidden_size, dtype=input.dtype, device=input.device\\n            )\\n        else:\\n            hx = hx.unsqueeze(0) if not is_batched else hx\\n\\n        ret = _VF.gru_cell(\\n            input,\\n            hx,\\n            self.weight_ih,\\n            self.weight_hh,\\n            self.bias_ih,\\n            self.bias_hh,\\n        )\\n\\n        if not is_batched:\\n            ret = ret.squeeze(0)\\n\\n        return ret\\n\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"PairwiseDistance\\\", \\\"CosineSimilarity\\\"]\\n\\n\\nclass PairwiseDistance(Module):\\n    r\\\"\\\"\\\"\\n    Computes the pairwise distance between input vectors, or between columns of input matrices.\\n\\n    Distances are computed using ``p``-norm, with constant ``eps`` added to avoid division by zero\\n    if ``p`` is negative, i.e.:\\n\\n    .. math ::\\n        \\\\mathrm{dist}\\\\left(x, y\\\\right) = \\\\left\\\\Vert x-y + \\\\epsilon e \\\\right\\\\Vert_p,\\n\\n    where :math:`e` is the vector of ones and the ``p``-norm is given by.\\n\\n    .. math ::\\n        \\\\Vert x \\\\Vert _p = \\\\left( \\\\sum_{i=1}^n  \\\\vert x_i \\\\vert ^ p \\\\right) ^ {1/p}.\\n\\n    Args:\\n        p (real, optional): the norm degree. Can be negative. Default: 2\\n        eps (float, optional): Small value to avoid division by zero.\\n            Default: 1e-6\\n        keepdim (bool, optional): Determines whether or not to keep the vector dimension.\\n            Default: False\\n    Shape:\\n        - Input1: :math:`(N, D)` or :math:`(D)` where `N = batch dimension` and `D = vector dimension`\\n        - Input2: :math:`(N, D)` or :math:`(D)`, same shape as the Input1\\n        - Output: :math:`(N)` or :math:`()` based on input dimension.\\n          If :attr:`keepdim` is ``True``, then :math:`(N, 1)` or :math:`(1)` based on input dimension.\\n\\n    Examples::\\n        >>> pdist = nn.PairwiseDistance(p=2)\\n        >>> input1 = torch.randn(100, 128)\\n        >>> input2 = torch.randn(100, 128)\\n        >>> output = pdist(input1, input2)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"norm\\\", \\\"eps\\\", \\\"keepdim\\\"]\\n    norm: float\\n    eps: float\\n    keepdim: bool\\n\\n    def __init__(\\n        self, p: float = 2.0, eps: float = 1e-6, keepdim: bool = False\\n    ) -> None:\\n        super().__init__()\\n        self.norm = p\\n        self.eps = eps\\n        self.keepdim = keepdim\\n\\n    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:\\n        return F.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim)\\n\\n\\nclass CosineSimilarity(Module):\\n    r\\\"\\\"\\\"Returns cosine similarity between :math:`x_1` and :math:`x_2`, computed along `dim`.\\n\\n    .. math ::\\n        \\\\text{similarity} = \\\\dfrac{x_1 \\\\cdot x_2}{\\\\max(\\\\Vert x_1 \\\\Vert _2 \\\\cdot \\\\Vert x_2 \\\\Vert _2, \\\\epsilon)}.\\n\\n    Args:\\n        dim (int, optional): Dimension where cosine similarity is computed. Default: 1\\n        eps (float, optional): Small value to avoid division by zero.\\n            Default: 1e-8\\n    Shape:\\n        - Input1: :math:`(\\\\ast_1, D, \\\\ast_2)` where D is at position `dim`\\n        - Input2: :math:`(\\\\ast_1, D, \\\\ast_2)`, same number of dimensions as x1, matching x1 size at dimension `dim`,\\n              and broadcastable with x1 at other dimensions.\\n        - Output: :math:`(\\\\ast_1, \\\\ast_2)`\\n    Examples::\\n        >>> input1 = torch.randn(100, 128)\\n        >>> input2 = torch.randn(100, 128)\\n        >>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)\\n        >>> output = cos(input1, input2)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"dim\\\", \\\"eps\\\"]\\n    dim: int\\n    eps: float\\n\\n    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:\\n        super().__init__()\\n        self.dim = dim\\n        self.eps = eps\\n\\n    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:\\n        return F.cosine_similarity(x1, x2, self.dim, self.eps)\\n\\n\\n# mypy: allow-untyped-defs\\nimport torch\\nimport torch.distributed as dist\\nfrom torch.autograd.function import Function\\n\\n\\nclass SyncBatchNorm(Function):\\n    @staticmethod\\n    def forward(\\n        self,\\n        input,\\n        weight,\\n        bias,\\n        running_mean,\\n        running_var,\\n        eps,\\n        momentum,\\n        process_group,\\n        world_size,\\n    ):\\n        if not (\\n            input.is_contiguous(memory_format=torch.channels_last)\\n            or input.is_contiguous(memory_format=torch.channels_last_3d)\\n        ):\\n            input = input.contiguous()\\n        if weight is not None:\\n            weight = weight.contiguous()\\n\\n        size = int(input.numel() // input.size(1))\\n        if size == 1 and world_size < 2:\\n            raise ValueError(\\n                f\\\"Expected more than 1 value per channel when training, got input size {size}\\\"\\n            )\\n\\n        num_channels = input.shape[1]\\n        if input.numel() > 0:\\n            # calculate mean/invstd for input.\\n            mean, invstd = torch.batch_norm_stats(input, eps)\\n\\n            count = torch.full(\\n                (1,),\\n                input.numel() // input.size(1),\\n                dtype=mean.dtype,\\n                device=mean.device,\\n            )\\n\\n            # C, C, 1 -> (2C + 1)\\n            combined = torch.cat([mean, invstd, count], dim=0)\\n        else:\\n            # for empty input, set stats and the count to zero. The stats with\\n            # zero count will be filtered out later when computing global mean\\n            # & invstd, but they still needs to participate the all_gather\\n            # collective communication to unblock other peer processes.\\n            combined = torch.zeros(\\n                2 * num_channels + 1, dtype=input.dtype, device=input.device\\n            )\\n\\n        # Use allgather instead of allreduce because count could be different across\\n        # ranks, simple all reduce op can not give correct results.\\n        # batch_norm_gather_stats_with_counts calculates global mean & invstd based on\\n        # all gathered mean, invstd and count.\\n        # for nccl backend, use the optimized version of all gather.\\n        # The Gloo backend does not support `all_gather_into_tensor`.\\n        if process_group._get_backend_name() != \\\"gloo\\\":\\n            # world_size * (2C + 1)\\n            combined_size = combined.numel()\\n            combined_flat = torch.empty(\\n                1,\\n                combined_size * world_size,\\n                dtype=combined.dtype,\\n                device=combined.device,\\n            )\\n            dist.all_gather_into_tensor(\\n                combined_flat, combined, process_group, async_op=False\\n            )\\n            combined = torch.reshape(combined_flat, (world_size, combined_size))\\n            # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1\\n            mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)\\n        else:\\n            # world_size * (2C + 1)\\n            combined_list = [torch.empty_like(combined) for _ in range(world_size)]\\n            dist.all_gather(combined_list, combined, process_group, async_op=False)\\n            combined = torch.stack(combined_list, dim=0)\\n            # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1\\n            mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)\\n\\n        if not (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()):\\n            # The lines below force a synchronization between CUDA and CPU, because\\n            # the shape of the result count_all depends on the values in mask tensor.\\n            # Such synchronizations break CUDA Graph capturing.\\n            # See https://github.com/pytorch/pytorch/issues/78549\\n            # FIXME: https://github.com/pytorch/pytorch/issues/78656 describes\\n            # a better longer-term solution.\\n\\n            # remove stats from empty inputs\\n            mask = count_all.squeeze(-1) >= 1\\n            count_all = count_all[mask]\\n            mean_all = mean_all[mask]\\n            invstd_all = invstd_all[mask]\\n\\n        # calculate global mean & invstd\\n        counts = count_all.view(-1)\\n        if running_mean is not None and counts.dtype != running_mean.dtype:\\n            counts = counts.to(running_mean.dtype)\\n        mean, invstd = torch.batch_norm_gather_stats_with_counts(\\n            input,\\n            mean_all,\\n            invstd_all,\\n            running_mean,\\n            running_var,\\n            momentum,\\n            eps,\\n            counts,\\n        )\\n\\n        self.save_for_backward(input, weight, mean, invstd, count_all.to(torch.int32))\\n        self.process_group = process_group\\n\\n        # apply element-wise normalization\\n        if input.numel() > 0:\\n            return torch.batch_norm_elemt(input, weight, bias, mean, invstd, eps)\\n        else:\\n            return torch.empty_like(input)\\n\\n    @staticmethod\\n    def backward(self, grad_output):\\n        if not (\\n            grad_output.is_contiguous(memory_format=torch.channels_last)\\n            or grad_output.is_contiguous(memory_format=torch.channels_last_3d)\\n        ):\\n            grad_output = grad_output.contiguous()\\n        saved_input, weight, mean, invstd, count_tensor = self.saved_tensors\\n        grad_input = grad_weight = grad_bias = None\\n        process_group = self.process_group\\n\\n        if saved_input.numel() > 0:\\n            # calculate local stats as well as grad_weight / grad_bias\\n            (\\n                sum_dy,\\n                sum_dy_xmu,\\n                grad_weight,\\n                grad_bias,\\n            ) = torch.batch_norm_backward_reduce(\\n                grad_output,\\n                saved_input,\\n                mean,\\n                invstd,\\n                weight,\\n                self.needs_input_grad[0],\\n                self.needs_input_grad[1],\\n                self.needs_input_grad[2],\\n            )\\n\\n            if self.needs_input_grad[0]:\\n                # synchronizing stats used to calculate input gradient.\\n                num_channels = sum_dy.shape[0]\\n                combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)\\n                torch.distributed.all_reduce(\\n                    combined,\\n                    torch.distributed.ReduceOp.SUM,\\n                    process_group,\\n                    async_op=False,\\n                )\\n                sum_dy, sum_dy_xmu = torch.split(combined, num_channels)\\n\\n                # backward pass for gradient calculation\\n                if weight is not None and weight.dtype != mean.dtype:\\n                    weight = weight.to(mean.dtype)\\n                grad_input = torch.batch_norm_backward_elemt(\\n                    grad_output,\\n                    saved_input,\\n                    mean,\\n                    invstd,\\n                    weight,\\n                    sum_dy,\\n                    sum_dy_xmu,\\n                    count_tensor,\\n                )\\n            # synchronizing of grad_weight / grad_bias is not needed as distributed\\n            # training would handle all reduce.\\n            if weight is None or not self.needs_input_grad[1]:\\n                grad_weight = None\\n\\n            if weight is None or not self.needs_input_grad[2]:\\n                grad_bias = None\\n        else:\\n            # This process got an empty input tensor in the forward pass.\\n            # Although this process can directly set grad_input as an empty\\n            # tensor of zeros, it still needs to participate in the collective\\n            # communication to unblock its peers, as other peer processes might\\n            # have received non-empty inputs.\\n            num_channels = saved_input.shape[1]\\n            if self.needs_input_grad[0]:\\n                # launch all_reduce to unblock other peer processes\\n                combined = torch.zeros(\\n                    2 * num_channels, dtype=saved_input.dtype, device=saved_input.device\\n                )\\n                torch.distributed.all_reduce(\\n                    combined,\\n                    torch.distributed.ReduceOp.SUM,\\n                    process_group,\\n                    async_op=False,\\n                )\\n\\n            # Leave grad_input, grad_weight and grad_bias as None, which will be\\n            # interpreted by the autograd engine as Tensors full of zeros.\\n\\n        return grad_input, grad_weight, grad_bias, None, None, None, None, None, None\\n\\n\\nclass CrossMapLRN2d(Function):\\n    @staticmethod\\n    def forward(ctx, input, size, alpha=1e-4, beta=0.75, k=1):\\n        ctx.size = size\\n        ctx.alpha = alpha\\n        ctx.beta = beta\\n        ctx.k = k\\n        ctx.scale = None\\n\\n        if input.dim() != 4:\\n            raise ValueError(\\n                f\\\"CrossMapLRN2d: Expected input to be 4D, got {input.dim()}D instead.\\\"\\n            )\\n\\n        ctx.scale = ctx.scale or input.new()\\n        output = input.new()\\n\\n        batch_size = input.size(0)\\n        channels = input.size(1)\\n        input_height = input.size(2)\\n        input_width = input.size(3)\\n\\n        output.resize_as_(input)\\n        ctx.scale.resize_as_(input)\\n\\n        # use output storage as temporary buffer\\n        input_square = output\\n        torch.pow(input, 2, out=input_square)\\n\\n        pre_pad = int((ctx.size - 1) / 2 + 1)\\n        pre_pad_crop = min(pre_pad, channels)\\n\\n        scale_first = ctx.scale.select(1, 0)\\n        scale_first.zero_()\\n        # compute first feature map normalization\\n        for c in range(pre_pad_crop):\\n            scale_first.add_(input_square.select(1, c))\\n\\n        # reuse computations for next feature maps normalization\\n        # by adding the next feature map and removing the previous\\n        for c in range(1, channels):\\n            scale_previous = ctx.scale.select(1, c - 1)\\n            scale_current = ctx.scale.select(1, c)\\n            scale_current.copy_(scale_previous)\\n            if c < channels - pre_pad + 1:\\n                square_next = input_square.select(1, c + pre_pad - 1)\\n                scale_current.add_(square_next, alpha=1)\\n\\n            if c > pre_pad:\\n                square_previous = input_square.select(1, c - pre_pad)\\n                scale_current.add_(square_previous, alpha=-1)\\n\\n        ctx.scale.mul_(ctx.alpha / ctx.size).add_(ctx.k)\\n\\n        torch.pow(ctx.scale, -ctx.beta, out=output)\\n        output.mul_(input)\\n\\n        ctx.save_for_backward(input, output)\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        input, output = ctx.saved_tensors\\n        grad_input = grad_output.new()\\n\\n        batch_size = input.size(0)\\n        channels = input.size(1)\\n        input_height = input.size(2)\\n        input_width = input.size(3)\\n\\n        paddded_ratio = input.new(channels + ctx.size - 1, input_height, input_width)\\n        accum_ratio = input.new(input_height, input_width)\\n\\n        cache_ratio_value = 2 * ctx.alpha * ctx.beta / ctx.size\\n        inversePrePad = int(ctx.size - (ctx.size - 1) / 2)\\n\\n        grad_input.resize_as_(input)\\n        torch.pow(ctx.scale, -ctx.beta, out=grad_input).mul_(grad_output)\\n\\n        paddded_ratio.zero_()\\n        padded_ratio_center = paddded_ratio.narrow(0, inversePrePad, channels)\\n        for n in range(batch_size):\\n            torch.mul(grad_output[n], output[n], out=padded_ratio_center)\\n            padded_ratio_center.div_(ctx.scale[n])\\n            torch.sum(\\n                paddded_ratio.narrow(0, 0, ctx.size - 1),\\n                0,\\n                keepdim=False,\\n                out=accum_ratio,\\n            )\\n            for c in range(channels):\\n                accum_ratio.add_(paddded_ratio[c + ctx.size - 1])\\n                grad_input[n][c].addcmul_(\\n                    input[n][c], accum_ratio, value=-cache_ratio_value\\n                )\\n                accum_ratio.add_(paddded_ratio[c], alpha=-1)\\n\\n        return grad_input, None, None, None, None\\n\\n\\nclass BackwardHookFunction(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, *args):\\n        ctx.mark_non_differentiable(*[arg for arg in args if not arg.requires_grad])\\n        return args\\n\\n    @staticmethod\\n    def backward(ctx, *args):\\n        return args\\n\\n\\nfrom typing import List, Optional\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.common_types import (\\n    _ratio_2_t,\\n    _ratio_3_t,\\n    _size_1_t,\\n    _size_2_opt_t,\\n    _size_2_t,\\n    _size_3_opt_t,\\n    _size_3_t,\\n    _size_any_opt_t,\\n    _size_any_t,\\n)\\n\\nfrom .module import Module\\nfrom .utils import _pair, _single, _triple\\n\\n\\n__all__ = [\\n    \\\"MaxPool1d\\\",\\n    \\\"MaxPool2d\\\",\\n    \\\"MaxPool3d\\\",\\n    \\\"MaxUnpool1d\\\",\\n    \\\"MaxUnpool2d\\\",\\n    \\\"MaxUnpool3d\\\",\\n    \\\"AvgPool1d\\\",\\n    \\\"AvgPool2d\\\",\\n    \\\"AvgPool3d\\\",\\n    \\\"FractionalMaxPool2d\\\",\\n    \\\"FractionalMaxPool3d\\\",\\n    \\\"LPPool1d\\\",\\n    \\\"LPPool2d\\\",\\n    \\\"LPPool3d\\\",\\n    \\\"AdaptiveMaxPool1d\\\",\\n    \\\"AdaptiveMaxPool2d\\\",\\n    \\\"AdaptiveMaxPool3d\\\",\\n    \\\"AdaptiveAvgPool1d\\\",\\n    \\\"AdaptiveAvgPool2d\\\",\\n    \\\"AdaptiveAvgPool3d\\\",\\n]\\n\\n\\nclass _MaxPoolNd(Module):\\n    __constants__ = [\\n        \\\"kernel_size\\\",\\n        \\\"stride\\\",\\n        \\\"padding\\\",\\n        \\\"dilation\\\",\\n        \\\"return_indices\\\",\\n        \\\"ceil_mode\\\",\\n    ]\\n    return_indices: bool\\n    ceil_mode: bool\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_any_t,\\n        stride: Optional[_size_any_t] = None,\\n        padding: _size_any_t = 0,\\n        dilation: _size_any_t = 1,\\n        return_indices: bool = False,\\n        ceil_mode: bool = False,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = kernel_size\\n        self.stride = stride if (stride is not None) else kernel_size\\n        self.padding = padding\\n        self.dilation = dilation\\n        self.return_indices = return_indices\\n        self.ceil_mode = ceil_mode\\n\\n    def extra_repr(self) -> str:\\n        return (\\n            \\\"kernel_size={kernel_size}, stride={stride}, padding={padding}\\\"\\n            \\\", dilation={dilation}, ceil_mode={ceil_mode}\\\".format(**self.__dict__)\\n        )\\n\\n\\nclass MaxPool1d(_MaxPoolNd):\\n    r\\\"\\\"\\\"Applies a 1D max pooling over an input signal composed of several input planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C, L)`\\n    and output :math:`(N, C, L_{out})` can be precisely described as:\\n\\n    .. math::\\n        out(N_i, C_j, k) = \\\\max_{m=0, \\\\ldots, \\\\text{kernel\\\\_size} - 1}\\n                input(N_i, C_j, stride \\\\times k + m)\\n\\n    If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides\\n    for :attr:`padding` number of points. :attr:`dilation` is the stride between the elements within the\\n    sliding window. This `link`_ has a nice visualization of the pooling parameters.\\n\\n    Note:\\n        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding\\n        or the input. Sliding windows that would start in the right padded region are ignored.\\n\\n    Args:\\n        kernel_size: The size of the sliding window, must be > 0.\\n        stride: The stride of the sliding window, must be > 0. Default value is :attr:`kernel_size`.\\n        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.\\n        dilation: The stride between elements within a sliding window, must be > 0.\\n        return_indices: If ``True``, will return the argmax along with the max values.\\n                        Useful for :class:`torch.nn.MaxUnpool1d` later\\n        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This\\n                   ensures that every element in the input tensor is covered by a sliding window.\\n\\n    Shape:\\n        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.\\n        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where\\n\\n          .. math::\\n              L_{out} = \\\\left\\\\lfloor \\\\frac{L_{in} + 2 \\\\times \\\\text{padding} - \\\\text{dilation}\\n                    \\\\times (\\\\text{kernel\\\\_size} - 1) - 1}{\\\\text{stride}} + 1\\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> # pool of size=3, stride=2\\n        >>> m = nn.MaxPool1d(3, stride=2)\\n        >>> input = torch.randn(20, 16, 50)\\n        >>> output = m(input)\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_1_t\\n    stride: _size_1_t\\n    padding: _size_1_t\\n    dilation: _size_1_t\\n\\n    def forward(self, input: Tensor):\\n        return F.max_pool1d(\\n            input,\\n            self.kernel_size,\\n            self.stride,\\n            self.padding,\\n            self.dilation,\\n            ceil_mode=self.ceil_mode,\\n            return_indices=self.return_indices,\\n        )\\n\\n\\nclass MaxPool2d(_MaxPoolNd):\\n    r\\\"\\\"\\\"Applies a 2D max pooling over an input signal composed of several input planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C, H, W)`,\\n    output :math:`(N, C, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kH, kW)`\\n    can be precisely described as:\\n\\n    .. math::\\n        \\\\begin{aligned}\\n            out(N_i, C_j, h, w) ={} & \\\\max_{m=0, \\\\ldots, kH-1} \\\\max_{n=0, \\\\ldots, kW-1} \\\\\\\\\\n                                    & \\\\text{input}(N_i, C_j, \\\\text{stride[0]} \\\\times h + m,\\n                                                   \\\\text{stride[1]} \\\\times w + n)\\n        \\\\end{aligned}\\n\\n    If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides\\n    for :attr:`padding` number of points. :attr:`dilation` controls the spacing between the kernel points.\\n    It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.\\n\\n    Note:\\n        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding\\n        or the input. Sliding windows that would start in the right padded region are ignored.\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the height and width dimension\\n        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,\\n          and the second `int` for the width dimension\\n\\n    Args:\\n        kernel_size: the size of the window to take a max over\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        padding: Implicit negative infinity padding to be added on both sides\\n        dilation: a parameter that controls the stride of elements in the window\\n        return_indices: if ``True``, will return the max indices along with the outputs.\\n                        Useful for :class:`torch.nn.MaxUnpool2d` later\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in} + 2 * \\\\text{padding[0]} - \\\\text{dilation[0]}\\n                    \\\\times (\\\\text{kernel\\\\_size[0]} - 1) - 1}{\\\\text{stride[0]}} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in} + 2 * \\\\text{padding[1]} - \\\\text{dilation[1]}\\n                    \\\\times (\\\\text{kernel\\\\_size[1]} - 1) - 1}{\\\\text{stride[1]}} + 1\\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> # pool of square window of size=3, stride=2\\n        >>> m = nn.MaxPool2d(3, stride=2)\\n        >>> # pool of non-square window\\n        >>> m = nn.MaxPool2d((3, 2), stride=(2, 1))\\n        >>> input = torch.randn(20, 16, 50, 32)\\n        >>> output = m(input)\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_2_t\\n    stride: _size_2_t\\n    padding: _size_2_t\\n    dilation: _size_2_t\\n\\n    def forward(self, input: Tensor):\\n        return F.max_pool2d(\\n            input,\\n            self.kernel_size,\\n            self.stride,\\n            self.padding,\\n            self.dilation,\\n            ceil_mode=self.ceil_mode,\\n            return_indices=self.return_indices,\\n        )\\n\\n\\nclass MaxPool3d(_MaxPoolNd):\\n    r\\\"\\\"\\\"Applies a 3D max pooling over an input signal composed of several input planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C, D, H, W)`,\\n    output :math:`(N, C, D_{out}, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kD, kH, kW)`\\n    can be precisely described as:\\n\\n    .. math::\\n        \\\\begin{aligned}\\n            \\\\text{out}(N_i, C_j, d, h, w) ={} & \\\\max_{k=0, \\\\ldots, kD-1} \\\\max_{m=0, \\\\ldots, kH-1} \\\\max_{n=0, \\\\ldots, kW-1} \\\\\\\\\\n                                              & \\\\text{input}(N_i, C_j, \\\\text{stride[0]} \\\\times d + k,\\n                                                             \\\\text{stride[1]} \\\\times h + m, \\\\text{stride[2]} \\\\times w + n)\\n        \\\\end{aligned}\\n\\n    If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides\\n    for :attr:`padding` number of points. :attr:`dilation` controls the spacing between the kernel points.\\n    It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.\\n\\n    Note:\\n        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding\\n        or the input. Sliding windows that would start in the right padded region are ignored.\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the depth, height and width dimension\\n        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,\\n          the second `int` for the height dimension and the third `int` for the width dimension\\n\\n    Args:\\n        kernel_size: the size of the window to take a max over\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        padding: Implicit negative infinity padding to be added on all three sides\\n        dilation: a parameter that controls the stride of elements in the window\\n        return_indices: if ``True``, will return the max indices along with the outputs.\\n                        Useful for :class:`torch.nn.MaxUnpool3d` later\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              D_{out} = \\\\left\\\\lfloor\\\\frac{D_{in} + 2 \\\\times \\\\text{padding}[0] - \\\\text{dilation}[0] \\\\times\\n                (\\\\text{kernel\\\\_size}[0] - 1) - 1}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in} + 2 \\\\times \\\\text{padding}[1] - \\\\text{dilation}[1] \\\\times\\n                (\\\\text{kernel\\\\_size}[1] - 1) - 1}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in} + 2 \\\\times \\\\text{padding}[2] - \\\\text{dilation}[2] \\\\times\\n                (\\\\text{kernel\\\\_size}[2] - 1) - 1}{\\\\text{stride}[2]} + 1\\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> # pool of square window of size=3, stride=2\\n        >>> m = nn.MaxPool3d(3, stride=2)\\n        >>> # pool of non-square window\\n        >>> m = nn.MaxPool3d((3, 2, 2), stride=(2, 1, 2))\\n        >>> input = torch.randn(20, 16, 50, 44, 31)\\n        >>> output = m(input)\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n    \\\"\\\"\\\"  # noqa: E501\\n\\n    kernel_size: _size_3_t\\n    stride: _size_3_t\\n    padding: _size_3_t\\n    dilation: _size_3_t\\n\\n    def forward(self, input: Tensor):\\n        return F.max_pool3d(\\n            input,\\n            self.kernel_size,\\n            self.stride,\\n            self.padding,\\n            self.dilation,\\n            ceil_mode=self.ceil_mode,\\n            return_indices=self.return_indices,\\n        )\\n\\n\\nclass _MaxUnpoolNd(Module):\\n    def extra_repr(self) -> str:\\n        return f\\\"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}\\\"\\n\\n\\nclass MaxUnpool1d(_MaxUnpoolNd):\\n    r\\\"\\\"\\\"Computes a partial inverse of :class:`MaxPool1d`.\\n\\n    :class:`MaxPool1d` is not fully invertible, since the non-maximal values are lost.\\n\\n    :class:`MaxUnpool1d` takes in as input the output of :class:`MaxPool1d`\\n    including the indices of the maximal values and computes a partial inverse\\n    in which all non-maximal values are set to zero.\\n\\n    Note:\\n        This operation may behave nondeterministically when the input indices has repeat values.\\n        See https://github.com/pytorch/pytorch/issues/80827 and :doc:`/notes/randomness` for more information.\\n\\n    .. note:: :class:`MaxPool1d` can map several input sizes to the same output\\n              sizes. Hence, the inversion process can get ambiguous.\\n              To accommodate this, you can provide the needed output size\\n              as an additional argument :attr:`output_size` in the forward call.\\n              See the Inputs and Example below.\\n\\n    Args:\\n        kernel_size (int or tuple): Size of the max pooling window.\\n        stride (int or tuple): Stride of the max pooling window.\\n            It is set to :attr:`kernel_size` by default.\\n        padding (int or tuple): Padding that was added to the input\\n\\n    Inputs:\\n        - `input`: the input Tensor to invert\\n        - `indices`: the indices given out by :class:`~torch.nn.MaxPool1d`\\n        - `output_size` (optional): the targeted output size\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in})` or :math:`(C, H_{in})`.\\n        - Output: :math:`(N, C, H_{out})` or :math:`(C, H_{out})`, where\\n\\n          .. math::\\n              H_{out} = (H_{in} - 1) \\\\times \\\\text{stride}[0] - 2 \\\\times \\\\text{padding}[0] + \\\\text{kernel\\\\_size}[0]\\n\\n          or as given by :attr:`output_size` in the call operator\\n\\n    Example::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"do other tests modify the global state?\\\")\\n        >>> pool = nn.MaxPool1d(2, stride=2, return_indices=True)\\n        >>> unpool = nn.MaxUnpool1d(2, stride=2)\\n        >>> input = torch.tensor([[[1., 2, 3, 4, 5, 6, 7, 8]]])\\n        >>> output, indices = pool(input)\\n        >>> unpool(output, indices)\\n        tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.]]])\\n\\n        >>> # Example showcasing the use of output_size\\n        >>> input = torch.tensor([[[1., 2, 3, 4, 5, 6, 7, 8, 9]]])\\n        >>> output, indices = pool(input)\\n        >>> unpool(output, indices, output_size=input.size())\\n        tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.,  0.]]])\\n\\n        >>> unpool(output, indices)\\n        tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.]]])\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_1_t\\n    stride: _size_1_t\\n    padding: _size_1_t\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_1_t,\\n        stride: Optional[_size_1_t] = None,\\n        padding: _size_1_t = 0,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = _single(kernel_size)\\n        self.stride = _single(stride if (stride is not None) else kernel_size)\\n        self.padding = _single(padding)\\n\\n    def forward(\\n        self, input: Tensor, indices: Tensor, output_size: Optional[List[int]] = None\\n    ) -> Tensor:\\n        return F.max_unpool1d(\\n            input, indices, self.kernel_size, self.stride, self.padding, output_size\\n        )\\n\\n\\nclass MaxUnpool2d(_MaxUnpoolNd):\\n    r\\\"\\\"\\\"Computes a partial inverse of :class:`MaxPool2d`.\\n\\n    :class:`MaxPool2d` is not fully invertible, since the non-maximal values are lost.\\n\\n    :class:`MaxUnpool2d` takes in as input the output of :class:`MaxPool2d`\\n    including the indices of the maximal values and computes a partial inverse\\n    in which all non-maximal values are set to zero.\\n\\n    Note:\\n        This operation may behave nondeterministically when the input indices has repeat values.\\n        See https://github.com/pytorch/pytorch/issues/80827 and :doc:`/notes/randomness` for more information.\\n\\n    .. note:: :class:`MaxPool2d` can map several input sizes to the same output\\n              sizes. Hence, the inversion process can get ambiguous.\\n              To accommodate this, you can provide the needed output size\\n              as an additional argument :attr:`output_size` in the forward call.\\n              See the Inputs and Example below.\\n\\n    Args:\\n        kernel_size (int or tuple): Size of the max pooling window.\\n        stride (int or tuple): Stride of the max pooling window.\\n            It is set to :attr:`kernel_size` by default.\\n        padding (int or tuple): Padding that was added to the input\\n\\n    Inputs:\\n        - `input`: the input Tensor to invert\\n        - `indices`: the indices given out by :class:`~torch.nn.MaxPool2d`\\n        - `output_size` (optional): the targeted output size\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          .. math::\\n            H_{out} = (H_{in} - 1) \\\\times \\\\text{stride[0]} - 2 \\\\times \\\\text{padding[0]} + \\\\text{kernel\\\\_size[0]}\\n\\n          .. math::\\n            W_{out} = (W_{in} - 1) \\\\times \\\\text{stride[1]} - 2 \\\\times \\\\text{padding[1]} + \\\\text{kernel\\\\_size[1]}\\n\\n          or as given by :attr:`output_size` in the call operator\\n\\n    Example::\\n\\n        >>> pool = nn.MaxPool2d(2, stride=2, return_indices=True)\\n        >>> unpool = nn.MaxUnpool2d(2, stride=2)\\n        >>> input = torch.tensor([[[[ 1.,  2.,  3.,  4.],\\n                                    [ 5.,  6.,  7.,  8.],\\n                                    [ 9., 10., 11., 12.],\\n                                    [13., 14., 15., 16.]]]])\\n        >>> output, indices = pool(input)\\n        >>> unpool(output, indices)\\n        tensor([[[[  0.,   0.,   0.,   0.],\\n                  [  0.,   6.,   0.,   8.],\\n                  [  0.,   0.,   0.,   0.],\\n                  [  0.,  14.,   0.,  16.]]]])\\n        >>> # Now using output_size to resolve an ambiguous size for the inverse\\n        >>> input = torch.tensor([[[[ 1.,  2.,  3.,  4.,  5.],\\n                                    [ 6.,  7.,  8.,  9., 10.],\\n                                    [11., 12., 13., 14., 15.],\\n                                    [16., 17., 18., 19., 20.]]]])\\n        >>> output, indices = pool(input)\\n        >>> # This call will not work without specifying output_size\\n        >>> unpool(output, indices, output_size=input.size())\\n        tensor([[[[ 0.,  0.,  0.,  0.,  0.],\\n                  [ 0.,  7.,  0.,  9.,  0.],\\n                  [ 0.,  0.,  0.,  0.,  0.],\\n                  [ 0., 17.,  0., 19.,  0.]]]])\\n\\n\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_2_t\\n    stride: _size_2_t\\n    padding: _size_2_t\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_2_t,\\n        stride: Optional[_size_2_t] = None,\\n        padding: _size_2_t = 0,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = _pair(kernel_size)\\n        self.stride = _pair(stride if (stride is not None) else kernel_size)\\n        self.padding = _pair(padding)\\n\\n    def forward(\\n        self, input: Tensor, indices: Tensor, output_size: Optional[List[int]] = None\\n    ) -> Tensor:\\n        return F.max_unpool2d(\\n            input, indices, self.kernel_size, self.stride, self.padding, output_size\\n        )\\n\\n\\nclass MaxUnpool3d(_MaxUnpoolNd):\\n    r\\\"\\\"\\\"Computes a partial inverse of :class:`MaxPool3d`.\\n\\n    :class:`MaxPool3d` is not fully invertible, since the non-maximal values are lost.\\n    :class:`MaxUnpool3d` takes in as input the output of :class:`MaxPool3d`\\n    including the indices of the maximal values and computes a partial inverse\\n    in which all non-maximal values are set to zero.\\n\\n    Note:\\n        This operation may behave nondeterministically when the input indices has repeat values.\\n        See https://github.com/pytorch/pytorch/issues/80827 and :doc:`/notes/randomness` for more information.\\n\\n    .. note:: :class:`MaxPool3d` can map several input sizes to the same output\\n              sizes. Hence, the inversion process can get ambiguous.\\n              To accommodate this, you can provide the needed output size\\n              as an additional argument :attr:`output_size` in the forward call.\\n              See the Inputs section below.\\n\\n    Args:\\n        kernel_size (int or tuple): Size of the max pooling window.\\n        stride (int or tuple): Stride of the max pooling window.\\n            It is set to :attr:`kernel_size` by default.\\n        padding (int or tuple): Padding that was added to the input\\n\\n    Inputs:\\n        - `input`: the input Tensor to invert\\n        - `indices`: the indices given out by :class:`~torch.nn.MaxPool3d`\\n        - `output_size` (optional): the targeted output size\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              D_{out} = (D_{in} - 1) \\\\times \\\\text{stride[0]} - 2 \\\\times \\\\text{padding[0]} + \\\\text{kernel\\\\_size[0]}\\n\\n          .. math::\\n              H_{out} = (H_{in} - 1) \\\\times \\\\text{stride[1]} - 2 \\\\times \\\\text{padding[1]} + \\\\text{kernel\\\\_size[1]}\\n\\n          .. math::\\n              W_{out} = (W_{in} - 1) \\\\times \\\\text{stride[2]} - 2 \\\\times \\\\text{padding[2]} + \\\\text{kernel\\\\_size[2]}\\n\\n          or as given by :attr:`output_size` in the call operator\\n\\n    Example::\\n\\n        >>> # pool of square window of size=3, stride=2\\n        >>> pool = nn.MaxPool3d(3, stride=2, return_indices=True)\\n        >>> unpool = nn.MaxUnpool3d(3, stride=2)\\n        >>> output, indices = pool(torch.randn(20, 16, 51, 33, 15))\\n        >>> unpooled_output = unpool(output, indices)\\n        >>> unpooled_output.size()\\n        torch.Size([20, 16, 51, 33, 15])\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_3_t\\n    stride: _size_3_t\\n    padding: _size_3_t\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_3_t,\\n        stride: Optional[_size_3_t] = None,\\n        padding: _size_3_t = 0,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = _triple(kernel_size)\\n        self.stride = _triple(stride if (stride is not None) else kernel_size)\\n        self.padding = _triple(padding)\\n\\n    def forward(\\n        self, input: Tensor, indices: Tensor, output_size: Optional[List[int]] = None\\n    ) -> Tensor:\\n        return F.max_unpool3d(\\n            input, indices, self.kernel_size, self.stride, self.padding, output_size\\n        )\\n\\n\\nclass _AvgPoolNd(Module):\\n    __constants__ = [\\n        \\\"kernel_size\\\",\\n        \\\"stride\\\",\\n        \\\"padding\\\",\\n        \\\"ceil_mode\\\",\\n        \\\"count_include_pad\\\",\\n    ]\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}\\\"\\n\\n\\nclass AvgPool1d(_AvgPoolNd):\\n    r\\\"\\\"\\\"Applies a 1D average pooling over an input signal composed of several input planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C, L)`,\\n    output :math:`(N, C, L_{out})` and :attr:`kernel_size` :math:`k`\\n    can be precisely described as:\\n\\n    .. math::\\n\\n        \\\\text{out}(N_i, C_j, l) = \\\\frac{1}{k} \\\\sum_{m=0}^{k-1}\\n                               \\\\text{input}(N_i, C_j, \\\\text{stride} \\\\times l + m)\\n\\n    If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides\\n    for :attr:`padding` number of points.\\n\\n    Note:\\n        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding\\n        or the input. Sliding windows that would start in the right padded region are ignored.\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding` can each be\\n    an ``int`` or a one-element tuple.\\n\\n    Args:\\n        kernel_size: the size of the window\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        padding: implicit zero padding to be added on both sides\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n        count_include_pad: when True, will include the zero-padding in the averaging calculation\\n\\n    Shape:\\n        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.\\n        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where\\n\\n          .. math::\\n              L_{out} = \\\\left\\\\lfloor \\\\frac{L_{in} +\\n              2 \\\\times \\\\text{padding} - \\\\text{kernel\\\\_size}}{\\\\text{stride}} + 1\\\\right\\\\rfloor\\n\\n          Per the note above, if ``ceil_mode`` is True and :math:`(L_{out} - 1) \\\\times \\\\text{stride} \\\\geq L_{in}\\n          + \\\\text{padding}`, we skip the last window as it would start in the right padded region, resulting in\\n          :math:`L_{out}` being reduced by one.\\n\\n    Examples::\\n\\n        >>> # pool with window of size=3, stride=2\\n        >>> m = nn.AvgPool1d(3, stride=2)\\n        >>> m(torch.tensor([[[1., 2, 3, 4, 5, 6, 7]]]))\\n        tensor([[[2., 4., 6.]]])\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_1_t\\n    stride: _size_1_t\\n    padding: _size_1_t\\n    ceil_mode: bool\\n    count_include_pad: bool\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_1_t,\\n        stride: _size_1_t = None,\\n        padding: _size_1_t = 0,\\n        ceil_mode: bool = False,\\n        count_include_pad: bool = True,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = _single(kernel_size)\\n        self.stride = _single(stride if stride is not None else kernel_size)\\n        self.padding = _single(padding)\\n        self.ceil_mode = ceil_mode\\n        self.count_include_pad = count_include_pad\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.avg_pool1d(\\n            input,\\n            self.kernel_size,\\n            self.stride,\\n            self.padding,\\n            self.ceil_mode,\\n            self.count_include_pad,\\n        )\\n\\n\\nclass AvgPool2d(_AvgPoolNd):\\n    r\\\"\\\"\\\"Applies a 2D average pooling over an input signal composed of several input planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C, H, W)`,\\n    output :math:`(N, C, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kH, kW)`\\n    can be precisely described as:\\n\\n    .. math::\\n\\n        out(N_i, C_j, h, w)  = \\\\frac{1}{kH * kW} \\\\sum_{m=0}^{kH-1} \\\\sum_{n=0}^{kW-1}\\n                               input(N_i, C_j, stride[0] \\\\times h + m, stride[1] \\\\times w + n)\\n\\n    If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides\\n    for :attr:`padding` number of points.\\n\\n    Note:\\n        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding\\n        or the input. Sliding windows that would start in the right padded region are ignored.\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the height and width dimension\\n        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,\\n          and the second `int` for the width dimension\\n\\n    Args:\\n        kernel_size: the size of the window\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        padding: implicit zero padding to be added on both sides\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n        count_include_pad: when True, will include the zero-padding in the averaging calculation\\n        divisor_override: if specified, it will be used as divisor, otherwise size of the pooling region will be used.\\n\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in}  + 2 \\\\times \\\\text{padding}[0] -\\n                \\\\text{kernel\\\\_size}[0]}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in}  + 2 \\\\times \\\\text{padding}[1] -\\n                \\\\text{kernel\\\\_size}[1]}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n          Per the note above, if ``ceil_mode`` is True and :math:`(H_{out} - 1)\\\\times \\\\text{stride}[0]\\\\geq H_{in}\\n          + \\\\text{padding}[0]`, we skip the last window as it would start in the bottom padded region,\\n          resulting in :math:`H_{out}` being reduced by one.\\n\\n          The same applies for :math:`W_{out}`.\\n\\n    Examples::\\n\\n        >>> # pool of square window of size=3, stride=2\\n        >>> m = nn.AvgPool2d(3, stride=2)\\n        >>> # pool of non-square window\\n        >>> m = nn.AvgPool2d((3, 2), stride=(2, 1))\\n        >>> input = torch.randn(20, 16, 50, 32)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\n        \\\"kernel_size\\\",\\n        \\\"stride\\\",\\n        \\\"padding\\\",\\n        \\\"ceil_mode\\\",\\n        \\\"count_include_pad\\\",\\n        \\\"divisor_override\\\",\\n    ]\\n\\n    kernel_size: _size_2_t\\n    stride: _size_2_t\\n    padding: _size_2_t\\n    ceil_mode: bool\\n    count_include_pad: bool\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_2_t,\\n        stride: Optional[_size_2_t] = None,\\n        padding: _size_2_t = 0,\\n        ceil_mode: bool = False,\\n        count_include_pad: bool = True,\\n        divisor_override: Optional[int] = None,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = kernel_size\\n        self.stride = stride if (stride is not None) else kernel_size\\n        self.padding = padding\\n        self.ceil_mode = ceil_mode\\n        self.count_include_pad = count_include_pad\\n        self.divisor_override = divisor_override\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.avg_pool2d(\\n            input,\\n            self.kernel_size,\\n            self.stride,\\n            self.padding,\\n            self.ceil_mode,\\n            self.count_include_pad,\\n            self.divisor_override,\\n        )\\n\\n\\nclass AvgPool3d(_AvgPoolNd):\\n    r\\\"\\\"\\\"Applies a 3D average pooling over an input signal composed of several input planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C, D, H, W)`,\\n    output :math:`(N, C, D_{out}, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kD, kH, kW)`\\n    can be precisely described as:\\n\\n    .. math::\\n        \\\\begin{aligned}\\n            \\\\text{out}(N_i, C_j, d, h, w) ={} & \\\\sum_{k=0}^{kD-1} \\\\sum_{m=0}^{kH-1} \\\\sum_{n=0}^{kW-1} \\\\\\\\\\n                                              & \\\\frac{\\\\text{input}(N_i, C_j, \\\\text{stride}[0] \\\\times d + k,\\n                                                      \\\\text{stride}[1] \\\\times h + m, \\\\text{stride}[2] \\\\times w + n)}\\n                                                     {kD \\\\times kH \\\\times kW}\\n        \\\\end{aligned}\\n\\n    If :attr:`padding` is non-zero, then the input is implicitly zero-padded on all three sides\\n    for :attr:`padding` number of points.\\n\\n    Note:\\n        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding\\n        or the input. Sliding windows that would start in the right padded region are ignored.\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the depth, height and width dimension\\n        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,\\n          the second `int` for the height dimension and the third `int` for the width dimension\\n\\n    Args:\\n        kernel_size: the size of the window\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        padding: implicit zero padding to be added on all three sides\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n        count_include_pad: when True, will include the zero-padding in the averaging calculation\\n        divisor_override: if specified, it will be used as divisor, otherwise :attr:`kernel_size` will be used\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or\\n          :math:`(C, D_{out}, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              D_{out} = \\\\left\\\\lfloor\\\\frac{D_{in} + 2 \\\\times \\\\text{padding}[0] -\\n                    \\\\text{kernel\\\\_size}[0]}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in} + 2 \\\\times \\\\text{padding}[1] -\\n                    \\\\text{kernel\\\\_size}[1]}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in} + 2 \\\\times \\\\text{padding}[2] -\\n                    \\\\text{kernel\\\\_size}[2]}{\\\\text{stride}[2]} + 1\\\\right\\\\rfloor\\n\\n          Per the note above, if ``ceil_mode`` is True and :math:`(D_{out} - 1)\\\\times \\\\text{stride}[0]\\\\geq D_{in}\\n          + \\\\text{padding}[0]`, we skip the last window as it would start in the padded region,\\n          resulting in :math:`D_{out}` being reduced by one.\\n\\n          The same applies for :math:`W_{out}` and :math:`H_{out}`.\\n\\n    Examples::\\n\\n        >>> # pool of square window of size=3, stride=2\\n        >>> m = nn.AvgPool3d(3, stride=2)\\n        >>> # pool of non-square window\\n        >>> m = nn.AvgPool3d((3, 2, 2), stride=(2, 1, 2))\\n        >>> input = torch.randn(20, 16, 50, 44, 31)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\n        \\\"kernel_size\\\",\\n        \\\"stride\\\",\\n        \\\"padding\\\",\\n        \\\"ceil_mode\\\",\\n        \\\"count_include_pad\\\",\\n        \\\"divisor_override\\\",\\n    ]\\n\\n    kernel_size: _size_3_t\\n    stride: _size_3_t\\n    padding: _size_3_t\\n    ceil_mode: bool\\n    count_include_pad: bool\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_3_t,\\n        stride: Optional[_size_3_t] = None,\\n        padding: _size_3_t = 0,\\n        ceil_mode: bool = False,\\n        count_include_pad: bool = True,\\n        divisor_override: Optional[int] = None,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = kernel_size\\n        self.stride = stride if (stride is not None) else kernel_size\\n        self.padding = padding\\n        self.ceil_mode = ceil_mode\\n        self.count_include_pad = count_include_pad\\n        self.divisor_override = divisor_override\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.avg_pool3d(\\n            input,\\n            self.kernel_size,\\n            self.stride,\\n            self.padding,\\n            self.ceil_mode,\\n            self.count_include_pad,\\n            self.divisor_override,\\n        )\\n\\n    def __setstate__(self, d):\\n        super().__setstate__(d)\\n        self.__dict__.setdefault(\\\"padding\\\", 0)\\n        self.__dict__.setdefault(\\\"ceil_mode\\\", False)\\n        self.__dict__.setdefault(\\\"count_include_pad\\\", True)\\n\\n\\nclass FractionalMaxPool2d(Module):\\n    r\\\"\\\"\\\"Applies a 2D fractional max pooling over an input signal composed of several input planes.\\n\\n    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham\\n\\n    The max-pooling operation is applied in :math:`kH \\\\times kW` regions by a stochastic\\n    step size determined by the target output size.\\n    The number of output features is equal to the number of input planes.\\n\\n    .. note:: Exactly one of ``output_size`` or ``output_ratio`` must be defined.\\n\\n    Args:\\n        kernel_size: the size of the window to take a max over.\\n                     Can be a single number k (for a square kernel of k x k) or a tuple `(kh, kw)`\\n        output_size: the target output size of the image of the form `oH x oW`.\\n                     Can be a tuple `(oH, oW)` or a single number oH for a square image `oH x oH`.\\n                     Note that we must have :math:`kH + oH - 1 <= H_{in}` and :math:`kW + oW - 1 <= W_{in}`\\n        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.\\n                      This has to be a number or tuple in the range (0, 1).\\n                      Note that we must have :math:`kH + (output\\\\_ratio\\\\_H * H_{in}) - 1 <= H_{in}`\\n                      and :math:`kW + (output\\\\_ratio\\\\_W * W_{in}) - 1 <= W_{in}`\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to :meth:`nn.MaxUnpool2d`. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n          :math:`(H_{out}, W_{out})=\\\\text{output\\\\_size}` or\\n          :math:`(H_{out}, W_{out})=\\\\text{output\\\\_ratio} \\\\times (H_{in}, W_{in})`.\\n\\n    Examples:\\n        >>> # pool of square window of size=3, and target output size 13x12\\n        >>> m = nn.FractionalMaxPool2d(3, output_size=(13, 12))\\n        >>> # pool of square window and target output size being half of input image size\\n        >>> m = nn.FractionalMaxPool2d(3, output_ratio=(0.5, 0.5))\\n        >>> input = torch.randn(20, 16, 50, 32)\\n        >>> output = m(input)\\n\\n    .. _Fractional MaxPooling:\\n        https://arxiv.org/abs/1412.6071\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"kernel_size\\\", \\\"return_indices\\\", \\\"output_size\\\", \\\"output_ratio\\\"]\\n\\n    kernel_size: _size_2_t\\n    return_indices: bool\\n    output_size: _size_2_t\\n    output_ratio: _ratio_2_t\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_2_t,\\n        output_size: Optional[_size_2_t] = None,\\n        output_ratio: Optional[_ratio_2_t] = None,\\n        return_indices: bool = False,\\n        _random_samples=None,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = _pair(kernel_size)\\n        self.return_indices = return_indices\\n        self.register_buffer(\\\"_random_samples\\\", _random_samples)\\n        self.output_size = _pair(output_size) if output_size is not None else None\\n        self.output_ratio = _pair(output_ratio) if output_ratio is not None else None\\n        if output_size is None and output_ratio is None:\\n            raise ValueError(\\n                \\\"FractionalMaxPool2d requires specifying either \\\"\\n                \\\"an output size, or a pooling ratio\\\"\\n            )\\n        if output_size is not None and output_ratio is not None:\\n            raise ValueError(\\n                \\\"only one of output_size and output_ratio may be specified\\\"\\n            )\\n        if self.output_ratio is not None:\\n            if not (0 < self.output_ratio[0] < 1 and 0 < self.output_ratio[1] < 1):\\n                raise ValueError(\\n                    f\\\"output_ratio must be between 0 and 1 (got {output_ratio})\\\"\\n                )\\n\\n    def forward(self, input: Tensor):\\n        return F.fractional_max_pool2d(\\n            input,\\n            self.kernel_size,\\n            self.output_size,\\n            self.output_ratio,\\n            self.return_indices,\\n            _random_samples=self._random_samples,\\n        )\\n\\n\\nclass FractionalMaxPool3d(Module):\\n    r\\\"\\\"\\\"Applies a 3D fractional max pooling over an input signal composed of several input planes.\\n\\n    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham\\n\\n    The max-pooling operation is applied in :math:`kT \\\\times kH \\\\times kW` regions by a stochastic\\n    step size determined by the target output size.\\n    The number of output features is equal to the number of input planes.\\n\\n    .. note:: Exactly one of ``output_size`` or ``output_ratio`` must be defined.\\n\\n    Args:\\n        kernel_size: the size of the window to take a max over.\\n                     Can be a single number k (for a square kernel of k x k x k) or a tuple `(kt x kh x kw)`\\n        output_size: the target output size of the image of the form `oT x oH x oW`.\\n                     Can be a tuple `(oT, oH, oW)` or a single number oH for a square image `oH x oH x oH`\\n        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.\\n                      This has to be a number or tuple in the range (0, 1)\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to :meth:`nn.MaxUnpool3d`. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, T_{in}, H_{in}, W_{in})` or :math:`(C, T_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, T_{out}, H_{out}, W_{out})` or :math:`(C, T_{out}, H_{out}, W_{out})`, where\\n          :math:`(T_{out}, H_{out}, W_{out})=\\\\text{output\\\\_size}` or\\n          :math:`(T_{out}, H_{out}, W_{out})=\\\\text{output\\\\_ratio} \\\\times (T_{in}, H_{in}, W_{in})`\\n\\n    Examples:\\n        >>> # pool of cubic window of size=3, and target output size 13x12x11\\n        >>> m = nn.FractionalMaxPool3d(3, output_size=(13, 12, 11))\\n        >>> # pool of cubic window and target output size being half of input size\\n        >>> m = nn.FractionalMaxPool3d(3, output_ratio=(0.5, 0.5, 0.5))\\n        >>> input = torch.randn(20, 16, 50, 32, 16)\\n        >>> output = m(input)\\n\\n    .. _Fractional MaxPooling:\\n        https://arxiv.org/abs/1412.6071\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"kernel_size\\\", \\\"return_indices\\\", \\\"output_size\\\", \\\"output_ratio\\\"]\\n    kernel_size: _size_3_t\\n    return_indices: bool\\n    output_size: _size_3_t\\n    output_ratio: _ratio_3_t\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_3_t,\\n        output_size: Optional[_size_3_t] = None,\\n        output_ratio: Optional[_ratio_3_t] = None,\\n        return_indices: bool = False,\\n        _random_samples=None,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = _triple(kernel_size)\\n        self.return_indices = return_indices\\n        self.register_buffer(\\\"_random_samples\\\", _random_samples)\\n        self.output_size = _triple(output_size) if output_size is not None else None\\n        self.output_ratio = _triple(output_ratio) if output_ratio is not None else None\\n        if output_size is None and output_ratio is None:\\n            raise ValueError(\\n                \\\"FractionalMaxPool3d requires specifying either \\\"\\n                \\\"an output size, or a pooling ratio\\\"\\n            )\\n        if output_size is not None and output_ratio is not None:\\n            raise ValueError(\\n                \\\"only one of output_size and output_ratio may be specified\\\"\\n            )\\n        if self.output_ratio is not None:\\n            if not (\\n                0 < self.output_ratio[0] < 1\\n                and 0 < self.output_ratio[1] < 1\\n                and 0 < self.output_ratio[2] < 1\\n            ):\\n                raise ValueError(\\n                    f\\\"output_ratio must be between 0 and 1 (got {output_ratio})\\\"\\n                )\\n\\n    def forward(self, input: Tensor):\\n        return F.fractional_max_pool3d(\\n            input,\\n            self.kernel_size,\\n            self.output_size,\\n            self.output_ratio,\\n            self.return_indices,\\n            _random_samples=self._random_samples,\\n        )\\n\\n\\nclass _LPPoolNd(Module):\\n    __constants__ = [\\\"norm_type\\\", \\\"kernel_size\\\", \\\"stride\\\", \\\"ceil_mode\\\"]\\n\\n    norm_type: float\\n    ceil_mode: bool\\n\\n    def __init__(\\n        self,\\n        norm_type: float,\\n        kernel_size: _size_any_t,\\n        stride: Optional[_size_any_t] = None,\\n        ceil_mode: bool = False,\\n    ) -> None:\\n        super().__init__()\\n        self.norm_type = norm_type\\n        self.kernel_size = kernel_size\\n        self.stride = stride\\n        self.ceil_mode = ceil_mode\\n\\n    def extra_repr(self) -> str:\\n        return (\\n            \\\"norm_type={norm_type}, kernel_size={kernel_size}, stride={stride}, \\\"\\n            \\\"ceil_mode={ceil_mode}\\\".format(**self.__dict__)\\n        )\\n\\n\\nclass LPPool1d(_LPPoolNd):\\n    r\\\"\\\"\\\"Applies a 1D power-average pooling over an input signal composed of several input planes.\\n\\n    On each window, the function computed is:\\n\\n    .. math::\\n        f(X) = \\\\sqrt[p]{\\\\sum_{x \\\\in X} x^{p}}\\n\\n    - At p = :math:`\\\\infty`, one gets Max Pooling\\n    - At p = 1, one gets Sum Pooling (which is proportional to Average Pooling)\\n\\n    .. note:: If the sum to the power of `p` is zero, the gradient of this function is\\n              not defined. This implementation will set the gradient to zero in this case.\\n\\n    Args:\\n        kernel_size: a single int, the size of the window\\n        stride: a single int, the stride of the window. Default value is :attr:`kernel_size`\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n\\n    Shape:\\n        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.\\n        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where\\n\\n          .. math::\\n              L_{out} = \\\\left\\\\lfloor\\\\frac{L_{in} - \\\\text{kernel\\\\_size}}{\\\\text{stride}} + 1\\\\right\\\\rfloor\\n\\n    Examples::\\n        >>> # power-2 pool of window of length 3, with stride 2.\\n        >>> m = nn.LPPool1d(2, 3, stride=2)\\n        >>> input = torch.randn(20, 16, 50)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_1_t\\n    stride: _size_1_t\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.lp_pool1d(\\n            input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode\\n        )\\n\\n\\nclass LPPool2d(_LPPoolNd):\\n    r\\\"\\\"\\\"Applies a 2D power-average pooling over an input signal composed of several input planes.\\n\\n    On each window, the function computed is:\\n\\n    .. math::\\n        f(X) = \\\\sqrt[p]{\\\\sum_{x \\\\in X} x^{p}}\\n\\n    - At p = :math:`\\\\infty`, one gets Max Pooling\\n    - At p = 1, one gets Sum Pooling (which is proportional to average pooling)\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the height and width dimension\\n        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,\\n          and the second `int` for the width dimension\\n\\n    .. note:: If the sum to the power of `p` is zero, the gradient of this function is\\n              not defined. This implementation will set the gradient to zero in this case.\\n\\n    Args:\\n        kernel_size: the size of the window\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in} - \\\\text{kernel\\\\_size}[0]}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in} - \\\\text{kernel\\\\_size}[1]}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> # power-2 pool of square window of size=3, stride=2\\n        >>> m = nn.LPPool2d(2, 3, stride=2)\\n        >>> # pool of non-square window of power 1.2\\n        >>> m = nn.LPPool2d(1.2, (3, 2), stride=(2, 1))\\n        >>> input = torch.randn(20, 16, 50, 32)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_2_t\\n    stride: _size_2_t\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.lp_pool2d(\\n            input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode\\n        )\\n\\n\\nclass LPPool3d(_LPPoolNd):\\n    r\\\"\\\"\\\"Applies a 3D power-average pooling over an input signal composed of several input planes.\\n\\n    On each window, the function computed is:\\n\\n    .. math::\\n        f(X) = \\\\sqrt[p]{\\\\sum_{x \\\\in X} x^{p}}\\n\\n    - At p = :math:`\\\\infty`, one gets Max Pooling\\n    - At p = 1, one gets Sum Pooling (which is proportional to average pooling)\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the height, width and depth dimension\\n        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,\\n          the second `int` for the height dimension and the third `int` for the width dimension\\n\\n    .. note:: If the sum to the power of `p` is zero, the gradient of this function is\\n              not defined. This implementation will set the gradient to zero in this case.\\n\\n    Args:\\n        kernel_size: the size of the window\\n        stride: the stride of the window. Default value is :attr:`kernel_size`\\n        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or\\n          :math:`(C, D_{out}, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              D_{out} = \\\\left\\\\lfloor\\\\frac{D_{in} - \\\\text{kernel\\\\_size}[0]}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in} - \\\\text{kernel\\\\_size}[1]}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in} - \\\\text{kernel\\\\_size}[2]}{\\\\text{stride}[2]} + 1\\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> # power-2 pool of square window of size=3, stride=2\\n        >>> m = nn.LPPool3d(2, 3, stride=2)\\n        >>> # pool of non-square window of power 1.2\\n        >>> m = nn.LPPool3d(1.2, (3, 2, 2), stride=(2, 1, 2))\\n        >>> input = torch.randn(20, 16, 50, 44, 31)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    kernel_size: _size_3_t\\n    stride: _size_3_t\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.lp_pool3d(\\n            input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode\\n        )\\n\\n\\nclass _AdaptiveMaxPoolNd(Module):\\n    __constants__ = [\\\"output_size\\\", \\\"return_indices\\\"]\\n    return_indices: bool\\n\\n    def __init__(\\n        self, output_size: _size_any_opt_t, return_indices: bool = False\\n    ) -> None:\\n        super().__init__()\\n        self.output_size = output_size\\n        self.return_indices = return_indices\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"output_size={self.output_size}\\\"\\n\\n\\n# FIXME (by @ssnl): Improve adaptive pooling docs: specify what the input and\\n#   output shapes are, and how the operation computes output.\\n\\n\\nclass AdaptiveMaxPool1d(_AdaptiveMaxPoolNd):\\n    r\\\"\\\"\\\"Applies a 1D adaptive max pooling over an input signal composed of several input planes.\\n\\n    The output size is :math:`L_{out}`, for any input size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        output_size: the target output size :math:`L_{out}`.\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to nn.MaxUnpool1d. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.\\n        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where\\n          :math:`L_{out}=\\\\text{output\\\\_size}`.\\n\\n    Examples:\\n        >>> # target output size of 5\\n        >>> m = nn.AdaptiveMaxPool1d(5)\\n        >>> input = torch.randn(1, 64, 8)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    output_size: _size_1_t\\n\\n    def forward(self, input: Tensor):\\n        return F.adaptive_max_pool1d(input, self.output_size, self.return_indices)\\n\\n\\nclass AdaptiveMaxPool2d(_AdaptiveMaxPoolNd):\\n    r\\\"\\\"\\\"Applies a 2D adaptive max pooling over an input signal composed of several input planes.\\n\\n    The output is of size :math:`H_{out} \\\\times W_{out}`, for any input size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        output_size: the target output size of the image of the form :math:`H_{out} \\\\times W_{out}`.\\n                     Can be a tuple :math:`(H_{out}, W_{out})` or a single :math:`H_{out}` for a\\n                     square image :math:`H_{out} \\\\times H_{out}`. :math:`H_{out}` and :math:`W_{out}`\\n                     can be either a ``int``, or ``None`` which means the size will be the same as that\\n                     of the input.\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to nn.MaxUnpool2d. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n          :math:`(H_{out}, W_{out})=\\\\text{output\\\\_size}`.\\n\\n    Examples:\\n        >>> # target output size of 5x7\\n        >>> m = nn.AdaptiveMaxPool2d((5, 7))\\n        >>> input = torch.randn(1, 64, 8, 9)\\n        >>> output = m(input)\\n        >>> # target output size of 7x7 (square)\\n        >>> m = nn.AdaptiveMaxPool2d(7)\\n        >>> input = torch.randn(1, 64, 10, 9)\\n        >>> output = m(input)\\n        >>> # target output size of 10x7\\n        >>> m = nn.AdaptiveMaxPool2d((None, 7))\\n        >>> input = torch.randn(1, 64, 10, 9)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    output_size: _size_2_opt_t\\n\\n    def forward(self, input: Tensor):\\n        return F.adaptive_max_pool2d(input, self.output_size, self.return_indices)\\n\\n\\nclass AdaptiveMaxPool3d(_AdaptiveMaxPoolNd):\\n    r\\\"\\\"\\\"Applies a 3D adaptive max pooling over an input signal composed of several input planes.\\n\\n    The output is of size :math:`D_{out} \\\\times H_{out} \\\\times W_{out}`, for any input size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        output_size: the target output size of the image of the form :math:`D_{out} \\\\times H_{out} \\\\times W_{out}`.\\n                     Can be a tuple :math:`(D_{out}, H_{out}, W_{out})` or a single\\n                     :math:`D_{out}` for a cube :math:`D_{out} \\\\times D_{out} \\\\times D_{out}`.\\n                     :math:`D_{out}`, :math:`H_{out}` and :math:`W_{out}` can be either a\\n                     ``int``, or ``None`` which means the size will be the same as that of the input.\\n\\n        return_indices: if ``True``, will return the indices along with the outputs.\\n                        Useful to pass to nn.MaxUnpool3d. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,\\n          where :math:`(D_{out}, H_{out}, W_{out})=\\\\text{output\\\\_size}`.\\n\\n    Examples:\\n        >>> # target output size of 5x7x9\\n        >>> m = nn.AdaptiveMaxPool3d((5, 7, 9))\\n        >>> input = torch.randn(1, 64, 8, 9, 10)\\n        >>> output = m(input)\\n        >>> # target output size of 7x7x7 (cube)\\n        >>> m = nn.AdaptiveMaxPool3d(7)\\n        >>> input = torch.randn(1, 64, 10, 9, 8)\\n        >>> output = m(input)\\n        >>> # target output size of 7x9x8\\n        >>> m = nn.AdaptiveMaxPool3d((7, None, None))\\n        >>> input = torch.randn(1, 64, 10, 9, 8)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    output_size: _size_3_opt_t\\n\\n    def forward(self, input: Tensor):\\n        return F.adaptive_max_pool3d(input, self.output_size, self.return_indices)\\n\\n\\nclass _AdaptiveAvgPoolNd(Module):\\n    __constants__ = [\\\"output_size\\\"]\\n\\n    def __init__(self, output_size: _size_any_opt_t) -> None:\\n        super().__init__()\\n        self.output_size = output_size\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"output_size={self.output_size}\\\"\\n\\n\\nclass AdaptiveAvgPool1d(_AdaptiveAvgPoolNd):\\n    r\\\"\\\"\\\"Applies a 1D adaptive average pooling over an input signal composed of several input planes.\\n\\n    The output size is :math:`L_{out}`, for any input size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        output_size: the target output size :math:`L_{out}`.\\n\\n    Shape:\\n        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.\\n        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where\\n          :math:`L_{out}=\\\\text{output\\\\_size}`.\\n\\n    Examples:\\n        >>> # target output size of 5\\n        >>> m = nn.AdaptiveAvgPool1d(5)\\n        >>> input = torch.randn(1, 64, 8)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    output_size: _size_1_t\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.adaptive_avg_pool1d(input, self.output_size)\\n\\n\\nclass AdaptiveAvgPool2d(_AdaptiveAvgPoolNd):\\n    r\\\"\\\"\\\"Applies a 2D adaptive average pooling over an input signal composed of several input planes.\\n\\n    The output is of size H x W, for any input size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        output_size: the target output size of the image of the form H x W.\\n                     Can be a tuple (H, W) or a single H for a square image H x H.\\n                     H and W can be either a ``int``, or ``None`` which means the size will\\n                     be the same as that of the input.\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, S_{0}, S_{1})` or :math:`(C, S_{0}, S_{1})`, where\\n          :math:`S=\\\\text{output\\\\_size}`.\\n\\n    Examples:\\n        >>> # target output size of 5x7\\n        >>> m = nn.AdaptiveAvgPool2d((5, 7))\\n        >>> input = torch.randn(1, 64, 8, 9)\\n        >>> output = m(input)\\n        >>> # target output size of 7x7 (square)\\n        >>> m = nn.AdaptiveAvgPool2d(7)\\n        >>> input = torch.randn(1, 64, 10, 9)\\n        >>> output = m(input)\\n        >>> # target output size of 10x7\\n        >>> m = nn.AdaptiveAvgPool2d((None, 7))\\n        >>> input = torch.randn(1, 64, 10, 9)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    output_size: _size_2_opt_t\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.adaptive_avg_pool2d(input, self.output_size)\\n\\n\\nclass AdaptiveAvgPool3d(_AdaptiveAvgPoolNd):\\n    r\\\"\\\"\\\"Applies a 3D adaptive average pooling over an input signal composed of several input planes.\\n\\n    The output is of size D x H x W, for any input size.\\n    The number of output features is equal to the number of input planes.\\n\\n    Args:\\n        output_size: the target output size of the form D x H x W.\\n                     Can be a tuple (D, H, W) or a single number D for a cube D x D x D.\\n                     D, H and W can be either a ``int``, or ``None`` which means the size will\\n                     be the same as that of the input.\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, S_{0}, S_{1}, S_{2})` or :math:`(C, S_{0}, S_{1}, S_{2})`,\\n          where :math:`S=\\\\text{output\\\\_size}`.\\n\\n    Examples:\\n        >>> # target output size of 5x7x9\\n        >>> m = nn.AdaptiveAvgPool3d((5, 7, 9))\\n        >>> input = torch.randn(1, 64, 8, 9, 10)\\n        >>> output = m(input)\\n        >>> # target output size of 7x7x7 (cube)\\n        >>> m = nn.AdaptiveAvgPool3d(7)\\n        >>> input = torch.randn(1, 64, 10, 9, 8)\\n        >>> output = m(input)\\n        >>> # target output size of 7x9x8\\n        >>> m = nn.AdaptiveAvgPool3d((7, None, None))\\n        >>> input = torch.randn(1, 64, 10, 9, 8)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    output_size: _size_3_opt_t\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.adaptive_avg_pool3d(input, self.output_size)\\n\\n\\n# mypy: allow-untyped-defs\\nimport numbers\\nfrom typing import List, Optional, Tuple, Union\\n\\nimport torch\\nfrom torch import Size, Tensor\\nfrom torch.nn import functional as F, init\\nfrom torch.nn.parameter import Parameter\\n\\nfrom ._functions import CrossMapLRN2d as _cross_map_lrn2d\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"LocalResponseNorm\\\", \\\"CrossMapLRN2d\\\", \\\"LayerNorm\\\", \\\"GroupNorm\\\", \\\"RMSNorm\\\"]\\n\\n\\nclass LocalResponseNorm(Module):\\n    r\\\"\\\"\\\"Applies local response normalization over an input signal.\\n\\n    The input signal is composed of several input planes, where channels occupy the second dimension.\\n    Applies normalization across channels.\\n\\n    .. math::\\n        b_{c} = a_{c}\\\\left(k + \\\\frac{\\\\alpha}{n}\\n        \\\\sum_{c'=\\\\max(0, c-n/2)}^{\\\\min(N-1,c+n/2)}a_{c'}^2\\\\right)^{-\\\\beta}\\n\\n    Args:\\n        size: amount of neighbouring channels used for normalization\\n        alpha: multiplicative factor. Default: 0.0001\\n        beta: exponent. Default: 0.75\\n        k: additive factor. Default: 1\\n\\n    Shape:\\n        - Input: :math:`(N, C, *)`\\n        - Output: :math:`(N, C, *)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> lrn = nn.LocalResponseNorm(2)\\n        >>> signal_2d = torch.randn(32, 5, 24, 24)\\n        >>> signal_4d = torch.randn(16, 5, 7, 7, 7, 7)\\n        >>> output_2d = lrn(signal_2d)\\n        >>> output_4d = lrn(signal_4d)\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"size\\\", \\\"alpha\\\", \\\"beta\\\", \\\"k\\\"]\\n    size: int\\n    alpha: float\\n    beta: float\\n    k: float\\n\\n    def __init__(\\n        self, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1.0\\n    ) -> None:\\n        super().__init__()\\n        self.size = size\\n        self.alpha = alpha\\n        self.beta = beta\\n        self.k = k\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.local_response_norm(input, self.size, self.alpha, self.beta, self.k)\\n\\n    def extra_repr(self):\\n        return \\\"{size}, alpha={alpha}, beta={beta}, k={k}\\\".format(**self.__dict__)\\n\\n\\nclass CrossMapLRN2d(Module):\\n    size: int\\n    alpha: float\\n    beta: float\\n    k: float\\n\\n    def __init__(\\n        self, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1\\n    ) -> None:\\n        super().__init__()\\n        self.size = size\\n        self.alpha = alpha\\n        self.beta = beta\\n        self.k = k\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return _cross_map_lrn2d.apply(input, self.size, self.alpha, self.beta, self.k)\\n\\n    def extra_repr(self) -> str:\\n        return \\\"{size}, alpha={alpha}, beta={beta}, k={k}\\\".format(**self.__dict__)\\n\\n\\n_shape_t = Union[int, List[int], Size]\\n\\n\\nclass LayerNorm(Module):\\n    r\\\"\\\"\\\"Applies Layer Normalization over a mini-batch of inputs.\\n\\n    This layer implements the operation as described in\\n    the paper `Layer Normalization <https://arxiv.org/abs/1607.06450>`__\\n\\n    .. math::\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated over the last `D` dimensions, where `D`\\n    is the dimension of :attr:`normalized_shape`. For example, if :attr:`normalized_shape`\\n    is ``(3, 5)`` (a 2-dimensional shape), the mean and standard-deviation are computed over\\n    the last 2 dimensions of the input (i.e. ``input.mean((-2, -1))``).\\n    :math:`\\\\gamma` and :math:`\\\\beta` are learnable affine transform parameters of\\n    :attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``.\\n    The standard-deviation is calculated via the biased estimator, equivalent to\\n    `torch.var(input, unbiased=False)`.\\n\\n    .. note::\\n        Unlike Batch Normalization and Instance Normalization, which applies\\n        scalar scale and bias for each entire channel/plane with the\\n        :attr:`affine` option, Layer Normalization applies per-element scale and\\n        bias with :attr:`elementwise_affine`.\\n\\n    This layer uses statistics computed from input data in both training and\\n    evaluation modes.\\n\\n    Args:\\n        normalized_shape (int or list or torch.Size): input shape from an expected input\\n            of size\\n\\n            .. math::\\n                [* \\\\times \\\\text{normalized\\\\_shape}[0] \\\\times \\\\text{normalized\\\\_shape}[1]\\n                    \\\\times \\\\ldots \\\\times \\\\text{normalized\\\\_shape}[-1]]\\n\\n            If a single integer is used, it is treated as a singleton list, and this module will\\n            normalize over the last dimension which is expected to be of that specific size.\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        elementwise_affine: a boolean value that when set to ``True``, this module\\n            has learnable per-element affine parameters initialized to ones (for weights)\\n            and zeros (for biases). Default: ``True``.\\n        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if\\n            :attr:`elementwise_affine` is ``True``). Default: ``True``.\\n\\n    Attributes:\\n        weight: the learnable weights of the module of shape\\n            :math:`\\\\text{normalized\\\\_shape}` when :attr:`elementwise_affine` is set to ``True``.\\n            The values are initialized to 1.\\n        bias:   the learnable bias of the module of shape\\n                :math:`\\\\text{normalized\\\\_shape}` when :attr:`elementwise_affine` is set to ``True``.\\n                The values are initialized to 0.\\n\\n    Shape:\\n        - Input: :math:`(N, *)`\\n        - Output: :math:`(N, *)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # NLP Example\\n        >>> batch, sentence_length, embedding_dim = 20, 5, 10\\n        >>> embedding = torch.randn(batch, sentence_length, embedding_dim)\\n        >>> layer_norm = nn.LayerNorm(embedding_dim)\\n        >>> # Activate module\\n        >>> layer_norm(embedding)\\n        >>>\\n        >>> # Image Example\\n        >>> N, C, H, W = 20, 5, 10, 10\\n        >>> input = torch.randn(N, C, H, W)\\n        >>> # Normalize over the last three dimensions (i.e. the channel and spatial dimensions)\\n        >>> # as shown in the image below\\n        >>> layer_norm = nn.LayerNorm([C, H, W])\\n        >>> output = layer_norm(input)\\n\\n    .. image:: ../_static/img/nn/layer_norm.jpg\\n        :scale: 50 %\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"normalized_shape\\\", \\\"eps\\\", \\\"elementwise_affine\\\"]\\n    normalized_shape: Tuple[int, ...]\\n    eps: float\\n    elementwise_affine: bool\\n\\n    def __init__(\\n        self,\\n        normalized_shape: _shape_t,\\n        eps: float = 1e-5,\\n        elementwise_affine: bool = True,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        if isinstance(normalized_shape, numbers.Integral):\\n            # mypy error: incompatible types in assignment\\n            normalized_shape = (normalized_shape,)  # type: ignore[assignment]\\n        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]\\n        self.eps = eps\\n        self.elementwise_affine = elementwise_affine\\n        if self.elementwise_affine:\\n            self.weight = Parameter(\\n                torch.empty(self.normalized_shape, **factory_kwargs)\\n            )\\n            if bias:\\n                self.bias = Parameter(\\n                    torch.empty(self.normalized_shape, **factory_kwargs)\\n                )\\n            else:\\n                self.register_parameter(\\\"bias\\\", None)\\n        else:\\n            self.register_parameter(\\\"weight\\\", None)\\n            self.register_parameter(\\\"bias\\\", None)\\n\\n        self.reset_parameters()\\n\\n    def reset_parameters(self) -> None:\\n        if self.elementwise_affine:\\n            init.ones_(self.weight)\\n            if self.bias is not None:\\n                init.zeros_(self.bias)\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.layer_norm(\\n            input, self.normalized_shape, self.weight, self.bias, self.eps\\n        )\\n\\n    def extra_repr(self) -> str:\\n        return (\\n            \\\"{normalized_shape}, eps={eps}, \\\"\\n            \\\"elementwise_affine={elementwise_affine}\\\".format(**self.__dict__)\\n        )\\n\\n\\nclass GroupNorm(Module):\\n    r\\\"\\\"\\\"Applies Group Normalization over a mini-batch of inputs.\\n\\n    This layer implements the operation as described in\\n    the paper `Group Normalization <https://arxiv.org/abs/1803.08494>`__\\n\\n    .. math::\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The input channels are separated into :attr:`num_groups` groups, each containing\\n    ``num_channels / num_groups`` channels. :attr:`num_channels` must be divisible by\\n    :attr:`num_groups`. The mean and standard-deviation are calculated\\n    separately over the each group. :math:`\\\\gamma` and :math:`\\\\beta` are learnable\\n    per-channel affine transform parameter vectors of size :attr:`num_channels` if\\n    :attr:`affine` is ``True``.\\n    The standard-deviation is calculated via the biased estimator, equivalent to\\n    `torch.var(input, unbiased=False)`.\\n\\n    This layer uses statistics computed from input data in both training and\\n    evaluation modes.\\n\\n    Args:\\n        num_groups (int): number of groups to separate the channels into\\n        num_channels (int): number of channels expected in input\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        affine: a boolean value that when set to ``True``, this module\\n            has learnable per-channel affine parameters initialized to ones (for weights)\\n            and zeros (for biases). Default: ``True``.\\n\\n    Shape:\\n        - Input: :math:`(N, C, *)` where :math:`C=\\\\text{num\\\\_channels}`\\n        - Output: :math:`(N, C, *)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> input = torch.randn(20, 6, 10, 10)\\n        >>> # Separate 6 channels into 3 groups\\n        >>> m = nn.GroupNorm(3, 6)\\n        >>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm)\\n        >>> m = nn.GroupNorm(6, 6)\\n        >>> # Put all 6 channels into a single group (equivalent with LayerNorm)\\n        >>> m = nn.GroupNorm(1, 6)\\n        >>> # Activating the module\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"num_groups\\\", \\\"num_channels\\\", \\\"eps\\\", \\\"affine\\\"]\\n    num_groups: int\\n    num_channels: int\\n    eps: float\\n    affine: bool\\n\\n    def __init__(\\n        self,\\n        num_groups: int,\\n        num_channels: int,\\n        eps: float = 1e-5,\\n        affine: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        if num_channels % num_groups != 0:\\n            raise ValueError(\\\"num_channels must be divisible by num_groups\\\")\\n\\n        self.num_groups = num_groups\\n        self.num_channels = num_channels\\n        self.eps = eps\\n        self.affine = affine\\n        if self.affine:\\n            self.weight = Parameter(torch.empty(num_channels, **factory_kwargs))\\n            self.bias = Parameter(torch.empty(num_channels, **factory_kwargs))\\n        else:\\n            self.register_parameter(\\\"weight\\\", None)\\n            self.register_parameter(\\\"bias\\\", None)\\n\\n        self.reset_parameters()\\n\\n    def reset_parameters(self) -> None:\\n        if self.affine:\\n            init.ones_(self.weight)\\n            init.zeros_(self.bias)\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)\\n\\n    def extra_repr(self) -> str:\\n        return \\\"{num_groups}, {num_channels}, eps={eps}, \\\" \\\"affine={affine}\\\".format(\\n            **self.__dict__\\n        )\\n\\n\\nclass RMSNorm(Module):\\n    r\\\"\\\"\\\"Applies Root Mean Square Layer Normalization over a mini-batch of inputs.\\n\\n    This layer implements the operation as described in\\n    the paper `Root Mean Square Layer Normalization <https://arxiv.org/pdf/1910.07467.pdf>`__\\n\\n    .. math::\\n        y = \\\\frac{x}{\\\\sqrt{\\\\mathrm{RMS}[x] + \\\\epsilon}} * \\\\gamma\\n\\n    The root mean squared norm is taken over the last ``D`` dimensions, where ``D``\\n    is the dimension of :attr:`normalized_shape`. For example, if :attr:`normalized_shape`\\n    is ``(3, 5)`` (a 2-dimensional shape), the rms norm is computed over\\n    the last 2 dimensions of the input.\\n\\n    Args:\\n        normalized_shape (int or list or torch.Size): input shape from an expected input\\n            of size\\n\\n            .. math::\\n                [* \\\\times \\\\text{normalized\\\\_shape}[0] \\\\times \\\\text{normalized\\\\_shape}[1]\\n                    \\\\times \\\\ldots \\\\times \\\\text{normalized\\\\_shape}[-1]]\\n\\n            If a single integer is used, it is treated as a singleton list, and this module will\\n            normalize over the last dimension which is expected to be of that specific size.\\n        eps: a value added to the denominator for numerical stability. Default: :func:`torch.finfo(x.dtype).eps`\\n        elementwise_affine: a boolean value that when set to ``True``, this module\\n            has learnable per-element affine parameters initialized to ones (for weights)\\n            and zeros (for biases). Default: ``True``.\\n\\n    Shape:\\n        - Input: :math:`(N, *)`\\n        - Output: :math:`(N, *)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> rms_norm = nn.RMSNorm([2, 3])\\n        >>> input = torch.randn(2, 2, 3)\\n        >>> rms_norm(input)\\n\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"normalized_shape\\\", \\\"eps\\\", \\\"elementwise_affine\\\"]\\n    normalized_shape: Tuple[int, ...]\\n    eps: Optional[float]\\n    elementwise_affine: bool\\n\\n    def __init__(\\n        self,\\n        normalized_shape: _shape_t,\\n        eps: Optional[float] = None,\\n        elementwise_affine: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        if isinstance(normalized_shape, numbers.Integral):\\n            # mypy error: incompatible types in assignment\\n            normalized_shape = (normalized_shape,)  # type: ignore[assignment]\\n        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]\\n        self.eps = eps\\n        self.elementwise_affine = elementwise_affine\\n        if self.elementwise_affine:\\n            self.weight = Parameter(\\n                torch.empty(self.normalized_shape, **factory_kwargs)\\n            )\\n        else:\\n            self.register_parameter(\\\"weight\\\", None)\\n        self.reset_parameters()\\n\\n    def reset_parameters(self) -> None:\\n        \\\"\\\"\\\"\\n        Resets parameters based on their initialization used in __init__.\\n        \\\"\\\"\\\"\\n        if self.elementwise_affine:\\n            init.ones_(self.weight)\\n\\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\\n        \\\"\\\"\\\"\\n        Runs forward pass.\\n        \\\"\\\"\\\"\\n        return F.rms_norm(x, self.normalized_shape, self.weight, self.eps)\\n\\n    def extra_repr(self) -> str:\\n        \\\"\\\"\\\"\\n        Extra information about the module.\\n        \\\"\\\"\\\"\\n        return (\\n            \\\"{normalized_shape}, eps={eps}, \\\"\\n            \\\"elementwise_affine={elementwise_affine}\\\".format(**self.__dict__)\\n        )\\n\\n\\n# TODO: ContrastiveNorm2d\\n# TODO: DivisiveNorm2d\\n# TODO: SubtractiveNorm2d\\n\\n\\n# mypy: allow-untyped-defs\\n\\nimport functools\\nimport inspect\\nimport itertools\\nimport warnings\\nimport weakref\\nfrom collections import namedtuple, OrderedDict\\nfrom typing import (\\n    Any,\\n    Callable,\\n    Dict,\\n    Iterator,\\n    List,\\n    Mapping,\\n    Optional,\\n    overload,\\n    Set,\\n    Tuple,\\n    TypeVar,\\n    Union,\\n)\\nfrom typing_extensions import Self\\n\\nimport torch\\nfrom torch import device, dtype, Tensor\\nfrom torch._prims_common import DeviceLikeType\\nfrom torch.nn.parameter import Buffer, Parameter\\nfrom torch.utils._python_dispatch import is_traceable_wrapper_subclass\\nfrom torch.utils.hooks import BackwardHook, RemovableHandle\\n\\n\\n__all__ = [\\n    \\\"register_module_forward_pre_hook\\\",\\n    \\\"register_module_forward_hook\\\",\\n    \\\"register_module_full_backward_pre_hook\\\",\\n    \\\"register_module_backward_hook\\\",\\n    \\\"register_module_full_backward_hook\\\",\\n    \\\"register_module_buffer_registration_hook\\\",\\n    \\\"register_module_module_registration_hook\\\",\\n    \\\"register_module_parameter_registration_hook\\\",\\n    \\\"Module\\\",\\n]\\n\\n_grad_t = Union[Tuple[Tensor, ...], Tensor]\\n# See https://mypy.readthedocs.io/en/latest/generics.html#generic-methods-and-generic-self for the use\\n# of `T` to annotate `self`. Many methods of `Module` return `self` and we want those return values to be\\n# the type of the subclass, not the looser type of `Module`.\\nT = TypeVar(\\\"T\\\", bound=\\\"Module\\\")\\n\\n\\nclass _IncompatibleKeys(\\n    namedtuple(\\\"IncompatibleKeys\\\", [\\\"missing_keys\\\", \\\"unexpected_keys\\\"]),\\n):\\n    def __repr__(self):\\n        if not self.missing_keys and not self.unexpected_keys:\\n            return \\\"<All keys matched successfully>\\\"\\n        return super().__repr__()\\n\\n    __str__ = __repr__\\n\\n\\ndef _addindent(s_, numSpaces):\\n    s = s_.split(\\\"\\\\n\\\")\\n    # don't do anything for single-line stuff\\n    if len(s) == 1:\\n        return s_\\n    first = s.pop(0)\\n    s = [(numSpaces * \\\" \\\") + line for line in s]\\n    s = \\\"\\\\n\\\".join(s)\\n    s = first + \\\"\\\\n\\\" + s\\n    return s\\n\\n\\nr\\\"\\\"\\\"This tracks hooks common to all modules that are executed immediately before\\n.registering the buffer/module/parameter\\\"\\\"\\\"\\n_global_buffer_registration_hooks: Dict[int, Callable] = OrderedDict()\\n_global_module_registration_hooks: Dict[int, Callable] = OrderedDict()\\n_global_parameter_registration_hooks: Dict[int, Callable] = OrderedDict()\\n\\n\\nclass _WrappedHook:\\n    def __init__(self, hook: Callable, module: Optional[\\\"Module\\\"] = None):\\n        self.hook: Callable = hook\\n        functools.update_wrapper(self, hook)\\n\\n        self.with_module: bool = False\\n\\n        if module is not None:\\n            self.module: weakref.ReferenceType[Module] = weakref.ref(module)\\n            self.with_module = True\\n\\n    def __call__(self, *args: Any, **kwargs: Any) -> Any:\\n        if self.with_module:\\n            module = self.module()\\n            if module is None:\\n                raise RuntimeError(\\\"You are trying to call the hook of a dead Module!\\\")\\n            return self.hook(module, *args, **kwargs)\\n        return self.hook(*args, **kwargs)\\n\\n    def __getstate__(self) -> Dict:\\n        result = {\\\"hook\\\": self.hook, \\\"with_module\\\": self.with_module}\\n        if self.with_module:\\n            result[\\\"module\\\"] = self.module()\\n\\n        return result\\n\\n    def __setstate__(self, state: Dict):\\n        self.hook = state[\\\"hook\\\"]\\n        self.with_module = state[\\\"with_module\\\"]\\n\\n        if self.with_module:\\n            if state[\\\"module\\\"] is None:\\n                raise RuntimeError(\\n                    \\\"You are trying to revive the hook of a dead Module!\\\"\\n                )\\n            self.module = weakref.ref(state[\\\"module\\\"])\\n\\n\\nr\\\"\\\"\\\"This tracks hooks common to all modules that are executed before/after\\ncalling forward and backward. This is global state used for debugging/profiling\\npurposes\\\"\\\"\\\"\\n_global_backward_pre_hooks: Dict[int, Callable] = OrderedDict()\\n_global_backward_hooks: Dict[int, Callable] = OrderedDict()\\n_global_is_full_backward_hook: Optional[bool] = None\\n_global_forward_pre_hooks: Dict[int, Callable] = OrderedDict()\\n_global_forward_hooks: Dict[int, Callable] = OrderedDict()\\n_global_forward_hooks_always_called: Dict[int, bool] = OrderedDict()\\n\\n_EXTRA_STATE_KEY_SUFFIX = \\\"_extra_state\\\"\\n\\n\\ndef register_module_buffer_registration_hook(\\n    hook: Callable[..., None],\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a buffer registration hook common to all modules.\\n\\n    .. warning ::\\n\\n        This adds global state to the `nn.Module` module\\n\\n    The hook will be called every time :func:`register_buffer` is invoked.\\n    It should have the following signature::\\n\\n        hook(module, name, buffer) -> None or new buffer\\n\\n    The hook can modify the input or return a single modified value in the hook.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n    \\\"\\\"\\\"\\n    handle = RemovableHandle(_global_buffer_registration_hooks)\\n    _global_buffer_registration_hooks[handle.id] = hook\\n    return handle\\n\\n\\ndef register_module_module_registration_hook(\\n    hook: Callable[..., None],\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a module registration hook common to all modules.\\n\\n    .. warning ::\\n\\n        This adds global state to the `nn.Module` module\\n\\n    The hook will be called every time :func:`register_module` is invoked.\\n    It should have the following signature::\\n\\n        hook(module, name, submodule) -> None or new submodule\\n\\n    The hook can modify the input or return a single modified value in the hook.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n    \\\"\\\"\\\"\\n    handle = RemovableHandle(_global_module_registration_hooks)\\n    _global_module_registration_hooks[handle.id] = hook\\n    return handle\\n\\n\\ndef register_module_parameter_registration_hook(\\n    hook: Callable[..., None],\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a parameter registration hook common to all modules.\\n\\n    .. warning ::\\n\\n        This adds global state to the `nn.Module` module\\n\\n    The hook will be called every time :func:`register_parameter` is invoked.\\n    It should have the following signature::\\n\\n        hook(module, name, param) -> None or new parameter\\n\\n    The hook can modify the input or return a single modified value in the hook.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n    \\\"\\\"\\\"\\n    handle = RemovableHandle(_global_parameter_registration_hooks)\\n    _global_parameter_registration_hooks[handle.id] = hook\\n    return handle\\n\\n\\ndef register_module_forward_pre_hook(hook: Callable[..., None]) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a forward pre-hook common to all modules.\\n\\n    .. warning ::\\n\\n        This adds global state to the `nn.module` module\\n        and it is only intended for debugging/profiling purposes.\\n\\n    The hook will be called every time before :func:`forward` is invoked.\\n    It should have the following signature::\\n\\n        hook(module, input) -> None or modified input\\n\\n    The input contains only the positional arguments given to the module.\\n    Keyword arguments won't be passed to the hooks and only to the ``forward``.\\n    The hook can modify the input. User can either return a tuple or a\\n    single modified value in the hook. We will wrap the value into a tuple\\n    if a single value is returned(unless that value is already a tuple).\\n\\n    This hook has precedence over the specific module hooks registered with\\n    ``register_forward_pre_hook``.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n    \\\"\\\"\\\"\\n    handle = RemovableHandle(_global_forward_pre_hooks)\\n    _global_forward_pre_hooks[handle.id] = hook\\n    return handle\\n\\n\\ndef register_module_forward_hook(\\n    hook: Callable[..., None],\\n    *,\\n    always_call: bool = False,\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a global forward hook for all the modules.\\n\\n    .. warning ::\\n\\n        This adds global state to the `nn.module` module\\n        and it is only intended for debugging/profiling purposes.\\n\\n    The hook will be called every time after :func:`forward` has computed an output.\\n    It should have the following signature::\\n\\n        hook(module, input, output) -> None or modified output\\n\\n    The input contains only the positional arguments given to the module.\\n    Keyword arguments won't be passed to the hooks and only to the ``forward``.\\n    The hook can modify the output. It can modify the input inplace but\\n    it will not have effect on forward since this is called after\\n    :func:`forward` is called.\\n\\n    Parameters:\\n        hook (Callable): The user defined hook to be registered.\\n        always_call (bool): If ``True`` the ``hook`` will be run regardless of\\n            whether an exception is raised while calling the Module.\\n            Default: ``False``\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n\\n    This hook will be executed before specific module hooks registered with\\n    ``register_forward_hook``.\\n    \\\"\\\"\\\"\\n    handle = RemovableHandle(\\n        _global_forward_hooks, extra_dict=_global_forward_hooks_always_called\\n    )\\n    _global_forward_hooks[handle.id] = hook\\n    if always_call:\\n        _global_forward_hooks_always_called[handle.id] = True\\n    return handle\\n\\n\\ndef register_module_backward_hook(\\n    hook: Callable[[\\\"Module\\\", _grad_t, _grad_t], Union[None, _grad_t]],\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a backward hook common to all the modules.\\n\\n    This function is deprecated in favor of\\n    :func:`torch.nn.modules.module.register_module_full_backward_hook`\\n    and the behavior of this function will change in future versions.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n\\n    \\\"\\\"\\\"\\n    global _global_is_full_backward_hook\\n    if _global_is_full_backward_hook is True:\\n        raise RuntimeError(\\n            \\\"Cannot use both regular backward hooks and full backward hooks as a \\\"\\n            \\\"global Module hook. Please use only one of them.\\\"\\n        )\\n\\n    _global_is_full_backward_hook = False\\n\\n    handle = RemovableHandle(_global_backward_hooks)\\n    _global_backward_hooks[handle.id] = hook\\n    return handle\\n\\n\\ndef register_module_full_backward_pre_hook(\\n    hook: Callable[[\\\"Module\\\", _grad_t], Union[None, _grad_t]],\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a backward pre-hook common to all the modules.\\n\\n    .. warning ::\\n        This adds global state to the `nn.module` module\\n        and it is only intended for debugging/profiling purposes.\\n\\n    Hooks registered using this function behave in the same way as those\\n    registered by :meth:`torch.nn.Module.register_full_backward_pre_hook`.\\n    Refer to its documentation for more details.\\n\\n    Hooks registered using this function will be called before hooks registered\\n    using :meth:`torch.nn.Module.register_full_backward_pre_hook`.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n\\n    \\\"\\\"\\\"\\n    handle = RemovableHandle(_global_backward_pre_hooks)\\n    _global_backward_pre_hooks[handle.id] = hook\\n    return handle\\n\\n\\ndef register_module_full_backward_hook(\\n    hook: Callable[[\\\"Module\\\", _grad_t, _grad_t], Union[None, _grad_t]],\\n) -> RemovableHandle:\\n    r\\\"\\\"\\\"Register a backward hook common to all the modules.\\n\\n    .. warning ::\\n        This adds global state to the `nn.module` module\\n        and it is only intended for debugging/profiling purposes.\\n\\n    Hooks registered using this function behave in the same way as those\\n    registered by :meth:`torch.nn.Module.register_full_backward_hook`.\\n    Refer to its documentation for more details.\\n\\n    Hooks registered using this function will be called before hooks registered\\n    using :meth:`torch.nn.Module.register_full_backward_hook`.\\n\\n    Returns:\\n        :class:`torch.utils.hooks.RemovableHandle`:\\n            a handle that can be used to remove the added hook by calling\\n            ``handle.remove()``\\n\\n    \\\"\\\"\\\"\\n    global _global_is_full_backward_hook\\n    if _global_is_full_backward_hook is False:\\n        raise RuntimeError(\\n            \\\"Cannot use both regular backward hooks and full backward hooks as a \\\"\\n            \\\"global Module hook. Please use only one of them.\\\"\\n        )\\n\\n    _global_is_full_backward_hook = True\\n\\n    handle = RemovableHandle(_global_backward_hooks)\\n    _global_backward_hooks[handle.id] = hook\\n    return handle\\n\\n\\n# Trick mypy into not applying contravariance rules to inputs by defining\\n# forward as a value, rather than a function.  See also\\n# https://github.com/python/mypy/issues/8795\\ndef _forward_unimplemented(self, *input: Any) -> None:\\n    r\\\"\\\"\\\"Define the computation performed at every call.\\n\\n    Should be overridden by all subclasses.\\n\\n    .. note::\\n        Although the recipe for forward pass needs to be defined within\\n        this function, one should call the :class:`Module` instance afterwards\\n        instead of this since the former takes care of running the\\n        registered hooks while the latter silently ignores them.\\n    \\\"\\\"\\\"\\n    raise NotImplementedError(\\n        f'Module [{type(self).__name__}] is missing the required \\\"forward\\\" function'\\n    )\\n\\n\\nclass Module:\\n    r\\\"\\\"\\\"Base class for all neural network modules.\\n\\n    Your models should also subclass this class.\\n\\n    Modules can also contain other Modules, allowing to nest them in\\n    a tree structure. You can assign the submodules as regular attributes::\\n\\n        import torch.nn as nn\\n        import torch.nn.functional as F\\n\\n        class Model(nn.Module):\\n            def __init__(self) -> None:\\n                super().__init__()\\n                self.conv1 = nn.Conv2d(1, 20, 5)\\n                self.conv2 = nn.Conv2d(20, 20, 5)\\n\\n            def forward(self, x):\\n                x = F.relu(self.conv1(x))\\n                return F.relu(self.conv2(x))\\n\\n    Submodules assigned in this way will be registered, and will have their\\n    parameters converted too when you call :meth:`to`, etc.\\n\\n    .. note::\\n        As per the example above, an ``__init__()`` call to the parent class\\n        must be made before assignment on the child.\\n\\n    :ivar training: Boolean represents whether this module is in training or\\n                    evaluation mode.\\n    :vartype training: bool\\n    \\\"\\\"\\\"\\n\\n    dump_patches: bool = False\\n\\n    _version: int = 1\\n    r\\\"\\\"\\\"This allows better BC support for :meth:`load_state_dict`. In\\n    :meth:`state_dict`, the version number will be saved as in the attribute\\n    `_metadata` of the returned state dict, and thus pickled. `_metadata` is a\\n    dictionary with keys that follow the naming convention of state dict. See\\n    ``_load_from_state_dict`` on how to use this information in loading.\\n\\n    If new parameters/buffers are added/removed from a module, this number shall\\n    be bumped, and the module's `_load_from_state_dict` method can compare the\\n    version number and do appropriate changes if the state dict is from before\\n    the change.\\\"\\\"\\\"\\n\\n    training: bool\\n    _parameters: Dict[str, Optional[Parameter]]\\n    _buffers: Dict[str, Optional[Tensor]]\\n    _non_persistent_buffers_set: Set[str]\\n    _backward_pre_hooks: Dict[int, Callable]\\n    _backward_hooks: Dict[int, Callable]\\n    _is_full_backward_hook: Optional[bool]\\n    _forward_hooks: Dict[int, Callable]\\n    # Marks whether the corresponding _forward_hooks accept kwargs or not.\\n    # As JIT does not support Set[int], this dict is used as a set, where all\\n    # hooks represented in this dict accept kwargs.\\n    _forward_hooks_with_kwargs: Dict[int, bool]\\n    # forward hooks that should always be called even if an exception is raised\\n    _forward_hooks_always_called: Dict[int, bool]\\n    _forward_pre_hooks: Dict[int, Callable]\\n    # Marks whether the corresponding _forward_hooks accept kwargs or not.\\n    # As JIT does not support Set[int], this dict is used as a set, where all\\n    # hooks represented in this dict accept kwargs.\\n    _forward_pre_hooks_with_kwargs: Dict[int, bool]\\n    _state_dict_hooks: Dict[int, Callable]\\n    _load_state_dict_pre_hooks: Dict[int, Callable]\\n    _state_dict_pre_hooks: Dict[int, Callable]\\n    _load_state_dict_post_hooks: Dict[int, Callable]\\n    _modules: Dict[str, Optional[\\\"Module\\\"]]\\n    call_super_init: bool = False\\n    _compiled_call_impl: Optional[Callable] = None\\n\\n    def __init__(self, *args, **kwargs) -> None:\\n        \\\"\\\"\\\"Initialize internal Module state, shared by both nn.Module and ScriptModule.\\\"\\\"\\\"\\n        torch._C._log_api_usage_once(\\\"python.nn_module\\\")\\n\\n        # Backward compatibility: no args used to be allowed when call_super_init=False\\n        if self.call_super_init is False and bool(kwargs):\\n            raise TypeError(\\n                f\\\"{type(self).__name__}.__init__() got an unexpected keyword argument '{next(iter(kwargs))}'\\\"\\n                \\\"\\\"\\n            )\\n\\n        if self.call_super_init is False and bool(args):\\n            raise TypeError(\\n                f\\\"{type(self).__name__}.__init__() takes 1 positional argument but {len(args) + 1} were\\\"\\n                \\\" given\\\"\\n            )\\n\\n        \\\"\\\"\\\"\\n        Calls super().__setattr__('a', a) instead of the typical self.a = a\\n        to avoid Module.__setattr__ overhead. Module's __setattr__ has special\\n        handling for parameters, submodules, and buffers but simply calls into\\n        super().__setattr__ for all other attributes.\\n        \\\"\\\"\\\"\\n        super().__setattr__(\\\"training\\\", True)\\n        super().__setattr__(\\\"_parameters\\\", {})\\n        super().__setattr__(\\\"_buffers\\\", {})\\n        super().__setattr__(\\\"_non_persistent_buffers_set\\\", set())\\n        super().__setattr__(\\\"_backward_pre_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_backward_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_is_full_backward_hook\\\", None)\\n        super().__setattr__(\\\"_forward_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_forward_hooks_with_kwargs\\\", OrderedDict())\\n        super().__setattr__(\\\"_forward_hooks_always_called\\\", OrderedDict())\\n        super().__setattr__(\\\"_forward_pre_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_forward_pre_hooks_with_kwargs\\\", OrderedDict())\\n        super().__setattr__(\\\"_state_dict_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_state_dict_pre_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_load_state_dict_pre_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_load_state_dict_post_hooks\\\", OrderedDict())\\n        super().__setattr__(\\\"_modules\\\", {})\\n\\n        if self.call_super_init:\\n            super().__init__(*args, **kwargs)\\n\\n    forward: Callable[..., Any] = _forward_unimplemented\\n\\n    def register_buffer(\\n        self, name: str, tensor: Optional[Tensor], persistent: bool = True\\n    ) -> None:\\n        r\\\"\\\"\\\"Add a buffer to the module.\\n\\n        This is typically used to register a buffer that should not to be\\n        considered a model parameter. For example, BatchNorm's ``running_mean``\\n        is not a parameter, but is part of the module's state. Buffers, by\\n        default, are persistent and will be saved alongside parameters. This\\n        behavior can be changed by setting :attr:`persistent` to ``False``. The\\n        only difference between a persistent buffer and a non-persistent buffer\\n        is that the latter will not be a part of this module's\\n        :attr:`state_dict`.\\n\\n        Buffers can be accessed as attributes using given names.\\n\\n        Args:\\n            name (str): name of the buffer. The buffer can be accessed\\n                from this module using the given name\\n            tensor (Tensor or None): buffer to be registered. If ``None``, then operations\\n                that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,\\n                the buffer is **not** included in the module's :attr:`state_dict`.\\n            persistent (bool): whether the buffer is part of this module's\\n                :attr:`state_dict`.\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> self.register_buffer('running_mean', torch.zeros(num_features))\\n\\n        \\\"\\\"\\\"\\n        if persistent is False and isinstance(self, torch.jit.ScriptModule):\\n            raise RuntimeError(\\\"ScriptModule does not support non-persistent buffers\\\")\\n\\n        if \\\"_buffers\\\" not in self.__dict__:\\n            raise AttributeError(\\\"cannot assign buffer before Module.__init__() call\\\")\\n        elif not isinstance(name, str):\\n            raise TypeError(\\n                f\\\"buffer name should be a string. Got {torch.typename(name)}\\\"\\n            )\\n        elif \\\".\\\" in name:\\n            raise KeyError('buffer name can\\\\'t contain \\\".\\\"')\\n        elif name == \\\"\\\":\\n            raise KeyError('buffer name can\\\\'t be empty string \\\"\\\"')\\n        elif hasattr(self, name) and name not in self._buffers:\\n            raise KeyError(f\\\"attribute '{name}' already exists\\\")\\n        elif tensor is not None and not isinstance(tensor, torch.Tensor):\\n            raise TypeError(\\n                f\\\"cannot assign '{torch.typename(tensor)}' object to buffer '{name}' \\\"\\n                \\\"(torch Tensor or None required)\\\"\\n            )\\n        else:\\n            for hook in _global_buffer_registration_hooks.values():\\n                output = hook(self, name, tensor)\\n                if output is not None:\\n                    tensor = output\\n            self._buffers[name] = tensor\\n            if persistent:\\n                self._non_persistent_buffers_set.discard(name)\\n            else:\\n                self._non_persistent_buffers_set.add(name)\\n\\n    def register_parameter(self, name: str, param: Optional[Parameter]) -> None:\\n        r\\\"\\\"\\\"Add a parameter to the module.\\n\\n        The parameter can be accessed as an attribute using given name.\\n\\n        Args:\\n            name (str): name of the parameter. The parameter can be accessed\\n                from this module using the given name\\n            param (Parameter or None): parameter to be added to the module. If\\n                ``None``, then operations that run on parameters, such as :attr:`cuda`,\\n                are ignored. If ``None``, the parameter is **not** included in the\\n                module's :attr:`state_dict`.\\n        \\\"\\\"\\\"\\n        if \\\"_parameters\\\" not in self.__dict__:\\n            raise AttributeError(\\n                \\\"cannot assign parameter before Module.__init__() call\\\"\\n            )\\n\\n        elif not isinstance(name, str):\\n            raise TypeError(\\n                f\\\"parameter name should be a string. Got {torch.typename(name)}\\\"\\n            )\\n        elif \\\".\\\" in name:\\n            raise KeyError('parameter name can\\\\'t contain \\\".\\\"')\\n        elif name == \\\"\\\":\\n            raise KeyError('parameter name can\\\\'t be empty string \\\"\\\"')\\n        elif hasattr(self, name) and name not in self._parameters:\\n            raise KeyError(f\\\"attribute '{name}' already exists\\\")\\n\\n        if param is None:\\n            self._parameters[name] = None\\n        elif not isinstance(param, Parameter):\\n            raise TypeError(\\n                f\\\"cannot assign '{torch.typename(param)}' object to parameter '{name}' \\\"\\n                \\\"(torch.nn.Parameter or None required)\\\"\\n            )\\n        elif param.grad_fn:\\n            raise ValueError(\\n                f\\\"Cannot assign non-leaf Tensor to parameter '{name}'. Model \\\"\\n                f\\\"parameters must be created explicitly. To express '{name}' \\\"\\n                \\\"as a function of another Tensor, compute the value in \\\"\\n                \\\"the forward() method.\\\"\\n            )\\n        else:\\n            for hook in _global_parameter_registration_hooks.values():\\n                output = hook(self, name, param)\\n                if output is not None:\\n                    param = output\\n            self._parameters[name] = param\\n\\n    def add_module(self, name: str, module: Optional[\\\"Module\\\"]) -> None:\\n        r\\\"\\\"\\\"Add a child module to the current module.\\n\\n        The module can be accessed as an attribute using the given name.\\n\\n        Args:\\n            name (str): name of the child module. The child module can be\\n                accessed from this module using the given name\\n            module (Module): child module to be added to the module.\\n        \\\"\\\"\\\"\\n        if not isinstance(module, Module) and module is not None:\\n            raise TypeError(f\\\"{torch.typename(module)} is not a Module subclass\\\")\\n        elif not isinstance(name, str):\\n            raise TypeError(\\n                f\\\"module name should be a string. Got {torch.typename(name)}\\\"\\n            )\\n        elif hasattr(self, name) and name not in self._modules:\\n            raise KeyError(f\\\"attribute '{name}' already exists\\\")\\n        elif \\\".\\\" in name:\\n            raise KeyError(f'module name can\\\\'t contain \\\".\\\", got: {name}')\\n        elif name == \\\"\\\":\\n            raise KeyError('module name can\\\\'t be empty string \\\"\\\"')\\n        for hook in _global_module_registration_hooks.values():\\n            output = hook(self, name, module)\\n            if output is not None:\\n                module = output\\n        self._modules[name] = module\\n\\n    def register_module(self, name: str, module: Optional[\\\"Module\\\"]) -> None:\\n        r\\\"\\\"\\\"Alias for :func:`add_module`.\\\"\\\"\\\"\\n        self.add_module(name, module)\\n\\n    def get_submodule(self, target: str) -> \\\"Module\\\":\\n        \\\"\\\"\\\"Return the submodule given by ``target`` if it exists, otherwise throw an error.\\n\\n        For example, let's say you have an ``nn.Module`` ``A`` that\\n        looks like this:\\n\\n        .. code-block:: text\\n\\n            A(\\n                (net_b): Module(\\n                    (net_c): Module(\\n                        (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))\\n                    )\\n                    (linear): Linear(in_features=100, out_features=200, bias=True)\\n                )\\n            )\\n\\n        (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested\\n        submodule ``net_b``, which itself has two submodules ``net_c``\\n        and ``linear``. ``net_c`` then has a submodule ``conv``.)\\n\\n        To check whether or not we have the ``linear`` submodule, we\\n        would call ``get_submodule(\\\"net_b.linear\\\")``. To check whether\\n        we have the ``conv`` submodule, we would call\\n        ``get_submodule(\\\"net_b.net_c.conv\\\")``.\\n\\n        The runtime of ``get_submodule`` is bounded by the degree\\n        of module nesting in ``target``. A query against\\n        ``named_modules`` achieves the same result, but it is O(N) in\\n        the number of transitive modules. So, for a simple check to see\\n        if some submodule exists, ``get_submodule`` should always be\\n        used.\\n\\n        Args:\\n            target: The fully-qualified string name of the submodule\\n                to look for. (See above example for how to specify a\\n                fully-qualified string.)\\n\\n        Returns:\\n            torch.nn.Module: The submodule referenced by ``target``\\n\\n        Raises:\\n            AttributeError: If the target string references an invalid\\n                path or resolves to something that is not an\\n                ``nn.Module``\\n        \\\"\\\"\\\"\\n        if target == \\\"\\\":\\n            return self\\n\\n        atoms: List[str] = target.split(\\\".\\\")\\n        mod: torch.nn.Module = self\\n\\n        for item in atoms:\\n            if not hasattr(mod, item):\\n                raise AttributeError(\\n                    mod._get_name() + \\\" has no \\\" \\\"attribute `\\\" + item + \\\"`\\\"\\n                )\\n\\n            mod = getattr(mod, item)\\n\\n            if not isinstance(mod, torch.nn.Module):\\n                raise AttributeError(\\\"`\\\" + item + \\\"` is not \\\" \\\"an nn.Module\\\")\\n\\n        return mod\\n\\n    def set_submodule(self, target: str, module: \\\"Module\\\") -> None:\\n        \\\"\\\"\\\"\\n        Set the submodule given by ``target`` if it exists, otherwise throw an error.\\n\\n        For example, let's say you have an ``nn.Module`` ``A`` that\\n        looks like this:\\n\\n        .. code-block:: text\\n\\n            A(\\n                (net_b): Module(\\n                    (net_c): Module(\\n                        (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))\\n                    )\\n                    (linear): Linear(in_features=100, out_features=200, bias=True)\\n                )\\n            )\\n\\n        (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested\\n        submodule ``net_b``, which itself has two submodules ``net_c``\\n        and ``linear``. ``net_c`` then has a submodule ``conv``.)\\n\\n        To overide the ``Conv2d`` with a new submodule ``Linear``, you\\n        would call\\n        ``set_submodule(\\\"net_b.net_c.conv\\\", nn.Linear(33, 16))``.\\n\\n        Args:\\n            target: The fully-qualified string name of the submodule\\n                to look for. (See above example for how to specify a\\n                fully-qualified string.)\\n            module: The module to set the submodule to.\\n\\n        Raises:\\n            ValueError: If the target string is empty\\n            AttributeError: If the target string references an invalid\\n                path or resolves to something that is not an\\n                ``nn.Module``\\n        \\\"\\\"\\\"\\n        if target == \\\"\\\":\\n            raise ValueError(\\\"Cannot set the submodule without a target name!\\\")\\n\\n        atoms: List[str] = target.split(\\\".\\\")\\n        name = atoms.pop(-1)\\n        mod: torch.nn.Module = self\\n\\n        for item in atoms:\\n            if not hasattr(mod, item):\\n                raise AttributeError(\\n                    mod._get_name() + \\\" has no attribute `\\\" + item + \\\"`\\\"\\n                )\\n\\n            mod = getattr(mod, item)\\n\\n            # Use isinstance instead of type here to also handle subclass of nn.Module\\n            if not isinstance(mod, torch.nn.Module):\\n                raise AttributeError(\\\"`\\\" + item + \\\"` is not an nn.Module\\\")\\n\\n        setattr(mod, name, module)\\n\\n    def get_parameter(self, target: str) -> \\\"Parameter\\\":\\n        \\\"\\\"\\\"Return the parameter given by ``target`` if it exists, otherwise throw an error.\\n\\n        See the docstring for ``get_submodule`` for a more detailed\\n        explanation of this method's functionality as well as how to\\n        correctly specify ``target``.\\n\\n        Args:\\n            target: The fully-qualified string name of the Parameter\\n                to look for. (See ``get_submodule`` for how to specify a\\n                fully-qualified string.)\\n\\n        Returns:\\n            torch.nn.Parameter: The Parameter referenced by ``target``\\n\\n        Raises:\\n            AttributeError: If the target string references an invalid\\n                path or resolves to something that is not an\\n                ``nn.Parameter``\\n        \\\"\\\"\\\"\\n        module_path, _, param_name = target.rpartition(\\\".\\\")\\n\\n        mod: torch.nn.Module = self.get_submodule(module_path)\\n\\n        if not hasattr(mod, param_name):\\n            raise AttributeError(\\n                mod._get_name() + \\\" has no attribute `\\\" + param_name + \\\"`\\\"\\n            )\\n\\n        param: torch.nn.Parameter = getattr(mod, param_name)\\n\\n        if not isinstance(param, torch.nn.Parameter):\\n            raise AttributeError(\\\"`\\\" + param_name + \\\"` is not an \\\" \\\"nn.Parameter\\\")\\n\\n        return param\\n\\n    def get_buffer(self, target: str) -> \\\"Tensor\\\":\\n        \\\"\\\"\\\"Return the buffer given by ``target`` if it exists, otherwise throw an error.\\n\\n        See the docstring for ``get_submodule`` for a more detailed\\n        explanation of this method's functionality as well as how to\\n        correctly specify ``target``.\\n\\n        Args:\\n            target: The fully-qualified string name of the buffer\\n                to look for. (See ``get_submodule`` for how to specify a\\n                fully-qualified string.)\\n\\n        Returns:\\n            torch.Tensor: The buffer referenced by ``target``\\n\\n        Raises:\\n            AttributeError: If the target string references an invalid\\n                path or resolves to something that is not a\\n                buffer\\n        \\\"\\\"\\\"\\n        module_path, _, buffer_name = target.rpartition(\\\".\\\")\\n\\n        mod: torch.nn.Module = self.get_submodule(module_path)\\n\\n        if not hasattr(mod, buffer_name):\\n            raise AttributeError(\\n                mod._get_name() + \\\" has no attribute `\\\" + buffer_name + \\\"`\\\"\\n            )\\n\\n        buffer: torch.Tensor = getattr(mod, buffer_name)\\n\\n        if buffer_name not in mod._buffers:\\n            raise AttributeError(\\\"`\\\" + buffer_name + \\\"` is not a buffer\\\")\\n\\n        return buffer\\n\\n    def get_extra_state(self) -> Any:\\n        \\\"\\\"\\\"Return any extra state to include in the module's state_dict.\\n\\n        Implement this and a corresponding :func:`set_extra_state` for your module\\n        if you need to store extra state. This function is called when building the\\n        module's `state_dict()`.\\n\\n        Note that extra state should be picklable to ensure working serialization\\n        of the state_dict. We only provide provide backwards compatibility guarantees\\n        for serializing Tensors; other objects may break backwards compatibility if\\n        their serialized pickled form changes.\\n\\n        Returns:\\n            object: Any extra state to store in the module's state_dict\\n        \\\"\\\"\\\"\\n        raise RuntimeError(\\n            \\\"Reached a code path in Module.get_extra_state() that should never be called. \\\"\\n            \\\"Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml \\\"\\n            \\\"to report this bug.\\\"\\n        )\\n\\n    def set_extra_state(self, state: Any) -> None:\\n        \\\"\\\"\\\"Set extra state contained in the loaded `state_dict`.\\n\\n        This function is called from :func:`load_state_dict` to handle any extra state\\n        found within the `state_dict`. Implement this function and a corresponding\\n        :func:`get_extra_state` for your module if you need to store extra state within its\\n        `state_dict`.\\n\\n        Args:\\n            state (dict): Extra state from the `state_dict`\\n        \\\"\\\"\\\"\\n        raise RuntimeError(\\n            \\\"Reached a code path in Module.set_extra_state() that should never be called. \\\"\\n            \\\"Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml \\\"\\n            \\\"to report this bug.\\\"\\n        )\\n\\n    def _apply(self, fn, recurse=True):\\n        if recurse:\\n            for module in self.children():\\n                module._apply(fn)\\n\\n        def compute_should_use_set_data(tensor, tensor_applied):\\n            if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):\\n                # If the new tensor has compatible tensor type as the existing tensor,\\n                # the current behavior is to change the tensor in-place using `.data =`,\\n                # and the future behavior is to overwrite the existing tensor. However,\\n                # changing the current behavior is a BC-breaking change, and we want it\\n                # to happen in future releases. So for now we introduce the\\n                # `torch.__future__.get_overwrite_module_params_on_conversion()`\\n                # global flag to let the user control whether they want the future\\n                # behavior of overwriting the existing tensor or not.\\n                return not torch.__future__.get_overwrite_module_params_on_conversion()\\n            else:\\n                return False\\n\\n        should_use_swap_tensors = (\\n            torch.__future__.get_swap_module_params_on_conversion()\\n        )\\n\\n        for key, param in self._parameters.items():\\n            if param is None:\\n                continue\\n            # Tensors stored in modules are graph leaves, and we don't want to\\n            # track autograd history of `param_applied`, so we have to use\\n            # `with torch.no_grad():`\\n            with torch.no_grad():\\n                param_applied = fn(param)\\n            p_should_use_set_data = compute_should_use_set_data(param, param_applied)\\n\\n            # subclasses may have multiple child tensors so we need to use swap_tensors\\n            p_should_use_swap_tensors = (\\n                should_use_swap_tensors or is_traceable_wrapper_subclass(param_applied)\\n            )\\n\\n            param_grad = param.grad\\n            if p_should_use_swap_tensors:\\n                try:\\n                    if param_grad is not None:\\n                        # Accessing param.grad makes its at::Tensor's use_count 2, which will prevent swapping.\\n                        # Decrement use count of the gradient by setting to None\\n                        param.grad = None\\n                    param_applied = torch.nn.Parameter(\\n                        param_applied, requires_grad=param.requires_grad\\n                    )\\n                    torch.utils.swap_tensors(param, param_applied)\\n                except Exception as e:\\n                    if param_grad is not None:\\n                        param.grad = param_grad\\n                    raise RuntimeError(\\n                        f\\\"_apply(): Couldn't swap {self._get_name()}.{key}\\\"\\n                    ) from e\\n                out_param = param\\n            elif p_should_use_set_data:\\n                param.data = param_applied\\n                out_param = param\\n            else:\\n                assert isinstance(param, Parameter)\\n                assert param.is_leaf\\n                out_param = Parameter(param_applied, param.requires_grad)\\n                self._parameters[key] = out_param\\n\\n            if param_grad is not None:\\n                with torch.no_grad():\\n                    grad_applied = fn(param_grad)\\n                g_should_use_set_data = compute_should_use_set_data(\\n                    param_grad, grad_applied\\n                )\\n                if p_should_use_swap_tensors:\\n                    grad_applied.requires_grad_(param_grad.requires_grad)\\n                    try:\\n                        torch.utils.swap_tensors(param_grad, grad_applied)\\n                    except Exception as e:\\n                        raise RuntimeError(\\n                            f\\\"_apply(): Couldn't swap {self._get_name()}.{key}.grad\\\"\\n                        ) from e\\n                    out_param.grad = param_grad\\n                elif g_should_use_set_data:\\n                    assert out_param.grad is not None\\n                    out_param.grad.data = grad_applied\\n                else:\\n                    assert param_grad.is_leaf\\n                    out_param.grad = grad_applied.requires_grad_(\\n                        param_grad.requires_grad\\n                    )\\n\\n        for key, buf in self._buffers.items():\\n            if buf is not None:\\n                self._buffers[key] = fn(buf)\\n\\n        return self\\n\\n    def apply(self: T, fn: Callable[[\\\"Module\\\"], None]) -> T:\\n        r\\\"\\\"\\\"Apply ``fn`` recursively to every submodule (as returned by ``.children()``) as well as self.\\n\\n        Typical use includes initializing the parameters of a model\\n        (see also :ref:`nn-init-doc`).\\n\\n        Args:\\n            fn (:class:`Module` -> None): function to be applied to each submodule\\n\\n        Returns:\\n            Module: self\\n\\n        Example::\\n\\n            >>> @torch.no_grad()\\n            >>> def init_weights(m):\\n            >>>     print(m)\\n            >>>     if type(m) == nn.Linear:\\n            >>>         m.weight.fill_(1.0)\\n            >>>         print(m.weight)\\n            >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))\\n            >>> net.apply(init_weights)\\n            Linear(in_features=2, out_features=2, bias=True)\\n            Parameter containing:\\n            tensor([[1., 1.],\\n                    [1., 1.]], requires_grad=True)\\n            Linear(in_features=2, out_features=2, bias=True)\\n            Parameter containing:\\n            tensor([[1., 1.],\\n                    [1., 1.]], requires_grad=True)\\n            Sequential(\\n              (0): Linear(in_features=2, out_features=2, bias=True)\\n              (1): Linear(in_features=2, out_features=2, bias=True)\\n            )\\n\\n        \\\"\\\"\\\"\\n        for module in self.children():\\n            module.apply(fn)\\n        fn(self)\\n        return self\\n\\n    def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:\\n        r\\\"\\\"\\\"Move all model parameters and buffers to the GPU.\\n\\n        This also makes associated parameters and buffers different objects. So\\n        it should be called before constructing optimizer if the module will\\n        live on GPU while being optimized.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Args:\\n            device (int, optional): if specified, all parameters will be\\n                copied to that device\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.cuda(device))\\n\\n    def ipu(self: T, device: Optional[Union[int, device]] = None) -> T:\\n        r\\\"\\\"\\\"Move all model parameters and buffers to the IPU.\\n\\n        This also makes associated parameters and buffers different objects. So\\n        it should be called before constructing optimizer if the module will\\n        live on IPU while being optimized.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Arguments:\\n            device (int, optional): if specified, all parameters will be\\n                copied to that device\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.ipu(device))\\n\\n    def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:\\n        r\\\"\\\"\\\"Move all model parameters and buffers to the XPU.\\n\\n        This also makes associated parameters and buffers different objects. So\\n        it should be called before constructing optimizer if the module will\\n        live on XPU while being optimized.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Arguments:\\n            device (int, optional): if specified, all parameters will be\\n                copied to that device\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.xpu(device))\\n\\n    def mtia(self: T, device: Optional[Union[int, device]] = None) -> T:\\n        r\\\"\\\"\\\"Move all model parameters and buffers to the MTIA.\\n\\n        This also makes associated parameters and buffers different objects. So\\n        it should be called before constructing optimizer if the module will\\n        live on MTIA while being optimized.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Arguments:\\n            device (int, optional): if specified, all parameters will be\\n                copied to that device\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.mtia(device))\\n\\n    def cpu(self: T) -> T:\\n        r\\\"\\\"\\\"Move all model parameters and buffers to the CPU.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.cpu())\\n\\n    def type(self: T, dst_type: Union[dtype, str]) -> T:\\n        r\\\"\\\"\\\"Casts all parameters and buffers to :attr:`dst_type`.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Args:\\n            dst_type (type or string): the desired type\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.type(dst_type))\\n\\n    def float(self: T) -> T:\\n        r\\\"\\\"\\\"Casts all floating point parameters and buffers to ``float`` datatype.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.float() if t.is_floating_point() else t)\\n\\n    def double(self: T) -> T:\\n        r\\\"\\\"\\\"Casts all floating point parameters and buffers to ``double`` datatype.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.double() if t.is_floating_point() else t)\\n\\n    def half(self: T) -> T:\\n        r\\\"\\\"\\\"Casts all floating point parameters and buffers to ``half`` datatype.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.half() if t.is_floating_point() else t)\\n\\n    def bfloat16(self: T) -> T:\\n        r\\\"\\\"\\\"Casts all floating point parameters and buffers to ``bfloat16`` datatype.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)\\n\\n    def to_empty(\\n        self: T, *, device: Optional[DeviceLikeType], recurse: bool = True\\n    ) -> T:\\n        r\\\"\\\"\\\"Move the parameters and buffers to the specified device without copying storage.\\n\\n        Args:\\n            device (:class:`torch.device`): The desired device of the parameters\\n                and buffers in this module.\\n            recurse (bool): Whether parameters and buffers of submodules should\\n                be recursively moved to the specified device.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self._apply(\\n            lambda t: torch.empty_like(t, device=device), recurse=recurse\\n        )\\n\\n    @overload\\n    def to(\\n        self,\\n        device: Optional[DeviceLikeType] = ...,\\n        dtype: Optional[dtype] = ...,\\n        non_blocking: bool = ...,\\n    ) -> Self:\\n        ...\\n\\n    @overload\\n    def to(self, dtype: dtype, non_blocking: bool = ...) -> Self:\\n        ...\\n\\n    @overload\\n    def to(self, tensor: Tensor, non_blocking: bool = ...) -> Self:\\n        ...\\n\\n    def to(self, *args, **kwargs):\\n        r\\\"\\\"\\\"Move and/or cast the parameters and buffers.\\n\\n        This can be called as\\n\\n        .. function:: to(device=None, dtype=None, non_blocking=False)\\n           :noindex:\\n\\n        .. function:: to(dtype, non_blocking=False)\\n           :noindex:\\n\\n        .. function:: to(tensor, non_blocking=False)\\n           :noindex:\\n\\n        .. function:: to(memory_format=torch.channels_last)\\n           :noindex:\\n\\n        Its signature is similar to :meth:`torch.Tensor.to`, but only accepts\\n        floating point or complex :attr:`dtype`\\\\ s. In addition, this method will\\n        only cast the floating point or complex parameters and buffers to :attr:`dtype`\\n        (if given). The integral parameters and buffers will be moved\\n        :attr:`device`, if that is given, but with dtypes unchanged. When\\n        :attr:`non_blocking` is set, it tries to convert/move asynchronously\\n        with respect to the host if possible, e.g., moving CPU Tensors with\\n        pinned memory to CUDA devices.\\n\\n        See below for examples.\\n\\n        .. note::\\n            This method modifies the module in-place.\\n\\n        Args:\\n            device (:class:`torch.device`): the desired device of the parameters\\n                and buffers in this module\\n            dtype (:class:`torch.dtype`): the desired floating point or complex dtype of\\n                the parameters and buffers in this module\\n            tensor (torch.Tensor): Tensor whose dtype and device are the desired\\n                dtype and device for all parameters and buffers in this module\\n            memory_format (:class:`torch.memory_format`): the desired memory\\n                format for 4D parameters and buffers in this module (keyword\\n                only argument)\\n\\n        Returns:\\n            Module: self\\n\\n        Examples::\\n\\n            >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n            >>> linear = nn.Linear(2, 2)\\n            >>> linear.weight\\n            Parameter containing:\\n            tensor([[ 0.1913, -0.3420],\\n                    [-0.5113, -0.2325]])\\n            >>> linear.to(torch.double)\\n            Linear(in_features=2, out_features=2, bias=True)\\n            >>> linear.weight\\n            Parameter containing:\\n            tensor([[ 0.1913, -0.3420],\\n                    [-0.5113, -0.2325]], dtype=torch.float64)\\n            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)\\n            >>> gpu1 = torch.device(\\\"cuda:1\\\")\\n            >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)\\n            Linear(in_features=2, out_features=2, bias=True)\\n            >>> linear.weight\\n            Parameter containing:\\n            tensor([[ 0.1914, -0.3420],\\n                    [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')\\n            >>> cpu = torch.device(\\\"cpu\\\")\\n            >>> linear.to(cpu)\\n            Linear(in_features=2, out_features=2, bias=True)\\n            >>> linear.weight\\n            Parameter containing:\\n            tensor([[ 0.1914, -0.3420],\\n                    [-0.5112, -0.2324]], dtype=torch.float16)\\n\\n            >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)\\n            >>> linear.weight\\n            Parameter containing:\\n            tensor([[ 0.3741+0.j,  0.2382+0.j],\\n                    [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)\\n            >>> linear(torch.ones(3, 2, dtype=torch.cdouble))\\n            tensor([[0.6122+0.j, 0.1150+0.j],\\n                    [0.6122+0.j, 0.1150+0.j],\\n                    [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)\\n\\n        \\\"\\\"\\\"\\n        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(\\n            *args, **kwargs\\n        )\\n\\n        if dtype is not None:\\n            if not (dtype.is_floating_point or dtype.is_complex):\\n                raise TypeError(\\n                    \\\"nn.Module.to only accepts floating point or complex \\\"\\n                    f\\\"dtypes, but got desired dtype={dtype}\\\"\\n                )\\n            if dtype.is_complex:\\n                warnings.warn(\\n                    \\\"Complex modules are a new feature under active development whose design may change, \\\"\\n                    \\\"and some modules might not work as expected when using complex tensors as parameters or buffers. \\\"\\n                    \\\"Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml \\\"\\n                    \\\"if a complex module does not work as expected.\\\"\\n                )\\n\\n        def convert(t):\\n            try:\\n                if convert_to_format is not None and t.dim() in (4, 5):\\n                    return t.to(\\n                        device,\\n                        dtype if t.is_floating_point() or t.is_complex() else None,\\n                        non_blocking,\\n                        memory_format=convert_to_format,\\n                    )\\n                return t.to(\\n                    device,\\n                    dtype if t.is_floating_point() or t.is_complex() else None,\\n                    non_blocking,\\n                )\\n            except NotImplementedError as e:\\n                if str(e) == \\\"Cannot copy out of meta tensor; no data!\\\":\\n                    raise NotImplementedError(\\n                        f\\\"{e} Please use torch.nn.Module.to_empty() instead of torch.nn.Module.to() \\\"\\n                        f\\\"when moving module from meta to a different device.\\\"\\n                    ) from None\\n                else:\\n                    raise\\n\\n        return self._apply(convert)\\n\\n    def register_full_backward_pre_hook(\\n        self,\\n        hook: Callable[[\\\"Module\\\", _grad_t], Union[None, _grad_t]],\\n        prepend: bool = False,\\n    ) -> RemovableHandle:\\n        r\\\"\\\"\\\"Register a backward pre-hook on the module.\\n\\n        The hook will be called every time the gradients for the module are computed.\\n        The hook should have the following signature::\\n\\n            hook(module, grad_output) -> tuple[Tensor] or None\\n\\n        The :attr:`grad_output` is a tuple. The hook should\\n        not modify its arguments, but it can optionally return a new gradient with\\n        respect to the output that will be used in place of :attr:`grad_output` in\\n        subsequent computations. Entries in :attr:`grad_output` will be ``None`` for\\n        all non-Tensor arguments.\\n\\n        For technical reasons, when this hook is applied to a Module, its forward function will\\n        receive a view of each Tensor passed to the Module. Similarly the caller will receive a view\\n        of each Tensor returned by the Module's forward function.\\n\\n        .. warning ::\\n            Modifying inputs inplace is not allowed when using backward hooks and\\n            will raise an error.\\n\\n        Args:\\n            hook (Callable): The user-defined hook to be registered.\\n            prepend (bool): If true, the provided ``hook`` will be fired before\\n                all existing ``backward_pre`` hooks on this\\n                :class:`torch.nn.modules.Module`. Otherwise, the provided\\n                ``hook`` will be fired after all existing ``backward_pre`` hooks\\n                on this :class:`torch.nn.modules.Module`. Note that global\\n                ``backward_pre`` hooks registered with\\n                :func:`register_module_full_backward_pre_hook` will fire before\\n                all hooks registered by this method.\\n\\n        Returns:\\n            :class:`torch.utils.hooks.RemovableHandle`:\\n                a handle that can be used to remove the added hook by calling\\n                ``handle.remove()``\\n\\n        \\\"\\\"\\\"\\n        handle = RemovableHandle(self._backward_pre_hooks)\\n        self._backward_pre_hooks[handle.id] = hook\\n        if prepend:\\n            self._backward_pre_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]\\n        return handle\\n\\n    def register_backward_hook(\\n        self, hook: Callable[[\\\"Module\\\", _grad_t, _grad_t], Union[None, _grad_t]]\\n    ) -> RemovableHandle:\\n        r\\\"\\\"\\\"Register a backward hook on the module.\\n\\n        This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and\\n        the behavior of this function will change in future versions.\\n\\n        Returns:\\n            :class:`torch.utils.hooks.RemovableHandle`:\\n                a handle that can be used to remove the added hook by calling\\n                ``handle.remove()``\\n\\n        \\\"\\\"\\\"\\n        if self._is_full_backward_hook is True:\\n            raise RuntimeError(\\n                \\\"Cannot use both regular backward hooks and full backward hooks on a \\\"\\n                \\\"single Module. Please use only one of them.\\\"\\n            )\\n\\n        self._is_full_backward_hook = False\\n\\n        handle = RemovableHandle(self._backward_hooks)\\n        self._backward_hooks[handle.id] = hook\\n        return handle\\n\\n    def register_full_backward_hook(\\n        self,\\n        hook: Callable[[\\\"Module\\\", _grad_t, _grad_t], Union[None, _grad_t]],\\n        prepend: bool = False,\\n    ) -> RemovableHandle:\\n        r\\\"\\\"\\\"Register a backward hook on the module.\\n\\n        The hook will be called every time the gradients with respect to a module\\n        are computed, i.e. the hook will execute if and only if the gradients with\\n        respect to module outputs are computed. The hook should have the following\\n        signature::\\n\\n            hook(module, grad_input, grad_output) -> tuple(Tensor) or None\\n\\n        The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients\\n        with respect to the inputs and outputs respectively. The hook should\\n        not modify its arguments, but it can optionally return a new gradient with\\n        respect to the input that will be used in place of :attr:`grad_input` in\\n        subsequent computations. :attr:`grad_input` will only correspond to the inputs given\\n        as positional arguments and all kwarg arguments are ignored. Entries\\n        in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor\\n        arguments.\\n\\n        For technical reasons, when this hook is applied to a Module, its forward function will\\n        receive a view of each Tensor passed to the Module. Similarly the caller will receive a view\\n        of each Tensor returned by the Module's forward function.\\n\\n        .. warning ::\\n            Modifying inputs or outputs inplace is not allowed when using backward hooks and\\n            will raise an error.\\n\\n        Args:\\n            hook (Callable): The user-defined hook to be registered.\\n            prepend (bool): If true, the provided ``hook`` will be fired before\\n                all existing ``backward`` hooks on this\\n                :class:`torch.nn.modules.Module`. Otherwise, the provided\\n                ``hook`` will be fired after all existing ``backward`` hooks on\\n                this :class:`torch.nn.modules.Module`. Note that global\\n                ``backward`` hooks registered with\\n                :func:`register_module_full_backward_hook` will fire before\\n                all hooks registered by this method.\\n\\n        Returns:\\n            :class:`torch.utils.hooks.RemovableHandle`:\\n                a handle that can be used to remove the added hook by calling\\n                ``handle.remove()``\\n\\n        \\\"\\\"\\\"\\n        if self._is_full_backward_hook is False:\\n            raise RuntimeError(\\n                \\\"Cannot use both regular backward hooks and full backward hooks on a \\\"\\n                \\\"single Module. Please use only one of them.\\\"\\n            )\\n\\n        self._is_full_backward_hook = True\\n\\n        handle = RemovableHandle(self._backward_hooks)\\n        self._backward_hooks[handle.id] = hook\\n        if prepend:\\n            self._backward_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]\\n        return handle\\n\\n    def _get_backward_hooks(self):\\n        r\\\"\\\"\\\"Return the backward hooks for use in the call function.\\n\\n        It returns two lists, one with the full backward hooks and one with the non-full\\n        backward hooks.\\n        \\\"\\\"\\\"\\n        full_backward_hooks: List[Callable] = []\\n        if _global_is_full_backward_hook is True:\\n            full_backward_hooks += _global_backward_hooks.values()\\n        if self._is_full_backward_hook is True:\\n            full_backward_hooks += self._backward_hooks.values()\\n\\n        non_full_backward_hooks: List[Callable] = []\\n        if _global_is_full_backward_hook is False:\\n            non_full_backward_hooks += _global_backward_hooks.values()\\n        if self._is_full_backward_hook is False:\\n            non_full_backward_hooks += self._backward_hooks.values()\\n\\n        return full_backward_hooks, non_full_backward_hooks\\n\\n    def _get_backward_pre_hooks(self):\\n        backward_pre_hooks: List[Callable] = []\\n        backward_pre_hooks += _global_backward_pre_hooks.values()\\n        backward_pre_hooks += self._backward_pre_hooks.values()\\n\\n        return backward_pre_hooks\\n\\n    def _maybe_warn_non_full_backward_hook(self, inputs, result, grad_fn):\\n        if not isinstance(result, torch.Tensor):\\n            if not (\\n                isinstance(result, tuple)\\n                and all(isinstance(r, torch.Tensor) for r in result)\\n            ):\\n                warnings.warn(\\n                    \\\"Using non-full backward hooks on a Module that does not return a \\\"\\n                    \\\"single Tensor or a tuple of Tensors is deprecated and will be removed \\\"\\n                    \\\"in future versions. This hook will be missing some of the grad_output. \\\"\\n                    \\\"Please use register_full_backward_hook to get the documented behavior.\\\",\\n                    FutureWarning,\\n                    stacklevel=2,\\n                )\\n                return\\n        else:\\n            result = (result,)\\n\\n        if not isinstance(inputs, torch.Tensor):\\n            if not (\\n                isinstance(inputs, tuple)\\n                and all(isinstance(i, torch.Tensor) for i in inputs)\\n            ):\\n                warnings.warn(\\n                    \\\"Using non-full backward hooks on a Module that does not take as input a \\\"\\n                    \\\"single Tensor or a tuple of Tensors is deprecated and will be removed \\\"\\n                    \\\"in future versions. This hook will be missing some of the grad_input. \\\"\\n                    \\\"Please use register_full_backward_hook to get the documented behavior.\\\",\\n                    FutureWarning,\\n                    stacklevel=2,\\n                )\\n                return\\n        else:\\n            inputs = (inputs,)\\n\\n        # At this point we are sure that inputs and result are tuple of Tensors\\n        out_grad_fn = {r.grad_fn for r in result if r.grad_fn is not None}\\n        if len(out_grad_fn) == 0 or (\\n            len(out_grad_fn) == 1 and grad_fn not in out_grad_fn\\n        ):\\n            warnings.warn(\\n                \\\"Using a non-full backward hook when outputs are nested in python data structure \\\"\\n                \\\"is deprecated and will be removed in future versions. This hook will be missing \\\"\\n                \\\"some grad_output.\\\",\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n        elif len(out_grad_fn) > 1:\\n            warnings.warn(\\n                \\\"Using a non-full backward hook when outputs are generated by different autograd Nodes \\\"\\n                \\\"is deprecated and will be removed in future versions. This hook will be missing \\\"\\n                \\\"some grad_output. Please use register_full_backward_hook to get the documented behavior.\\\",\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n        else:\\n            # At this point the grad_output part of the hook will most likely be correct\\n            inputs_grad_fn = {i.grad_fn for i in inputs if i.grad_fn is not None}\\n\\n            next_functions = {n[0] for n in grad_fn.next_functions}\\n\\n            if inputs_grad_fn != next_functions:\\n                warnings.warn(\\n                    \\\"Using a non-full backward hook when the forward contains multiple autograd Nodes \\\"\\n                    \\\"is deprecated and will be removed in future versions. This hook will be missing \\\"\\n                    \\\"some grad_input. Please use register_full_backward_hook to get the documented \\\"\\n                    \\\"behavior.\\\",\\n                    FutureWarning,\\n                    stacklevel=2,\\n                )\\n\\n    def register_forward_pre_hook(\\n        self,\\n        hook: Union[\\n            Callable[[T, Tuple[Any, ...]], Optional[Any]],\\n            Callable[\\n                [T, Tuple[Any, ...], Dict[str, Any]],\\n                Optional[Tuple[Any, Dict[str, Any]]],\\n            ],\\n        ],\\n        *,\\n        prepend: bool = False,\\n        with_kwargs: bool = False,\\n    ) -> RemovableHandle:\\n        r\\\"\\\"\\\"Register a forward pre-hook on the module.\\n\\n        The hook will be called every time before :func:`forward` is invoked.\\n\\n\\n        If ``with_kwargs`` is false or not specified, the input contains only\\n        the positional arguments given to the module. Keyword arguments won't be\\n        passed to the hooks and only to the ``forward``. The hook can modify the\\n        input. User can either return a tuple or a single modified value in the\\n        hook. We will wrap the value into a tuple if a single value is returned\\n        (unless that value is already a tuple). The hook should have the\\n        following signature::\\n\\n            hook(module, args) -> None or modified input\\n\\n        If ``with_kwargs`` is true, the forward pre-hook will be passed the\\n        kwargs given to the forward function. And if the hook modifies the\\n        input, both the args and kwargs should be returned. The hook should have\\n        the following signature::\\n\\n            hook(module, args, kwargs) -> None or a tuple of modified input and kwargs\\n\\n        Args:\\n            hook (Callable): The user defined hook to be registered.\\n            prepend (bool): If true, the provided ``hook`` will be fired before\\n                all existing ``forward_pre`` hooks on this\\n                :class:`torch.nn.modules.Module`. Otherwise, the provided\\n                ``hook`` will be fired after all existing ``forward_pre`` hooks\\n                on this :class:`torch.nn.modules.Module`. Note that global\\n                ``forward_pre`` hooks registered with\\n                :func:`register_module_forward_pre_hook` will fire before all\\n                hooks registered by this method.\\n                Default: ``False``\\n            with_kwargs (bool): If true, the ``hook`` will be passed the kwargs\\n                given to the forward function.\\n                Default: ``False``\\n\\n        Returns:\\n            :class:`torch.utils.hooks.RemovableHandle`:\\n                a handle that can be used to remove the added hook by calling\\n                ``handle.remove()``\\n        \\\"\\\"\\\"\\n        handle = RemovableHandle(\\n            self._forward_pre_hooks, extra_dict=self._forward_pre_hooks_with_kwargs\\n        )\\n        self._forward_pre_hooks[handle.id] = hook\\n        if with_kwargs:\\n            self._forward_pre_hooks_with_kwargs[handle.id] = True\\n\\n        if prepend:\\n            self._forward_pre_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]\\n        return handle\\n\\n    def register_forward_hook(\\n        self,\\n        hook: Union[\\n            Callable[[T, Tuple[Any, ...], Any], Optional[Any]],\\n            Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],\\n        ],\\n        *,\\n        prepend: bool = False,\\n        with_kwargs: bool = False,\\n        always_call: bool = False,\\n    ) -> RemovableHandle:\\n        r\\\"\\\"\\\"Register a forward hook on the module.\\n\\n        The hook will be called every time after :func:`forward` has computed an output.\\n\\n        If ``with_kwargs`` is ``False`` or not specified, the input contains only\\n        the positional arguments given to the module. Keyword arguments won't be\\n        passed to the hooks and only to the ``forward``. The hook can modify the\\n        output. It can modify the input inplace but it will not have effect on\\n        forward since this is called after :func:`forward` is called. The hook\\n        should have the following signature::\\n\\n            hook(module, args, output) -> None or modified output\\n\\n        If ``with_kwargs`` is ``True``, the forward hook will be passed the\\n        ``kwargs`` given to the forward function and be expected to return the\\n        output possibly modified. The hook should have the following signature::\\n\\n            hook(module, args, kwargs, output) -> None or modified output\\n\\n        Args:\\n            hook (Callable): The user defined hook to be registered.\\n            prepend (bool): If ``True``, the provided ``hook`` will be fired\\n                before all existing ``forward`` hooks on this\\n                :class:`torch.nn.modules.Module`. Otherwise, the provided\\n                ``hook`` will be fired after all existing ``forward`` hooks on\\n                this :class:`torch.nn.modules.Module`. Note that global\\n                ``forward`` hooks registered with\\n                :func:`register_module_forward_hook` will fire before all hooks\\n                registered by this method.\\n                Default: ``False``\\n            with_kwargs (bool): If ``True``, the ``hook`` will be passed the\\n                kwargs given to the forward function.\\n                Default: ``False``\\n            always_call (bool): If ``True`` the ``hook`` will be run regardless of\\n                whether an exception is raised while calling the Module.\\n                Default: ``False``\\n\\n        Returns:\\n            :class:`torch.utils.hooks.RemovableHandle`:\\n                a handle that can be used to remove the added hook by calling\\n                ``handle.remove()``\\n        \\\"\\\"\\\"\\n        handle = RemovableHandle(\\n            self._forward_hooks,\\n            extra_dict=[\\n                self._forward_hooks_with_kwargs,\\n                self._forward_hooks_always_called,\\n            ],\\n        )\\n        self._forward_hooks[handle.id] = hook\\n        if with_kwargs:\\n            self._forward_hooks_with_kwargs[handle.id] = True\\n        if always_call:\\n            self._forward_hooks_always_called[handle.id] = True\\n        if prepend:\\n            self._forward_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]\\n        return handle\\n\\n    def _slow_forward(self, *input, **kwargs):\\n        tracing_state = torch._C._get_tracing_state()\\n        if not tracing_state or isinstance(self.forward, torch._C.ScriptMethod):\\n            return self.forward(*input, **kwargs)\\n        recording_scopes = torch.jit._trace._trace_module_map is not None\\n        if recording_scopes:\\n            # type ignore was added because at this point one knows that\\n            # torch.jit._trace._trace_module_map is not Optional and has type Dict[Any, Any]\\n            name = torch.jit._trace._trace_module_map[self] if self in torch.jit._trace._trace_module_map else None  # type: ignore[index, operator] # noqa: B950\\n            if name:\\n                tracing_state.push_scope(name)\\n            else:\\n                recording_scopes = False\\n        try:\\n            result = self.forward(*input, **kwargs)\\n        finally:\\n            if recording_scopes:\\n                tracing_state.pop_scope()\\n        return result\\n\\n    def _wrapped_call_impl(self, *args, **kwargs):\\n        if self._compiled_call_impl is not None:\\n            return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\\n        else:\\n            return self._call_impl(*args, **kwargs)\\n\\n    # torchrec tests the code consistency with the following code\\n    # fmt: off\\n    def _call_impl(self, *args, **kwargs):\\n        forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)\\n        # If we don't have any hooks, we want to skip the rest of the logic in\\n        # this function, and just call forward.\\n        if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks\\n                or _global_backward_pre_hooks or _global_backward_hooks\\n                or _global_forward_hooks or _global_forward_pre_hooks):\\n            return forward_call(*args, **kwargs)\\n\\n        result = None\\n        called_always_called_hooks = set()\\n\\n        def inner():\\n            nonlocal result, args, kwargs\\n\\n            full_backward_hooks, non_full_backward_hooks = [], []\\n            backward_pre_hooks = []\\n            if self._backward_pre_hooks or _global_backward_pre_hooks:\\n                backward_pre_hooks = self._get_backward_pre_hooks()\\n\\n            if self._backward_hooks or _global_backward_hooks:\\n                full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()\\n\\n            if _global_forward_pre_hooks or self._forward_pre_hooks:\\n                for hook_id, hook in (\\n                    *_global_forward_pre_hooks.items(),\\n                    *self._forward_pre_hooks.items(),\\n                ):\\n                    if hook_id in self._forward_pre_hooks_with_kwargs:\\n                        args_kwargs_result = hook(self, args, kwargs)  # type: ignore[misc]\\n                        if args_kwargs_result is not None:\\n                            if isinstance(args_kwargs_result, tuple) and len(args_kwargs_result) == 2:\\n                                args, kwargs = args_kwargs_result\\n                            else:\\n                                raise RuntimeError(\\n                                    \\\"forward pre-hook must return None or a tuple \\\"\\n                                    f\\\"of (new_args, new_kwargs), but got {args_kwargs_result}.\\\"\\n                                )\\n                    else:\\n                        args_result = hook(self, args)\\n                        if args_result is not None:\\n                            if not isinstance(args_result, tuple):\\n                                args_result = (args_result,)\\n                            args = args_result\\n\\n            bw_hook = None\\n            if full_backward_hooks or backward_pre_hooks:\\n                bw_hook = BackwardHook(self, full_backward_hooks, backward_pre_hooks)\\n                args = bw_hook.setup_input_hook(args)\\n\\n            result = forward_call(*args, **kwargs)\\n            if _global_forward_hooks or self._forward_hooks:\\n                for hook_id, hook in (\\n                    *_global_forward_hooks.items(),\\n                    *self._forward_hooks.items(),\\n                ):\\n                    # mark that always called hook is run\\n                    if hook_id in self._forward_hooks_always_called or hook_id in _global_forward_hooks_always_called:\\n                        called_always_called_hooks.add(hook_id)\\n\\n                    if hook_id in self._forward_hooks_with_kwargs:\\n                        hook_result = hook(self, args, kwargs, result)\\n                    else:\\n                        hook_result = hook(self, args, result)\\n\\n                    if hook_result is not None:\\n                        result = hook_result\\n\\n            if bw_hook:\\n                if not isinstance(result, (torch.Tensor, tuple)):\\n                    warnings.warn(\\\"For backward hooks to be called,\\\"\\n                                  \\\" module output should be a Tensor or a tuple of Tensors\\\"\\n                                  f\\\" but received {type(result)}\\\")\\n                result = bw_hook.setup_output_hook(result)\\n\\n            # Handle the non-full backward hooks\\n            if non_full_backward_hooks:\\n                var = result\\n                while not isinstance(var, torch.Tensor):\\n                    if isinstance(var, dict):\\n                        var = next(v for v in var.values() if isinstance(v, torch.Tensor))\\n                    else:\\n                        var = var[0]\\n                grad_fn = var.grad_fn\\n                if grad_fn is not None:\\n                    for hook in non_full_backward_hooks:\\n                        grad_fn.register_hook(_WrappedHook(hook, self))\\n                    self._maybe_warn_non_full_backward_hook(args, result, grad_fn)\\n\\n            return result\\n\\n        from torch.compiler import is_compiling\\n\\n        # This is technically not behavior equivalent when compiling, but it's\\n        # incredibly unlikely we will ever support throwing an exception in NN\\n        # module, and then catching it here, and then reraising it, and then\\n        # catching it again, and expecting the resulting frame to be compiled.\\n        # The reraise here just gunks up our exception handling for no good\\n        # reason.  Don't try to run the always called hooks in event of\\n        # exception.\\n        if is_compiling():\\n            return inner()\\n\\n        try:\\n            return inner()\\n        except Exception:\\n            # run always called hooks if they have not already been run\\n            # For now only forward hooks have the always_call option but perhaps\\n            # this functionality should be added to full backward hooks as well.\\n            for hook_id, hook in _global_forward_hooks.items():\\n                if hook_id in _global_forward_hooks_always_called and hook_id not in called_always_called_hooks:  # type: ignore[possibly-undefined]\\n                    try:\\n                        hook_result = hook(self, args, result)  # type: ignore[possibly-undefined]\\n                        if hook_result is not None:\\n                            result = hook_result\\n                    except Exception as e:\\n                        warnings.warn(\\\"global module forward hook with ``always_call=True`` raised an exception \\\"\\n                                      f\\\"that was silenced as another error was raised in forward: {str(e)}\\\")\\n                        continue\\n\\n            for hook_id, hook in self._forward_hooks.items():\\n                if hook_id in self._forward_hooks_always_called and hook_id not in called_always_called_hooks:  # type: ignore[possibly-undefined]\\n                    try:\\n                        if hook_id in self._forward_hooks_with_kwargs:\\n                            hook_result = hook(self, args, kwargs, result)  # type: ignore[possibly-undefined]\\n                        else:\\n                            hook_result = hook(self, args, result)  # type: ignore[possibly-undefined]\\n                        if hook_result is not None:\\n                            result = hook_result\\n                    except Exception as e:\\n                        warnings.warn(\\\"module forward hook with ``always_call=True`` raised an exception \\\"\\n                                      f\\\"that was silenced as another error was raised in forward: {str(e)}\\\")\\n                        continue\\n            # raise exception raised in try block\\n            raise\\n    # fmt: on\\n\\n    __call__: Callable[..., Any] = _wrapped_call_impl\\n\\n    def __getstate__(self):\\n        state = self.__dict__.copy()\\n        state.pop(\\\"_compiled_call_impl\\\", None)\\n        return state\\n\\n    def __setstate__(self, state):\\n        self.__dict__.update(state)\\n\\n        # Support loading old checkpoints that don't have the following attrs:\\n        if \\\"_forward_pre_hooks\\\" not in self.__dict__:\\n            self._forward_pre_hooks = OrderedDict()\\n        if \\\"_forward_pre_hooks_with_kwargs\\\" not in self.__dict__:\\n            self._forward_pre_hooks_with_kwargs = OrderedDict()\\n        if \\\"_forward_hooks_with_kwargs\\\" not in self.__dict__:\\n            self._forward_hooks_with_kwargs = OrderedDict()\\n        if \\\"_forward_hooks_always_called\\\" not in self.__dict__:\\n            self._forward_hooks_always_called = OrderedDict()\\n        if \\\"_state_dict_hooks\\\" not in self.__dict__:\\n            self._state_dict_hooks = OrderedDict()\\n        if \\\"_state_dict_pre_hooks\\\" not in self.__dict__:\\n            self._state_dict_pre_hooks = OrderedDict()\\n        if \\\"_load_state_dict_pre_hooks\\\" not in self.__dict__:\\n            self._load_state_dict_pre_hooks = OrderedDict()\\n        if \\\"_load_state_dict_post_hooks\\\" not in self.__dict__:\\n            self._load_state_dict_post_hooks = OrderedDict()\\n        if \\\"_non_persistent_buffers_set\\\" not in self.__dict__:\\n            self._non_persistent_buffers_set = set()\\n        if \\\"_is_full_backward_hook\\\" not in self.__dict__:\\n            self._is_full_backward_hook = None\\n        if \\\"_backward_pre_hooks\\\" not in self.__dict__:\\n            self._backward_pre_hooks = OrderedDict()\\n\\n    # On the return type:\\n    # We choose to return `Any` in the `__getattr__` type signature instead of a more strict `Union[Tensor, Module]`.\\n    # This is done for better interop with various type checkers for the end users.\\n    # Having a stricter return type doesn't play nicely with `register_buffer()` and forces\\n    # people to excessively use type-ignores, asserts, casts, etc.\\n    # See full discussion on the problems with returning `Union` here\\n    # https://github.com/microsoft/pyright/issues/4213\\n    def __getattr__(self, name: str) -> Any:\\n        if \\\"_parameters\\\" in self.__dict__:\\n            _parameters = self.__dict__[\\\"_parameters\\\"]\\n            if name in _parameters:\\n                return _parameters[name]\\n        if \\\"_buffers\\\" in self.__dict__:\\n            _buffers = self.__dict__[\\\"_buffers\\\"]\\n            if name in _buffers:\\n                return _buffers[name]\\n        if \\\"_modules\\\" in self.__dict__:\\n            modules = self.__dict__[\\\"_modules\\\"]\\n            if name in modules:\\n                return modules[name]\\n        raise AttributeError(\\n            f\\\"'{type(self).__name__}' object has no attribute '{name}'\\\"\\n        )\\n\\n    def __setattr__(self, name: str, value: Union[Tensor, \\\"Module\\\"]) -> None:\\n        def remove_from(*dicts_or_sets):\\n            for d in dicts_or_sets:\\n                if name in d:\\n                    if isinstance(d, dict):\\n                        del d[name]\\n                    else:\\n                        d.discard(name)\\n\\n        params = self.__dict__.get(\\\"_parameters\\\")\\n        if isinstance(value, Parameter):\\n            if params is None:\\n                raise AttributeError(\\n                    \\\"cannot assign parameters before Module.__init__() call\\\"\\n                )\\n            remove_from(\\n                self.__dict__,\\n                self._buffers,\\n                self._modules,\\n                self._non_persistent_buffers_set,\\n            )\\n            self.register_parameter(name, value)\\n        elif params is not None and name in params:\\n            if value is not None:\\n                raise TypeError(\\n                    f\\\"cannot assign '{torch.typename(value)}' as parameter '{name}' \\\"\\n                    \\\"(torch.nn.Parameter or None expected)\\\"\\n                )\\n            self.register_parameter(name, value)\\n        else:\\n            modules = self.__dict__.get(\\\"_modules\\\")\\n            if isinstance(value, Module):\\n                if modules is None:\\n                    raise AttributeError(\\n                        \\\"cannot assign module before Module.__init__() call\\\"\\n                    )\\n                remove_from(\\n                    self.__dict__,\\n                    self._parameters,\\n                    self._buffers,\\n                    self._non_persistent_buffers_set,\\n                )\\n                for hook in _global_module_registration_hooks.values():\\n                    output = hook(self, name, value)\\n                    if output is not None:\\n                        value = output\\n                modules[name] = value\\n            elif modules is not None and name in modules:\\n                if value is not None:\\n                    raise TypeError(\\n                        f\\\"cannot assign '{torch.typename(value)}' as child module '{name}' \\\"\\n                        \\\"(torch.nn.Module or None expected)\\\"\\n                    )\\n                for hook in _global_module_registration_hooks.values():\\n                    output = hook(self, name, value)\\n                    if output is not None:\\n                        value = output\\n                modules[name] = value\\n            else:\\n                buffers = self.__dict__.get(\\\"_buffers\\\")\\n                if isinstance(value, Buffer) or buffers is not None and name in buffers:\\n                    if value is not None and not isinstance(value, torch.Tensor):\\n                        raise TypeError(\\n                            f\\\"cannot assign '{torch.typename(value)}' as buffer '{name}' \\\"\\n                            \\\"(torch.nn.Buffer, torch.Tensor or None expected)\\\"\\n                        )\\n                    if isinstance(value, Buffer):\\n                        persistent = value.persistent\\n                    else:\\n                        persistent = name not in self._non_persistent_buffers_set\\n                    # === HACK ===\\n                    # This whole block below should just be:\\n                    # self.register_buffer(name, value, persistent)\\n\\n                    # But to support subclasses of nn.Module that (wrongfully) implement a\\n                    # register_buffer() method that doesn't have the \\\"persistent\\\"\\n                    # argument. Only pass it in if it is accepted otherwise assume\\n                    # it is always true\\n                    if self.register_buffer is torch.nn.Module.register_buffer:\\n                        self.register_buffer(name, value, persistent)\\n                    else:\\n                        sign = inspect.signature(self.register_buffer)\\n                        if \\\"persistent\\\" in sign.parameters:\\n                            self.register_buffer(name, value, persistent)\\n                        else:\\n                            if not persistent:\\n                                raise RuntimeError(\\n                                    \\\"Registering a non-persistent buffer \\\"\\n                                    \\\"on a Module subclass that implements \\\"\\n                                    \\\"register_buffer() without the persistent \\\"\\n                                    \\\"argument is not allowed.\\\"\\n                                )\\n                            # Assume that the implementation without the argument has the\\n                            # behavior from before the argument was added: persistent=True\\n                            self.register_buffer(name, value)\\n                    # === HACK END ===\\n                else:\\n                    super().__setattr__(name, value)\\n\\n    def __delattr__(self, name):\\n        if name in self._parameters:\\n            del self._parameters[name]\\n        elif name in self._buffers:\\n            del self._buffers[name]\\n            self._non_persistent_buffers_set.discard(name)\\n        elif name in self._modules:\\n            del self._modules[name]\\n        else:\\n            super().__delattr__(name)\\n\\n    def _register_state_dict_hook(self, hook):\\n        r\\\"\\\"\\\"Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.\\n\\n        It should have the following signature::\\n            hook(module, state_dict, prefix, local_metadata) -> None or state_dict\\n\\n        The registered hooks can modify the ``state_dict`` inplace or return a new one.\\n        If a new ``state_dict`` is returned, it will only be respected if it is the root\\n        module that :meth:`~nn.Module.state_dict` is called from.\\n        \\\"\\\"\\\"\\n        if getattr(hook, \\\"_from_public_api\\\", False):\\n            raise RuntimeError(\\n                \\\"Cannot register the same function as the state dict post hook that was \\\"\\n                \\\"previously registered via register_state_dict_post_hook\\\"\\n            )\\n        handle = RemovableHandle(self._state_dict_hooks)\\n        self._state_dict_hooks[handle.id] = hook\\n        return handle\\n\\n    def register_state_dict_post_hook(self, hook):\\n        r\\\"\\\"\\\"Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.\\n\\n        It should have the following signature::\\n            hook(module, state_dict, prefix, local_metadata) -> None\\n\\n        The registered hooks can modify the ``state_dict`` inplace.\\n        \\\"\\\"\\\"\\n        # In _register_state_dict_hook there was a bug described in\\n        # https://github.com/pytorch/pytorch/issues/117437 where the return value\\n        # was only respected for the root module but not child submodules.\\n        # We fix this in this public version by only allowing inplace modifications on\\n        # the state_dict by the hook. However, since hooks registered via both these\\n        # APIs will be added to `_state_dict_hooks` and the type of `_state_dict_hooks`\\n        # cannot be changed due to many dependencies on it, we mark a hook\\n        # as being registered via the public API by setting `_from_public_api` on it.\\n        # In the implementation of `state_dict`, if the callable does not have this\\n        # flag, the old behavior of respecting the return value will be preserved\\n        # for the root module, otherwise, we ensure that the hook returns None.\\n        hook._from_public_api = True\\n        handle = RemovableHandle(self._state_dict_hooks)\\n        self._state_dict_hooks[handle.id] = hook\\n        return handle\\n\\n    def register_state_dict_pre_hook(self, hook):\\n        r\\\"\\\"\\\"Register a pre-hook for the :meth:`~torch.nn.Module.state_dict` method.\\n\\n        It should have the following signature::\\n            hook(module, prefix, keep_vars) -> None\\n\\n        The registered hooks can be used to perform pre-processing before the ``state_dict``\\n        call is made.\\n        \\\"\\\"\\\"\\n        handle = RemovableHandle(self._state_dict_pre_hooks)\\n        self._state_dict_pre_hooks[handle.id] = hook\\n        return handle\\n\\n    def _save_to_state_dict(self, destination, prefix, keep_vars):\\n        r\\\"\\\"\\\"Save module state to the `destination` dictionary.\\n\\n        The `destination` dictionary will contain the state\\n        of the module, but not its descendants. This is called on every\\n        submodule in :meth:`~torch.nn.Module.state_dict`.\\n\\n        In rare cases, subclasses can achieve class-specific behavior by\\n        overriding this method with custom logic.\\n\\n        Args:\\n            destination (dict): a dict where state will be stored\\n            prefix (str): the prefix for parameters and buffers used in this\\n                module\\n        \\\"\\\"\\\"\\n        for name, param in self._parameters.items():\\n            if param is not None:\\n                destination[prefix + name] = param if keep_vars else param.detach()\\n        for name, buf in self._buffers.items():\\n            if buf is not None and name not in self._non_persistent_buffers_set:\\n                destination[prefix + name] = buf if keep_vars else buf.detach()\\n        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX\\n        if (\\n            getattr(self.__class__, \\\"get_extra_state\\\", Module.get_extra_state)\\n            is not Module.get_extra_state\\n        ):\\n            destination[extra_state_key] = self.get_extra_state()\\n\\n    # The user can pass an optional arbitrary mappable object to `state_dict`, in which case `state_dict` returns\\n    # back that same object. But if they pass nothing, an `OrderedDict` is created and returned.\\n    T_destination = TypeVar(\\\"T_destination\\\", bound=Dict[str, Any])\\n\\n    @overload\\n    def state_dict(\\n        self, *, destination: T_destination, prefix: str = ..., keep_vars: bool = ...\\n    ) -> T_destination:\\n        ...\\n\\n    @overload\\n    def state_dict(self, *, prefix: str = ..., keep_vars: bool = ...) -> Dict[str, Any]:\\n        ...\\n\\n    # TODO: Change `*args` to `*` and remove the corresponding warning in docs when BC allows.\\n    # Also remove the logic for arg parsing together.\\n    def state_dict(self, *args, destination=None, prefix=\\\"\\\", keep_vars=False):\\n        r\\\"\\\"\\\"Return a dictionary containing references to the whole state of the module.\\n\\n        Both parameters and persistent buffers (e.g. running averages) are\\n        included. Keys are corresponding parameter and buffer names.\\n        Parameters and buffers set to ``None`` are not included.\\n\\n        .. note::\\n            The returned object is a shallow copy. It contains references\\n            to the module's parameters and buffers.\\n\\n        .. warning::\\n            Currently ``state_dict()`` also accepts positional arguments for\\n            ``destination``, ``prefix`` and ``keep_vars`` in order. However,\\n            this is being deprecated and keyword arguments will be enforced in\\n            future releases.\\n\\n        .. warning::\\n            Please avoid the use of argument ``destination`` as it is not\\n            designed for end-users.\\n\\n        Args:\\n            destination (dict, optional): If provided, the state of module will\\n                be updated into the dict and the same object is returned.\\n                Otherwise, an ``OrderedDict`` will be created and returned.\\n                Default: ``None``.\\n            prefix (str, optional): a prefix added to parameter and buffer\\n                names to compose the keys in state_dict. Default: ``''``.\\n            keep_vars (bool, optional): by default the :class:`~torch.Tensor` s\\n                returned in the state dict are detached from autograd. If it's\\n                set to ``True``, detaching will not be performed.\\n                Default: ``False``.\\n\\n        Returns:\\n            dict:\\n                a dictionary containing a whole state of the module\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> module.state_dict().keys()\\n            ['bias', 'weight']\\n\\n        \\\"\\\"\\\"\\n        # TODO: Remove `args` and the parsing logic when BC allows.\\n        if len(args) > 0:\\n            # DeprecationWarning is ignored by default\\n            warnings.warn(\\n                \\\"Positional args are being deprecated, use kwargs instead. Refer to \\\"\\n                \\\"https://pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.state_dict\\\"\\n                \\\" for details.\\\",\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n            if destination is None:\\n                destination = args[0]\\n            if len(args) > 1 and prefix == \\\"\\\":\\n                prefix = args[1]\\n            if len(args) > 2 and keep_vars is False:\\n                keep_vars = args[2]\\n\\n        if destination is None:\\n            destination = OrderedDict()\\n            destination._metadata = OrderedDict()\\n\\n        local_metadata = dict(version=self._version)\\n        if hasattr(destination, \\\"_metadata\\\"):\\n            destination._metadata[prefix[:-1]] = local_metadata\\n\\n        for hook in self._state_dict_pre_hooks.values():\\n            hook(self, prefix, keep_vars)\\n        self._save_to_state_dict(destination, prefix, keep_vars)\\n        for name, module in self._modules.items():\\n            if module is not None:\\n                module.state_dict(\\n                    destination=destination,\\n                    prefix=prefix + name + \\\".\\\",\\n                    keep_vars=keep_vars,\\n                )\\n        for hook in self._state_dict_hooks.values():\\n            hook_result = hook(self, destination, prefix, local_metadata)\\n            if not getattr(hook, \\\"_from_public_api\\\", False):\\n                if hook_result is not None:\\n                    destination = hook_result\\n            else:\\n                if hook_result is not None:\\n                    raise RuntimeError(\\\"state_dict post-hook must return None\\\")\\n        return destination\\n\\n    def _register_load_state_dict_pre_hook(self, hook, with_module=False):\\n        r\\\"\\\"\\\"See :meth:`~torch.nn.Module.register_load_state_dict_pre_hook` for details.\\n\\n        A subtle difference is that if ``with_module`` is set to ``False``, then the\\n        hook will not take the ``module`` as the first argument whereas\\n        :meth:`~torch.nn.Module.register_load_state_dict_pre_hook` always takes the\\n        ``module`` as the first argument.\\n\\n        Arguments:\\n            hook (Callable): Callable hook that will be invoked before\\n                loading the state dict.\\n            with_module (bool, optional): Whether or not to pass the module\\n                instance to the hook as the first parameter.\\n        \\\"\\\"\\\"\\n        handle = RemovableHandle(self._load_state_dict_pre_hooks)\\n        self._load_state_dict_pre_hooks[handle.id] = _WrappedHook(\\n            hook, self if with_module else None\\n        )\\n        return handle\\n\\n    def register_load_state_dict_pre_hook(self, hook):\\n        r\\\"\\\"\\\"Register a pre-hook to be run before module's :meth:`~nn.Module.load_state_dict` is called.\\n\\n        It should have the following signature::\\n            hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None  # noqa: B950\\n\\n        Arguments:\\n            hook (Callable): Callable hook that will be invoked before\\n                loading the state dict.\\n        \\\"\\\"\\\"\\n        return self._register_load_state_dict_pre_hook(hook, with_module=True)\\n\\n    def register_load_state_dict_post_hook(self, hook):\\n        r\\\"\\\"\\\"Register a post-hook to be run after module's :meth:`~nn.Module.load_state_dict` is called.\\n\\n        It should have the following signature::\\n            hook(module, incompatible_keys) -> None\\n\\n        The ``module`` argument is the current module that this hook is registered\\n        on, and the ``incompatible_keys`` argument is a ``NamedTuple`` consisting\\n        of attributes ``missing_keys`` and ``unexpected_keys``. ``missing_keys``\\n        is a ``list`` of ``str`` containing the missing keys and\\n        ``unexpected_keys`` is a ``list`` of ``str`` containing the unexpected keys.\\n\\n        The given incompatible_keys can be modified inplace if needed.\\n\\n        Note that the checks performed when calling :func:`load_state_dict` with\\n        ``strict=True`` are affected by modifications the hook makes to\\n        ``missing_keys`` or ``unexpected_keys``, as expected. Additions to either\\n        set of keys will result in an error being thrown when ``strict=True``, and\\n        clearing out both missing and unexpected keys will avoid an error.\\n\\n        Returns:\\n            :class:`torch.utils.hooks.RemovableHandle`:\\n                a handle that can be used to remove the added hook by calling\\n                ``handle.remove()``\\n        \\\"\\\"\\\"\\n        handle = RemovableHandle(self._load_state_dict_post_hooks)\\n        self._load_state_dict_post_hooks[handle.id] = hook\\n        return handle\\n\\n    def _load_from_state_dict(\\n        self,\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ):\\n        r\\\"\\\"\\\"Copy parameters and buffers from :attr:`state_dict` into only this module, but not its descendants.\\n\\n        This is called on every submodule\\n        in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this\\n        module in input :attr:`state_dict` is provided as :attr:`local_metadata`.\\n        For state dicts without metadata, :attr:`local_metadata` is empty.\\n        Subclasses can achieve class-specific backward compatible loading using\\n        the version number at `local_metadata.get(\\\"version\\\", None)`.\\n        Additionally, :attr:`local_metadata` can also contain the key\\n        `assign_to_params_buffers` that indicates whether keys should be\\n        assigned their corresponding tensor in the state_dict.\\n\\n        .. note::\\n            :attr:`state_dict` is not the same object as the input\\n            :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So\\n            it can be modified.\\n\\n        Args:\\n            state_dict (dict): a dict containing parameters and\\n                persistent buffers.\\n            prefix (str): the prefix for parameters and buffers used in this\\n                module\\n            local_metadata (dict): a dict containing the metadata for this module.\\n                See\\n            strict (bool): whether to strictly enforce that the keys in\\n                :attr:`state_dict` with :attr:`prefix` match the names of\\n                parameters and buffers in this module\\n            missing_keys (list of str): if ``strict=True``, add missing keys to\\n                this list\\n            unexpected_keys (list of str): if ``strict=True``, add unexpected\\n                keys to this list\\n            error_msgs (list of str): error messages should be added to this\\n                list, and will be reported together in\\n                :meth:`~torch.nn.Module.load_state_dict`\\n        \\\"\\\"\\\"\\n        for hook in self._load_state_dict_pre_hooks.values():\\n            hook(\\n                state_dict,\\n                prefix,\\n                local_metadata,\\n                strict,\\n                missing_keys,\\n                unexpected_keys,\\n                error_msgs,\\n            )\\n\\n        persistent_buffers = {\\n            k: v\\n            for k, v in self._buffers.items()\\n            if k not in self._non_persistent_buffers_set\\n        }\\n        local_name_params = itertools.chain(\\n            self._parameters.items(), persistent_buffers.items()\\n        )\\n        local_state = {k: v for k, v in local_name_params if v is not None}\\n        assign_to_params_buffers = local_metadata.get(\\\"assign_to_params_buffers\\\", False)\\n        use_swap_tensors = torch.__future__.get_swap_module_params_on_conversion()\\n\\n        for name, param in local_state.items():\\n            key = prefix + name\\n            if key in state_dict:\\n                input_param = state_dict[key]\\n                if not torch.overrides.is_tensor_like(input_param):\\n                    error_msgs.append(\\n                        f'While copying the parameter named \\\"{key}\\\", '\\n                        \\\"expected torch.Tensor or Tensor-like object from checkpoint but \\\"\\n                        f\\\"received {type(input_param)}\\\"\\n                    )\\n                    continue\\n\\n                # This is used to avoid copying uninitialized parameters into\\n                # non-lazy modules, since they dont have the hook to do the checks\\n                # in such case, it will error when accessing the .shape attribute.\\n                is_param_lazy = torch.nn.parameter.is_lazy(param)\\n                # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+\\n                if (\\n                    not is_param_lazy\\n                    and len(param.shape) == 0\\n                    and len(input_param.shape) == 1\\n                ):\\n                    input_param = input_param[0]\\n\\n                if not is_param_lazy and input_param.shape != param.shape:\\n                    # local shape should match the one in checkpoint\\n                    error_msgs.append(\\n                        f\\\"size mismatch for {key}: copying a param with shape {input_param.shape} from checkpoint, \\\"\\n                        f\\\"the shape in current model is {param.shape}.\\\"\\n                    )\\n                    continue\\n\\n                if (\\n                    param.is_meta\\n                    and not input_param.is_meta\\n                    and not assign_to_params_buffers\\n                ):\\n                    warnings.warn(\\n                        f\\\"for {key}: copying from a non-meta parameter in the checkpoint to a meta \\\"\\n                        \\\"parameter in the current model, which is a no-op. (Did you mean to \\\"\\n                        \\\"pass `assign=True` to assign items in the state dictionary to their \\\"\\n                        \\\"corresponding key in the module instead of copying them in place?)\\\"\\n                    )\\n\\n                try:\\n                    with torch.no_grad():\\n                        if use_swap_tensors:\\n                            new_input_param = param.module_load(\\n                                input_param, assign=assign_to_params_buffers\\n                            )\\n                            if id(new_input_param) == id(input_param) or id(\\n                                new_input_param\\n                            ) == id(param):\\n                                raise RuntimeError(\\n                                    \\\"module_load returned one of self or other, please .detach() \\\"\\n                                    \\\"the result if returning one of the inputs in module_load\\\"\\n                                )\\n                            if isinstance(param, torch.nn.Parameter):\\n                                if not isinstance(new_input_param, torch.nn.Parameter):\\n                                    new_input_param = torch.nn.Parameter(\\n                                        new_input_param,\\n                                        requires_grad=param.requires_grad,\\n                                    )\\n                                else:\\n                                    new_input_param.requires_grad_(param.requires_grad)\\n                            torch.utils.swap_tensors(param, new_input_param)\\n                            del new_input_param\\n                        elif assign_to_params_buffers:\\n                            # Shape checks are already done above\\n                            if isinstance(param, torch.nn.Parameter):\\n                                if not isinstance(input_param, torch.nn.Parameter):\\n                                    input_param = torch.nn.Parameter(\\n                                        input_param, requires_grad=param.requires_grad\\n                                    )\\n                                else:\\n                                    input_param.requires_grad_(param.requires_grad)\\n                            setattr(self, name, input_param)\\n                        else:\\n                            param.copy_(input_param)\\n                except Exception as ex:\\n                    action = \\\"swapping\\\" if use_swap_tensors else \\\"copying\\\"\\n                    error_msgs.append(\\n                        f'While {action} the parameter named \\\"{key}\\\", '\\n                        f\\\"whose dimensions in the model are {param.size()} and \\\"\\n                        f\\\"whose dimensions in the checkpoint are {input_param.size()}, \\\"\\n                        f\\\"an exception occurred : {ex.args}.\\\"\\n                    )\\n            elif strict:\\n                missing_keys.append(key)\\n\\n        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX\\n        if (\\n            getattr(self.__class__, \\\"set_extra_state\\\", Module.set_extra_state)\\n            is not Module.set_extra_state\\n        ):\\n            if extra_state_key in state_dict:\\n                self.set_extra_state(state_dict[extra_state_key])\\n            elif strict:\\n                missing_keys.append(extra_state_key)\\n        elif strict and (extra_state_key in state_dict):\\n            unexpected_keys.append(extra_state_key)\\n\\n        if strict:\\n            for key in state_dict.keys():\\n                if key.startswith(prefix) and key != extra_state_key:\\n                    input_name = key[len(prefix) :].split(\\\".\\\", 1)\\n                    # Must be Module if it have attributes\\n                    if len(input_name) > 1:\\n                        if input_name[0] not in self._modules:\\n                            unexpected_keys.append(key)\\n                    elif input_name[0] not in local_state:\\n                        unexpected_keys.append(key)\\n\\n    def load_state_dict(\\n        self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False\\n    ):\\n        r\\\"\\\"\\\"Copy parameters and buffers from :attr:`state_dict` into this module and its descendants.\\n\\n        If :attr:`strict` is ``True``, then\\n        the keys of :attr:`state_dict` must exactly match the keys returned\\n        by this module's :meth:`~torch.nn.Module.state_dict` function.\\n\\n        .. warning::\\n            If :attr:`assign` is ``True`` the optimizer must be created after\\n            the call to :attr:`load_state_dict` unless\\n            :func:`~torch.__future__.get_swap_module_params_on_conversion` is ``True``.\\n\\n        Args:\\n            state_dict (dict): a dict containing parameters and\\n                persistent buffers.\\n            strict (bool, optional): whether to strictly enforce that the keys\\n                in :attr:`state_dict` match the keys returned by this module's\\n                :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\\n            assign (bool, optional): When ``False``, the properties of the tensors\\n                in the current module are preserved while when ``True``, the\\n                properties of the Tensors in the state dict are preserved. The only\\n                exception is the ``requires_grad`` field of :class:`~torch.nn.Parameter`s\\n                for which the value from the module is preserved.\\n                Default: ``False``\\n\\n        Returns:\\n            ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\\n                * **missing_keys** is a list of str containing any keys that are expected\\n                    by this module but missing from the provided ``state_dict``.\\n                * **unexpected_keys** is a list of str containing the keys that are not\\n                    expected by this module but present in the provided ``state_dict``.\\n\\n        Note:\\n            If a parameter or buffer is registered as ``None`` and its corresponding key\\n            exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a\\n            ``RuntimeError``.\\n        \\\"\\\"\\\"\\n        if not isinstance(state_dict, Mapping):\\n            raise TypeError(\\n                f\\\"Expected state_dict to be dict-like, got {type(state_dict)}.\\\"\\n            )\\n\\n        missing_keys: List[str] = []\\n        unexpected_keys: List[str] = []\\n        error_msgs: List[str] = []\\n\\n        # copy state_dict so _load_from_state_dict can modify it\\n        metadata = getattr(state_dict, \\\"_metadata\\\", None)\\n        state_dict = OrderedDict(state_dict)\\n        if metadata is not None:\\n            # mypy isn't aware that \\\"_metadata\\\" exists in state_dict\\n            state_dict._metadata = metadata  # type: ignore[attr-defined]\\n\\n        def load(module, local_state_dict, prefix=\\\"\\\"):\\n            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})\\n            if assign:\\n                local_metadata[\\\"assign_to_params_buffers\\\"] = assign\\n            module._load_from_state_dict(\\n                local_state_dict,\\n                prefix,\\n                local_metadata,\\n                True,\\n                missing_keys,\\n                unexpected_keys,\\n                error_msgs,\\n            )\\n            for name, child in module._modules.items():\\n                if child is not None:\\n                    child_prefix = prefix + name + \\\".\\\"\\n                    child_state_dict = {\\n                        k: v\\n                        for k, v in local_state_dict.items()\\n                        if k.startswith(child_prefix)\\n                    }\\n                    load(child, child_state_dict, child_prefix)  # noqa: F821\\n\\n            # Note that the hook can modify missing_keys and unexpected_keys.\\n            incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys)\\n            for hook in module._load_state_dict_post_hooks.values():\\n                out = hook(module, incompatible_keys)\\n                assert out is None, (\\n                    \\\"Hooks registered with ``register_load_state_dict_post_hook`` are not\\\"\\n                    \\\"expected to return new values, if incompatible_keys need to be modified,\\\"\\n                    \\\"it should be done inplace.\\\"\\n                )\\n\\n        load(self, state_dict)\\n        del load\\n\\n        if strict:\\n            if len(unexpected_keys) > 0:\\n                error_msgs.insert(\\n                    0,\\n                    \\\"Unexpected key(s) in state_dict: {}. \\\".format(\\n                        \\\", \\\".join(f'\\\"{k}\\\"' for k in unexpected_keys)\\n                    ),\\n                )\\n            if len(missing_keys) > 0:\\n                error_msgs.insert(\\n                    0,\\n                    \\\"Missing key(s) in state_dict: {}. \\\".format(\\n                        \\\", \\\".join(f'\\\"{k}\\\"' for k in missing_keys)\\n                    ),\\n                )\\n\\n        if len(error_msgs) > 0:\\n            raise RuntimeError(\\n                \\\"Error(s) in loading state_dict for {}:\\\\n\\\\t{}\\\".format(\\n                    self.__class__.__name__, \\\"\\\\n\\\\t\\\".join(error_msgs)\\n                )\\n            )\\n        return _IncompatibleKeys(missing_keys, unexpected_keys)\\n\\n    def _named_members(\\n        self, get_members_fn, prefix=\\\"\\\", recurse=True, remove_duplicate: bool = True\\n    ):\\n        r\\\"\\\"\\\"Help yield various names + members of modules.\\\"\\\"\\\"\\n        memo = set()\\n        modules = (\\n            self.named_modules(prefix=prefix, remove_duplicate=remove_duplicate)\\n            if recurse\\n            else [(prefix, self)]\\n        )\\n        for module_prefix, module in modules:\\n            members = get_members_fn(module)\\n            for k, v in members:\\n                if v is None or v in memo:\\n                    continue\\n                if remove_duplicate:\\n                    memo.add(v)\\n                name = module_prefix + (\\\".\\\" if module_prefix else \\\"\\\") + k\\n                yield name, v\\n\\n    def parameters(self, recurse: bool = True) -> Iterator[Parameter]:\\n        r\\\"\\\"\\\"Return an iterator over module parameters.\\n\\n        This is typically passed to an optimizer.\\n\\n        Args:\\n            recurse (bool): if True, then yields parameters of this module\\n                and all submodules. Otherwise, yields only parameters that\\n                are direct members of this module.\\n\\n        Yields:\\n            Parameter: module parameter\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> for param in model.parameters():\\n            >>>     print(type(param), param.size())\\n            <class 'torch.Tensor'> (20L,)\\n            <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\\n\\n        \\\"\\\"\\\"\\n        for name, param in self.named_parameters(recurse=recurse):\\n            yield param\\n\\n    def named_parameters(\\n        self, prefix: str = \\\"\\\", recurse: bool = True, remove_duplicate: bool = True\\n    ) -> Iterator[Tuple[str, Parameter]]:\\n        r\\\"\\\"\\\"Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.\\n\\n        Args:\\n            prefix (str): prefix to prepend to all parameter names.\\n            recurse (bool): if True, then yields parameters of this module\\n                and all submodules. Otherwise, yields only parameters that\\n                are direct members of this module.\\n            remove_duplicate (bool, optional): whether to remove the duplicated\\n                parameters in the result. Defaults to True.\\n\\n        Yields:\\n            (str, Parameter): Tuple containing the name and parameter\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> for name, param in self.named_parameters():\\n            >>>     if name in ['bias']:\\n            >>>         print(param.size())\\n\\n        \\\"\\\"\\\"\\n        gen = self._named_members(\\n            lambda module: module._parameters.items(),\\n            prefix=prefix,\\n            recurse=recurse,\\n            remove_duplicate=remove_duplicate,\\n        )\\n        yield from gen\\n\\n    def buffers(self, recurse: bool = True) -> Iterator[Tensor]:\\n        r\\\"\\\"\\\"Return an iterator over module buffers.\\n\\n        Args:\\n            recurse (bool): if True, then yields buffers of this module\\n                and all submodules. Otherwise, yields only buffers that\\n                are direct members of this module.\\n\\n        Yields:\\n            torch.Tensor: module buffer\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> for buf in model.buffers():\\n            >>>     print(type(buf), buf.size())\\n            <class 'torch.Tensor'> (20L,)\\n            <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\\n\\n        \\\"\\\"\\\"\\n        for _, buf in self.named_buffers(recurse=recurse):\\n            yield buf\\n\\n    def named_buffers(\\n        self, prefix: str = \\\"\\\", recurse: bool = True, remove_duplicate: bool = True\\n    ) -> Iterator[Tuple[str, Tensor]]:\\n        r\\\"\\\"\\\"Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.\\n\\n        Args:\\n            prefix (str): prefix to prepend to all buffer names.\\n            recurse (bool, optional): if True, then yields buffers of this module\\n                and all submodules. Otherwise, yields only buffers that\\n                are direct members of this module. Defaults to True.\\n            remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.\\n\\n        Yields:\\n            (str, torch.Tensor): Tuple containing the name and buffer\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> for name, buf in self.named_buffers():\\n            >>>     if name in ['running_var']:\\n            >>>         print(buf.size())\\n\\n        \\\"\\\"\\\"\\n        gen = self._named_members(\\n            lambda module: module._buffers.items(),\\n            prefix=prefix,\\n            recurse=recurse,\\n            remove_duplicate=remove_duplicate,\\n        )\\n        yield from gen\\n\\n    def children(self) -> Iterator[\\\"Module\\\"]:\\n        r\\\"\\\"\\\"Return an iterator over immediate children modules.\\n\\n        Yields:\\n            Module: a child module\\n        \\\"\\\"\\\"\\n        for name, module in self.named_children():\\n            yield module\\n\\n    def named_children(self) -> Iterator[Tuple[str, \\\"Module\\\"]]:\\n        r\\\"\\\"\\\"Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.\\n\\n        Yields:\\n            (str, Module): Tuple containing a name and child module\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined vars\\\")\\n            >>> for name, module in model.named_children():\\n            >>>     if name in ['conv4', 'conv5']:\\n            >>>         print(module)\\n\\n        \\\"\\\"\\\"\\n        memo = set()\\n        for name, module in self._modules.items():\\n            if module is not None and module not in memo:\\n                memo.add(module)\\n                yield name, module\\n\\n    def modules(self) -> Iterator[\\\"Module\\\"]:\\n        r\\\"\\\"\\\"Return an iterator over all modules in the network.\\n\\n        Yields:\\n            Module: a module in the network\\n\\n        Note:\\n            Duplicate modules are returned only once. In the following\\n            example, ``l`` will be returned only once.\\n\\n        Example::\\n\\n            >>> l = nn.Linear(2, 2)\\n            >>> net = nn.Sequential(l, l)\\n            >>> for idx, m in enumerate(net.modules()):\\n            ...     print(idx, '->', m)\\n\\n            0 -> Sequential(\\n              (0): Linear(in_features=2, out_features=2, bias=True)\\n              (1): Linear(in_features=2, out_features=2, bias=True)\\n            )\\n            1 -> Linear(in_features=2, out_features=2, bias=True)\\n\\n        \\\"\\\"\\\"\\n        for _, module in self.named_modules():\\n            yield module\\n\\n    def named_modules(\\n        self,\\n        memo: Optional[Set[\\\"Module\\\"]] = None,\\n        prefix: str = \\\"\\\",\\n        remove_duplicate: bool = True,\\n    ):\\n        r\\\"\\\"\\\"Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.\\n\\n        Args:\\n            memo: a memo to store the set of modules already added to the result\\n            prefix: a prefix that will be added to the name of the module\\n            remove_duplicate: whether to remove the duplicated module instances in the result\\n                or not\\n\\n        Yields:\\n            (str, Module): Tuple of name and module\\n\\n        Note:\\n            Duplicate modules are returned only once. In the following\\n            example, ``l`` will be returned only once.\\n\\n        Example::\\n\\n            >>> l = nn.Linear(2, 2)\\n            >>> net = nn.Sequential(l, l)\\n            >>> for idx, m in enumerate(net.named_modules()):\\n            ...     print(idx, '->', m)\\n\\n            0 -> ('', Sequential(\\n              (0): Linear(in_features=2, out_features=2, bias=True)\\n              (1): Linear(in_features=2, out_features=2, bias=True)\\n            ))\\n            1 -> ('0', Linear(in_features=2, out_features=2, bias=True))\\n\\n        \\\"\\\"\\\"\\n        if memo is None:\\n            memo = set()\\n        if self not in memo:\\n            if remove_duplicate:\\n                memo.add(self)\\n            yield prefix, self\\n            for name, module in self._modules.items():\\n                if module is None:\\n                    continue\\n                submodule_prefix = prefix + (\\\".\\\" if prefix else \\\"\\\") + name\\n                yield from module.named_modules(\\n                    memo, submodule_prefix, remove_duplicate\\n                )\\n\\n    def train(self: T, mode: bool = True) -> T:\\n        r\\\"\\\"\\\"Set the module in training mode.\\n\\n        This has any effect only on certain modules. See documentations of\\n        particular modules for details of their behaviors in training/evaluation\\n        mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\\n        etc.\\n\\n        Args:\\n            mode (bool): whether to set training mode (``True``) or evaluation\\n                         mode (``False``). Default: ``True``.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        if not isinstance(mode, bool):\\n            raise ValueError(\\\"training mode is expected to be boolean\\\")\\n        self.training = mode\\n        for module in self.children():\\n            module.train(mode)\\n        return self\\n\\n    def eval(self: T) -> T:\\n        r\\\"\\\"\\\"Set the module in evaluation mode.\\n\\n        This has any effect only on certain modules. See documentations of\\n        particular modules for details of their behaviors in training/evaluation\\n        mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\\n        etc.\\n\\n        This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.\\n\\n        See :ref:`locally-disable-grad-doc` for a comparison between\\n        `.eval()` and several similar mechanisms that may be confused with it.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        return self.train(False)\\n\\n    def requires_grad_(self: T, requires_grad: bool = True) -> T:\\n        r\\\"\\\"\\\"Change if autograd should record operations on parameters in this module.\\n\\n        This method sets the parameters' :attr:`requires_grad` attributes\\n        in-place.\\n\\n        This method is helpful for freezing part of the module for finetuning\\n        or training parts of a model individually (e.g., GAN training).\\n\\n        See :ref:`locally-disable-grad-doc` for a comparison between\\n        `.requires_grad_()` and several similar mechanisms that may be confused with it.\\n\\n        Args:\\n            requires_grad (bool): whether autograd should record operations on\\n                                  parameters in this module. Default: ``True``.\\n\\n        Returns:\\n            Module: self\\n        \\\"\\\"\\\"\\n        for p in self.parameters():\\n            p.requires_grad_(requires_grad)\\n        return self\\n\\n    def zero_grad(self, set_to_none: bool = True) -> None:\\n        r\\\"\\\"\\\"Reset gradients of all model parameters.\\n\\n        See similar function under :class:`torch.optim.Optimizer` for more context.\\n\\n        Args:\\n            set_to_none (bool): instead of setting to zero, set the grads to None.\\n                See :meth:`torch.optim.Optimizer.zero_grad` for details.\\n        \\\"\\\"\\\"\\n        if getattr(self, \\\"_is_replica\\\", False):\\n            warnings.warn(\\n                \\\"Calling .zero_grad() from a module created with nn.DataParallel() has no effect. \\\"\\n                \\\"The parameters are copied (in a differentiable manner) from the original module. \\\"\\n                \\\"This means they are not leaf nodes in autograd and so don't accumulate gradients. \\\"\\n                \\\"If you need gradients in your forward method, consider using autograd.grad instead.\\\"\\n            )\\n\\n        for p in self.parameters():\\n            if p.grad is not None:\\n                if set_to_none:\\n                    p.grad = None\\n                else:\\n                    if p.grad.grad_fn is not None:\\n                        p.grad.detach_()\\n                    else:\\n                        p.grad.requires_grad_(False)\\n                    p.grad.zero_()\\n\\n    def share_memory(self: T) -> T:\\n        r\\\"\\\"\\\"See :meth:`torch.Tensor.share_memory_`.\\\"\\\"\\\"\\n        return self._apply(lambda t: t.share_memory_())\\n\\n    def _get_name(self):\\n        return self.__class__.__name__\\n\\n    def extra_repr(self) -> str:\\n        r\\\"\\\"\\\"Set the extra representation of the module.\\n\\n        To print customized extra information, you should re-implement\\n        this method in your own modules. Both single-line and multi-line\\n        strings are acceptable.\\n        \\\"\\\"\\\"\\n        return \\\"\\\"\\n\\n    def __repr__(self):\\n        # We treat the extra repr like the sub-module, one item per line\\n        extra_lines = []\\n        extra_repr = self.extra_repr()\\n        # empty string will be split into list ['']\\n        if extra_repr:\\n            extra_lines = extra_repr.split(\\\"\\\\n\\\")\\n        child_lines = []\\n        for key, module in self._modules.items():\\n            mod_str = repr(module)\\n            mod_str = _addindent(mod_str, 2)\\n            child_lines.append(\\\"(\\\" + key + \\\"): \\\" + mod_str)\\n        lines = extra_lines + child_lines\\n\\n        main_str = self._get_name() + \\\"(\\\"\\n        if lines:\\n            # simple one-liner info, which most builtin Modules will use\\n            if len(extra_lines) == 1 and not child_lines:\\n                main_str += extra_lines[0]\\n            else:\\n                main_str += \\\"\\\\n  \\\" + \\\"\\\\n  \\\".join(lines) + \\\"\\\\n\\\"\\n\\n        main_str += \\\")\\\"\\n        return main_str\\n\\n    def __dir__(self):\\n        module_attrs = dir(self.__class__)\\n        attrs = list(self.__dict__.keys())\\n        parameters = list(self._parameters.keys())\\n        modules = list(self._modules.keys())\\n        buffers = list(self._buffers.keys())\\n        keys = module_attrs + attrs + parameters + modules + buffers\\n\\n        # Eliminate attrs that are not legal Python variable names\\n        keys = [key for key in keys if not key[0].isdigit()]\\n\\n        return sorted(keys)\\n\\n    def _replicate_for_data_parallel(self):\\n        replica = self.__new__(type(self))\\n        replica.__dict__ = self.__dict__.copy()\\n\\n        # replicas do not have parameters themselves, the replicas reference the original\\n        # module.\\n        replica._parameters = {}\\n        replica._buffers = replica._buffers.copy()\\n        replica._modules = replica._modules.copy()\\n        replica._is_replica = True  # type: ignore[assignment]\\n\\n        return replica\\n\\n    def compile(self, *args, **kwargs):\\n        \\\"\\\"\\\"\\n        Compile this Module's forward using :func:`torch.compile`.\\n\\n        This Module's `__call__` method is compiled and all arguments are passed as-is\\n        to :func:`torch.compile`.\\n\\n        See :func:`torch.compile` for details on the arguments for this function.\\n        \\\"\\\"\\\"\\n        self._compiled_call_impl = torch.compile(self._call_impl, *args, **kwargs)\\n\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"ChannelShuffle\\\"]\\n\\n\\nclass ChannelShuffle(Module):\\n    r\\\"\\\"\\\"Divides and rearranges the channels in a tensor.\\n\\n    This operation divides the channels in a tensor of shape :math:`(N, C, *)`\\n    into g groups as :math:`(N, \\\\frac{C}{g}, g, *)` and shuffles them,\\n    while retaining the original tensor shape in the final output.\\n\\n    Args:\\n        groups (int): number of groups to divide channels in.\\n\\n    Examples::\\n\\n        >>> channel_shuffle = nn.ChannelShuffle(2)\\n        >>> input = torch.arange(1, 17, dtype=torch.float32).view(1, 4, 2, 2)\\n        >>> input\\n        tensor([[[[ 1.,  2.],\\n                  [ 3.,  4.]],\\n                 [[ 5.,  6.],\\n                  [ 7.,  8.]],\\n                 [[ 9., 10.],\\n                  [11., 12.]],\\n                 [[13., 14.],\\n                  [15., 16.]]]])\\n        >>> output = channel_shuffle(input)\\n        >>> output\\n        tensor([[[[ 1.,  2.],\\n                  [ 3.,  4.]],\\n                 [[ 9., 10.],\\n                  [11., 12.]],\\n                 [[ 5.,  6.],\\n                  [ 7.,  8.]],\\n                 [[13., 14.],\\n                  [15., 16.]]]])\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"groups\\\"]\\n    groups: int\\n\\n    def __init__(self, groups: int) -> None:\\n        super().__init__()\\n        self.groups = groups\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.channel_shuffle(input, self.groups)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"groups={self.groups}\\\"\\n\\n\\n# mypy: allow-untyped-defs\\nimport collections\\nfrom itertools import repeat\\nfrom typing import Any, Dict, List\\n\\n\\n__all__ = [\\\"consume_prefix_in_state_dict_if_present\\\"]\\n\\n\\ndef _ntuple(n, name=\\\"parse\\\"):\\n    def parse(x):\\n        if isinstance(x, collections.abc.Iterable):\\n            return tuple(x)\\n        return tuple(repeat(x, n))\\n\\n    parse.__name__ = name\\n    return parse\\n\\n\\n_single = _ntuple(1, \\\"_single\\\")\\n_pair = _ntuple(2, \\\"_pair\\\")\\n_triple = _ntuple(3, \\\"_triple\\\")\\n_quadruple = _ntuple(4, \\\"_quadruple\\\")\\n\\n\\ndef _reverse_repeat_tuple(t, n):\\n    r\\\"\\\"\\\"Reverse the order of `t` and repeat each element for `n` times.\\n\\n    This can be used to translate padding arg used by Conv and Pooling modules\\n    to the ones used by `F.pad`.\\n    \\\"\\\"\\\"\\n    return tuple(x for x in reversed(t) for _ in range(n))\\n\\n\\ndef _list_with_default(out_size: List[int], defaults: List[int]) -> List[int]:\\n    import torch\\n\\n    if isinstance(out_size, (int, torch.SymInt)):\\n        return out_size\\n    if len(defaults) <= len(out_size):\\n        raise ValueError(f\\\"Input dimension should be at least {len(out_size) + 1}\\\")\\n    return [\\n        v if v is not None else d for v, d in zip(out_size, defaults[-len(out_size) :])\\n    ]\\n\\n\\ndef consume_prefix_in_state_dict_if_present(\\n    state_dict: Dict[str, Any],\\n    prefix: str,\\n) -> None:\\n    r\\\"\\\"\\\"Strip the prefix in state_dict in place, if any.\\n\\n    ..note::\\n        Given a `state_dict` from a DP/DDP model, a local model can load it by applying\\n        `consume_prefix_in_state_dict_if_present(state_dict, \\\"module.\\\")` before calling\\n        :meth:`torch.nn.Module.load_state_dict`.\\n\\n    Args:\\n        state_dict (OrderedDict): a state-dict to be loaded to the model.\\n        prefix (str): prefix.\\n    \\\"\\\"\\\"\\n    keys = list(state_dict.keys())\\n    for key in keys:\\n        if key.startswith(prefix):\\n            newkey = key[len(prefix) :]\\n            state_dict[newkey] = state_dict.pop(key)\\n\\n    # also strip the prefix in metadata if any.\\n    if hasattr(state_dict, \\\"_metadata\\\"):\\n        keys = list(state_dict._metadata.keys())\\n        for key in keys:\\n            # for the metadata dict, the key can be:\\n            # '': for the DDP module, which we want to remove.\\n            # 'module': for the actual model.\\n            # 'module.xx.xx': for the rest.\\n            if len(key) == 0:\\n                continue\\n            # handling both, 'module' case and  'module.' cases\\n            if key == prefix.replace(\\\".\\\", \\\"\\\") or key.startswith(prefix):\\n                newkey = key[len(prefix) :]\\n                state_dict._metadata[newkey] = state_dict._metadata.pop(key)\\n\\n\\n# mypy: allow-untyped-defs\\n\\nimport warnings\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\n\\nfrom .batchnorm import _LazyNormBase, _NormBase\\n\\n\\n__all__ = [\\n    \\\"InstanceNorm1d\\\",\\n    \\\"InstanceNorm2d\\\",\\n    \\\"InstanceNorm3d\\\",\\n    \\\"LazyInstanceNorm1d\\\",\\n    \\\"LazyInstanceNorm2d\\\",\\n    \\\"LazyInstanceNorm3d\\\",\\n]\\n\\n\\nclass _InstanceNorm(_NormBase):\\n    def __init__(\\n        self,\\n        num_features: int,\\n        eps: float = 1e-5,\\n        momentum: float = 0.1,\\n        affine: bool = False,\\n        track_running_stats: bool = False,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs\\n        )\\n\\n    def _check_input_dim(self, input):\\n        raise NotImplementedError\\n\\n    def _get_no_batch_dim(self):\\n        raise NotImplementedError\\n\\n    def _handle_no_batch_input(self, input):\\n        return self._apply_instance_norm(input.unsqueeze(0)).squeeze(0)\\n\\n    def _apply_instance_norm(self, input):\\n        return F.instance_norm(\\n            input,\\n            self.running_mean,\\n            self.running_var,\\n            self.weight,\\n            self.bias,\\n            self.training or not self.track_running_stats,\\n            self.momentum if self.momentum is not None else 0.0,\\n            self.eps,\\n        )\\n\\n    def _load_from_state_dict(\\n        self,\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ):\\n        version = local_metadata.get(\\\"version\\\", None)\\n        # at version 1: removed running_mean and running_var when\\n        # track_running_stats=False (default)\\n        if version is None and not self.track_running_stats:\\n            running_stats_keys = []\\n            for name in (\\\"running_mean\\\", \\\"running_var\\\"):\\n                key = prefix + name\\n                if key in state_dict:\\n                    running_stats_keys.append(key)\\n            if len(running_stats_keys) > 0:\\n                error_msgs.append(\\n                    \\\"Unexpected running stats buffer(s) {names} for {klass} \\\"\\n                    \\\"with track_running_stats=False. If state_dict is a \\\"\\n                    \\\"checkpoint saved before 0.4.0, this may be expected \\\"\\n                    \\\"because {klass} does not track running stats by default \\\"\\n                    \\\"since 0.4.0. Please remove these keys from state_dict. If \\\"\\n                    \\\"the running stats are actually needed, instead set \\\"\\n                    \\\"track_running_stats=True in {klass} to enable them. See \\\"\\n                    \\\"the documentation of {klass} for details.\\\".format(\\n                        names=\\\" and \\\".join(f'\\\"{k}\\\"' for k in running_stats_keys),\\n                        klass=self.__class__.__name__,\\n                    )\\n                )\\n                for key in running_stats_keys:\\n                    state_dict.pop(key)\\n\\n        super()._load_from_state_dict(\\n            state_dict,\\n            prefix,\\n            local_metadata,\\n            strict,\\n            missing_keys,\\n            unexpected_keys,\\n            error_msgs,\\n        )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        self._check_input_dim(input)\\n\\n        feature_dim = input.dim() - self._get_no_batch_dim()\\n        if input.size(feature_dim) != self.num_features:\\n            if self.affine:\\n                raise ValueError(\\n                    f\\\"expected input's size at dim={feature_dim} to match num_features\\\"\\n                    f\\\" ({self.num_features}), but got: {input.size(feature_dim)}.\\\"\\n                )\\n            else:\\n                warnings.warn(\\n                    f\\\"input's size at dim={feature_dim} does not match num_features. \\\"\\n                    \\\"You can silence this warning by not passing in num_features, \\\"\\n                    \\\"which is not used because affine=False\\\"\\n                )\\n\\n        if input.dim() == self._get_no_batch_dim():\\n            return self._handle_no_batch_input(input)\\n\\n        return self._apply_instance_norm(input)\\n\\n\\nclass InstanceNorm1d(_InstanceNorm):\\n    r\\\"\\\"\\\"Applies Instance Normalization.\\n\\n    This operation applies Instance Normalization\\n    over a 2D (unbatched) or 3D (batched) input as described in the paper\\n    `Instance Normalization: The Missing Ingredient for Fast Stylization\\n    <https://arxiv.org/abs/1607.08022>`__.\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension separately\\n    for each object in a mini-batch. :math:`\\\\gamma` and :math:`\\\\beta` are learnable parameter vectors\\n    of size `C` (where `C` is the number of features or channels of the input) if :attr:`affine` is ``True``.\\n    The standard-deviation is calculated via the biased estimator, equivalent to\\n    `torch.var(input, unbiased=False)`.\\n\\n    By default, this layer uses instance statistics computed from input data in\\n    both training and evaluation modes.\\n\\n    If :attr:`track_running_stats` is set to ``True``, during training this\\n    layer keeps running estimates of its computed mean and variance, which are\\n    then used for normalization during evaluation. The running estimates are\\n    kept with a default :attr:`momentum` of 0.1.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    .. note::\\n        :class:`InstanceNorm1d` and :class:`LayerNorm` are very similar, but\\n        have some subtle differences. :class:`InstanceNorm1d` is applied\\n        on each channel of channeled data like multidimensional time series, but\\n        :class:`LayerNorm` is usually applied on entire sample and often in NLP\\n        tasks. Additionally, :class:`LayerNorm` applies elementwise affine\\n        transform, while :class:`InstanceNorm1d` usually don't apply affine\\n        transform.\\n\\n    Args:\\n        num_features: number of features or channels :math:`C` of the input\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        momentum: the value used for the running_mean and running_var computation. Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters, initialized the same way as done for batch normalization.\\n            Default: ``False``.\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics and always uses batch\\n            statistics in both training and eval modes. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, L)` or :math:`(C, L)`\\n        - Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.InstanceNorm1d(100)\\n        >>> # With Learnable Parameters\\n        >>> m = nn.InstanceNorm1d(100, affine=True)\\n        >>> input = torch.randn(20, 100, 40)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def _get_no_batch_dim(self):\\n        return 2\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() not in (2, 3):\\n            raise ValueError(f\\\"expected 2D or 3D input (got {input.dim()}D input)\\\")\\n\\n\\nclass LazyInstanceNorm1d(_LazyNormBase, _InstanceNorm):\\n    r\\\"\\\"\\\"A :class:`torch.nn.InstanceNorm1d` module with lazy initialization of the ``num_features`` argument.\\n\\n    The ``num_features`` argument of the :class:`InstanceNorm1d` is inferred from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight`, `bias`, `running_mean` and `running_var`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, L)` or :math:`(C, L)`\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        momentum: the value used for the running_mean and running_var computation. Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters, initialized the same way as done for batch normalization.\\n            Default: ``False``.\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics and always uses batch\\n            statistics in both training and eval modes. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, L)` or :math:`(C, L)`\\n        - Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input)\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = InstanceNorm1d  # type: ignore[assignment]\\n\\n    def _get_no_batch_dim(self):\\n        return 2\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() not in (2, 3):\\n            raise ValueError(f\\\"expected 2D or 3D input (got {input.dim()}D input)\\\")\\n\\n\\nclass InstanceNorm2d(_InstanceNorm):\\n    r\\\"\\\"\\\"Applies Instance Normalization.\\n\\n    This operation applies Instance Normalization\\n    over a 4D input (a mini-batch of 2D inputs\\n    with additional channel dimension) as described in the paper\\n    `Instance Normalization: The Missing Ingredient for Fast Stylization\\n    <https://arxiv.org/abs/1607.08022>`__.\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension separately\\n    for each object in a mini-batch. :math:`\\\\gamma` and :math:`\\\\beta` are learnable parameter vectors\\n    of size `C` (where `C` is the input size) if :attr:`affine` is ``True``.\\n    The standard-deviation is calculated via the biased estimator, equivalent to\\n    `torch.var(input, unbiased=False)`.\\n\\n    By default, this layer uses instance statistics computed from input data in\\n    both training and evaluation modes.\\n\\n    If :attr:`track_running_stats` is set to ``True``, during training this\\n    layer keeps running estimates of its computed mean and variance, which are\\n    then used for normalization during evaluation. The running estimates are\\n    kept with a default :attr:`momentum` of 0.1.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    .. note::\\n        :class:`InstanceNorm2d` and :class:`LayerNorm` are very similar, but\\n        have some subtle differences. :class:`InstanceNorm2d` is applied\\n        on each channel of channeled data like RGB images, but\\n        :class:`LayerNorm` is usually applied on entire sample and often in NLP\\n        tasks. Additionally, :class:`LayerNorm` applies elementwise affine\\n        transform, while :class:`InstanceNorm2d` usually don't apply affine\\n        transform.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, H, W)` or :math:`(C, H, W)`\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        momentum: the value used for the running_mean and running_var computation. Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters, initialized the same way as done for batch normalization.\\n            Default: ``False``.\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics and always uses batch\\n            statistics in both training and eval modes. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`\\n        - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.InstanceNorm2d(100)\\n        >>> # With Learnable Parameters\\n        >>> m = nn.InstanceNorm2d(100, affine=True)\\n        >>> input = torch.randn(20, 100, 35, 45)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def _get_no_batch_dim(self):\\n        return 3\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() not in (3, 4):\\n            raise ValueError(f\\\"expected 3D or 4D input (got {input.dim()}D input)\\\")\\n\\n\\nclass LazyInstanceNorm2d(_LazyNormBase, _InstanceNorm):\\n    r\\\"\\\"\\\"A :class:`torch.nn.InstanceNorm2d` module with lazy initialization of the ``num_features`` argument.\\n\\n    The ``num_features`` argument of the :class:`InstanceNorm2d` is inferred from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight`, `bias`,\\n    `running_mean` and `running_var`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, H, W)` or :math:`(C, H, W)`\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        momentum: the value used for the running_mean and running_var computation. Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters, initialized the same way as done for batch normalization.\\n            Default: ``False``.\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics and always uses batch\\n            statistics in both training and eval modes. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`\\n        - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = InstanceNorm2d  # type: ignore[assignment]\\n\\n    def _get_no_batch_dim(self):\\n        return 3\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() not in (3, 4):\\n            raise ValueError(f\\\"expected 3D or 4D input (got {input.dim()}D input)\\\")\\n\\n\\nclass InstanceNorm3d(_InstanceNorm):\\n    r\\\"\\\"\\\"Applies Instance Normalization.\\n\\n    This operation applies Instance Normalization\\n    over a 5D input (a mini-batch of 3D inputs with additional channel dimension) as described in the paper\\n    `Instance Normalization: The Missing Ingredient for Fast Stylization\\n    <https://arxiv.org/abs/1607.08022>`__.\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension separately\\n    for each object in a mini-batch. :math:`\\\\gamma` and :math:`\\\\beta` are learnable parameter vectors\\n    of size C (where C is the input size) if :attr:`affine` is ``True``.\\n    The standard-deviation is calculated via the biased estimator, equivalent to\\n    `torch.var(input, unbiased=False)`.\\n\\n    By default, this layer uses instance statistics computed from input data in\\n    both training and evaluation modes.\\n\\n    If :attr:`track_running_stats` is set to ``True``, during training this\\n    layer keeps running estimates of its computed mean and variance, which are\\n    then used for normalization during evaluation. The running estimates are\\n    kept with a default :attr:`momentum` of 0.1.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    .. note::\\n        :class:`InstanceNorm3d` and :class:`LayerNorm` are very similar, but\\n        have some subtle differences. :class:`InstanceNorm3d` is applied\\n        on each channel of channeled data like 3D models with RGB color, but\\n        :class:`LayerNorm` is usually applied on entire sample and often in NLP\\n        tasks. Additionally, :class:`LayerNorm` applies elementwise affine\\n        transform, while :class:`InstanceNorm3d` usually don't apply affine\\n        transform.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        momentum: the value used for the running_mean and running_var computation. Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters, initialized the same way as done for batch normalization.\\n            Default: ``False``.\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics and always uses batch\\n            statistics in both training and eval modes. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`\\n        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.InstanceNorm3d(100)\\n        >>> # With Learnable Parameters\\n        >>> m = nn.InstanceNorm3d(100, affine=True)\\n        >>> input = torch.randn(20, 100, 35, 45, 10)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def _get_no_batch_dim(self):\\n        return 4\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() not in (4, 5):\\n            raise ValueError(f\\\"expected 4D or 5D input (got {input.dim()}D input)\\\")\\n\\n\\nclass LazyInstanceNorm3d(_LazyNormBase, _InstanceNorm):\\n    r\\\"\\\"\\\"A :class:`torch.nn.InstanceNorm3d` module with lazy initialization of the ``num_features`` argument.\\n\\n    The ``num_features`` argument of the :class:`InstanceNorm3d` is inferred from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight`, `bias`,\\n    `running_mean` and `running_var`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`\\n        eps: a value added to the denominator for numerical stability. Default: 1e-5\\n        momentum: the value used for the running_mean and running_var computation. Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters, initialized the same way as done for batch normalization.\\n            Default: ``False``.\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics and always uses batch\\n            statistics in both training and eval modes. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`\\n        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input)\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = InstanceNorm3d  # type: ignore[assignment]\\n\\n    def _get_no_batch_dim(self):\\n        return 4\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() not in (4, 5):\\n            raise ValueError(f\\\"expected 4D or 5D input (got {input.dim()}D input)\\\")\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Tuple, Union\\n\\nfrom torch import Tensor\\nfrom torch.types import _size\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"Flatten\\\", \\\"Unflatten\\\"]\\n\\n\\nclass Flatten(Module):\\n    r\\\"\\\"\\\"\\n    Flattens a contiguous range of dims into a tensor.\\n\\n    For use with :class:`~nn.Sequential`, see :meth:`torch.flatten` for details.\\n\\n    Shape:\\n        - Input: :math:`(*, S_{\\\\text{start}},..., S_{i}, ..., S_{\\\\text{end}}, *)`,'\\n          where :math:`S_{i}` is the size at dimension :math:`i` and :math:`*` means any\\n          number of dimensions including none.\\n        - Output: :math:`(*, \\\\prod_{i=\\\\text{start}}^{\\\\text{end}} S_{i}, *)`.\\n\\n    Args:\\n        start_dim: first dim to flatten (default = 1).\\n        end_dim: last dim to flatten (default = -1).\\n\\n    Examples::\\n        >>> input = torch.randn(32, 1, 5, 5)\\n        >>> # With default parameters\\n        >>> m = nn.Flatten()\\n        >>> output = m(input)\\n        >>> output.size()\\n        torch.Size([32, 25])\\n        >>> # With non-default parameters\\n        >>> m = nn.Flatten(0, 2)\\n        >>> output = m(input)\\n        >>> output.size()\\n        torch.Size([160, 5])\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"start_dim\\\", \\\"end_dim\\\"]\\n    start_dim: int\\n    end_dim: int\\n\\n    def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:\\n        super().__init__()\\n        self.start_dim = start_dim\\n        self.end_dim = end_dim\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return input.flatten(self.start_dim, self.end_dim)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"start_dim={self.start_dim}, end_dim={self.end_dim}\\\"\\n\\n\\nclass Unflatten(Module):\\n    r\\\"\\\"\\\"\\n    Unflattens a tensor dim expanding it to a desired shape. For use with :class:`~nn.Sequential`.\\n\\n    * :attr:`dim` specifies the dimension of the input tensor to be unflattened, and it can\\n      be either `int` or `str` when `Tensor` or `NamedTensor` is used, respectively.\\n\\n    * :attr:`unflattened_size` is the new shape of the unflattened dimension of the tensor and it can be\\n      a `tuple` of ints or a `list` of ints or `torch.Size` for `Tensor` input;  a `NamedShape`\\n      (tuple of `(name, size)` tuples) for `NamedTensor` input.\\n\\n    Shape:\\n        - Input: :math:`(*, S_{\\\\text{dim}}, *)`, where :math:`S_{\\\\text{dim}}` is the size at\\n          dimension :attr:`dim` and :math:`*` means any number of dimensions including none.\\n        - Output: :math:`(*, U_1, ..., U_n, *)`, where :math:`U` = :attr:`unflattened_size` and\\n          :math:`\\\\prod_{i=1}^n U_i = S_{\\\\text{dim}}`.\\n\\n    Args:\\n        dim (Union[int, str]): Dimension to be unflattened\\n        unflattened_size (Union[torch.Size, Tuple, List, NamedShape]): New shape of the unflattened dimension\\n\\n    Examples:\\n        >>> input = torch.randn(2, 50)\\n        >>> # With tuple of ints\\n        >>> m = nn.Sequential(\\n        >>>     nn.Linear(50, 50),\\n        >>>     nn.Unflatten(1, (2, 5, 5))\\n        >>> )\\n        >>> output = m(input)\\n        >>> output.size()\\n        torch.Size([2, 2, 5, 5])\\n        >>> # With torch.Size\\n        >>> m = nn.Sequential(\\n        >>>     nn.Linear(50, 50),\\n        >>>     nn.Unflatten(1, torch.Size([2, 5, 5]))\\n        >>> )\\n        >>> output = m(input)\\n        >>> output.size()\\n        torch.Size([2, 2, 5, 5])\\n        >>> # With namedshape (tuple of tuples)\\n        >>> input = torch.randn(2, 50, names=('N', 'features'))\\n        >>> unflatten = nn.Unflatten('features', (('C', 2), ('H', 5), ('W', 5)))\\n        >>> output = unflatten(input)\\n        >>> output.size()\\n        torch.Size([2, 2, 5, 5])\\n    \\\"\\\"\\\"\\n\\n    NamedShape = Tuple[Tuple[str, int]]\\n\\n    __constants__ = [\\\"dim\\\", \\\"unflattened_size\\\"]\\n    dim: Union[int, str]\\n    unflattened_size: Union[_size, NamedShape]\\n\\n    def __init__(\\n        self, dim: Union[int, str], unflattened_size: Union[_size, NamedShape]\\n    ) -> None:\\n        super().__init__()\\n\\n        if isinstance(dim, int):\\n            self._require_tuple_int(unflattened_size)\\n        elif isinstance(dim, str):\\n            self._require_tuple_tuple(unflattened_size)\\n        else:\\n            raise TypeError(\\\"invalid argument type for dim parameter\\\")\\n\\n        self.dim = dim\\n        self.unflattened_size = unflattened_size\\n\\n    def _require_tuple_tuple(self, input):\\n        if isinstance(input, tuple):\\n            for idx, elem in enumerate(input):\\n                if not isinstance(elem, tuple):\\n                    raise TypeError(\\n                        \\\"unflattened_size must be tuple of tuples, \\\"\\n                        + f\\\"but found element of type {type(elem).__name__} at pos {idx}\\\"\\n                    )\\n            return\\n        raise TypeError(\\n            \\\"unflattened_size must be a tuple of tuples, \\\"\\n            + f\\\"but found type {type(input).__name__}\\\"\\n        )\\n\\n    def _require_tuple_int(self, input):\\n        if isinstance(input, (tuple, list)):\\n            for idx, elem in enumerate(input):\\n                if not isinstance(elem, int):\\n                    raise TypeError(\\n                        \\\"unflattened_size must be tuple of ints, \\\"\\n                        + f\\\"but found element of type {type(elem).__name__} at pos {idx}\\\"\\n                    )\\n            return\\n        raise TypeError(\\n            f\\\"unflattened_size must be a tuple of ints, but found type {type(input).__name__}\\\"\\n        )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return input.unflatten(self.dim, self.unflattened_size)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"dim={self.dim}, unflattened_size={self.unflattened_size}\\\"\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Any, Optional\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch.nn import functional as F, init\\nfrom torch.nn.parameter import Parameter, UninitializedBuffer, UninitializedParameter\\n\\nfrom ._functions import SyncBatchNorm as sync_batch_norm\\nfrom .lazy import LazyModuleMixin\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"BatchNorm1d\\\",\\n    \\\"LazyBatchNorm1d\\\",\\n    \\\"BatchNorm2d\\\",\\n    \\\"LazyBatchNorm2d\\\",\\n    \\\"BatchNorm3d\\\",\\n    \\\"LazyBatchNorm3d\\\",\\n    \\\"SyncBatchNorm\\\",\\n]\\n\\n\\nclass _NormBase(Module):\\n    \\\"\\\"\\\"Common base of _InstanceNorm and _BatchNorm.\\\"\\\"\\\"\\n\\n    _version = 2\\n    __constants__ = [\\\"track_running_stats\\\", \\\"momentum\\\", \\\"eps\\\", \\\"num_features\\\", \\\"affine\\\"]\\n    num_features: int\\n    eps: float\\n    momentum: Optional[float]\\n    affine: bool\\n    track_running_stats: bool\\n    # WARNING: weight and bias purposely not defined here.\\n    # See https://github.com/pytorch/pytorch/issues/39670\\n\\n    def __init__(\\n        self,\\n        num_features: int,\\n        eps: float = 1e-5,\\n        momentum: Optional[float] = 0.1,\\n        affine: bool = True,\\n        track_running_stats: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.num_features = num_features\\n        self.eps = eps\\n        self.momentum = momentum\\n        self.affine = affine\\n        self.track_running_stats = track_running_stats\\n        if self.affine:\\n            self.weight = Parameter(torch.empty(num_features, **factory_kwargs))\\n            self.bias = Parameter(torch.empty(num_features, **factory_kwargs))\\n        else:\\n            self.register_parameter(\\\"weight\\\", None)\\n            self.register_parameter(\\\"bias\\\", None)\\n        if self.track_running_stats:\\n            self.register_buffer(\\n                \\\"running_mean\\\", torch.zeros(num_features, **factory_kwargs)\\n            )\\n            self.register_buffer(\\n                \\\"running_var\\\", torch.ones(num_features, **factory_kwargs)\\n            )\\n            self.running_mean: Optional[Tensor]\\n            self.running_var: Optional[Tensor]\\n            self.register_buffer(\\n                \\\"num_batches_tracked\\\",\\n                torch.tensor(\\n                    0,\\n                    dtype=torch.long,\\n                    **{k: v for k, v in factory_kwargs.items() if k != \\\"dtype\\\"},\\n                ),\\n            )\\n            self.num_batches_tracked: Optional[Tensor]\\n        else:\\n            self.register_buffer(\\\"running_mean\\\", None)\\n            self.register_buffer(\\\"running_var\\\", None)\\n            self.register_buffer(\\\"num_batches_tracked\\\", None)\\n        self.reset_parameters()\\n\\n    def reset_running_stats(self) -> None:\\n        if self.track_running_stats:\\n            # running_mean/running_var/num_batches... are registered at runtime depending\\n            # if self.track_running_stats is on\\n            self.running_mean.zero_()  # type: ignore[union-attr]\\n            self.running_var.fill_(1)  # type: ignore[union-attr]\\n            self.num_batches_tracked.zero_()  # type: ignore[union-attr,operator]\\n\\n    def reset_parameters(self) -> None:\\n        self.reset_running_stats()\\n        if self.affine:\\n            init.ones_(self.weight)\\n            init.zeros_(self.bias)\\n\\n    def _check_input_dim(self, input):\\n        raise NotImplementedError\\n\\n    def extra_repr(self):\\n        return (\\n            \\\"{num_features}, eps={eps}, momentum={momentum}, affine={affine}, \\\"\\n            \\\"track_running_stats={track_running_stats}\\\".format(**self.__dict__)\\n        )\\n\\n    def _load_from_state_dict(\\n        self,\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ):\\n        version = local_metadata.get(\\\"version\\\", None)\\n\\n        if (version is None or version < 2) and self.track_running_stats:\\n            # at version 2: added num_batches_tracked buffer\\n            #               this should have a default value of 0\\n            num_batches_tracked_key = prefix + \\\"num_batches_tracked\\\"\\n            if num_batches_tracked_key not in state_dict:\\n                state_dict[num_batches_tracked_key] = (\\n                    self.num_batches_tracked\\n                    if self.num_batches_tracked is not None\\n                    and self.num_batches_tracked.device != torch.device(\\\"meta\\\")\\n                    else torch.tensor(0, dtype=torch.long)\\n                )\\n\\n        super()._load_from_state_dict(\\n            state_dict,\\n            prefix,\\n            local_metadata,\\n            strict,\\n            missing_keys,\\n            unexpected_keys,\\n            error_msgs,\\n        )\\n\\n\\nclass _BatchNorm(_NormBase):\\n    def __init__(\\n        self,\\n        num_features: int,\\n        eps: float = 1e-5,\\n        momentum: Optional[float] = 0.1,\\n        affine: bool = True,\\n        track_running_stats: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs\\n        )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        self._check_input_dim(input)\\n\\n        # exponential_average_factor is set to self.momentum\\n        # (when it is available) only so that it gets updated\\n        # in ONNX graph when this node is exported to ONNX.\\n        if self.momentum is None:\\n            exponential_average_factor = 0.0\\n        else:\\n            exponential_average_factor = self.momentum\\n\\n        if self.training and self.track_running_stats:\\n            # TODO: if statement only here to tell the jit to skip emitting this when it is None\\n            if self.num_batches_tracked is not None:  # type: ignore[has-type]\\n                self.num_batches_tracked.add_(1)  # type: ignore[has-type]\\n                if self.momentum is None:  # use cumulative moving average\\n                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)\\n                else:  # use exponential moving average\\n                    exponential_average_factor = self.momentum\\n\\n        r\\\"\\\"\\\"\\n        Decide whether the mini-batch stats should be used for normalization rather than the buffers.\\n        Mini-batch stats are used in training mode, and in eval mode when buffers are None.\\n        \\\"\\\"\\\"\\n        if self.training:\\n            bn_training = True\\n        else:\\n            bn_training = (self.running_mean is None) and (self.running_var is None)\\n\\n        r\\\"\\\"\\\"\\n        Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be\\n        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are\\n        used for normalization (i.e. in eval mode when buffers are not None).\\n        \\\"\\\"\\\"\\n        return F.batch_norm(\\n            input,\\n            # If buffers are not to be tracked, ensure that they won't be updated\\n            self.running_mean\\n            if not self.training or self.track_running_stats\\n            else None,\\n            self.running_var if not self.training or self.track_running_stats else None,\\n            self.weight,\\n            self.bias,\\n            bn_training,\\n            exponential_average_factor,\\n            self.eps,\\n        )\\n\\n\\nclass _LazyNormBase(LazyModuleMixin, _NormBase):\\n    weight: UninitializedParameter  # type: ignore[assignment]\\n    bias: UninitializedParameter  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        eps=1e-5,\\n        momentum=0.1,\\n        affine=True,\\n        track_running_stats=True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            # affine and track_running_stats are hardcoded to False to\\n            # avoid creating tensors that will soon be overwritten.\\n            0,\\n            eps,\\n            momentum,\\n            False,\\n            False,\\n            **factory_kwargs,\\n        )\\n        self.affine = affine\\n        self.track_running_stats = track_running_stats\\n        if self.affine:\\n            self.weight = UninitializedParameter(**factory_kwargs)\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n        if self.track_running_stats:\\n            self.running_mean = UninitializedBuffer(**factory_kwargs)\\n            self.running_var = UninitializedBuffer(**factory_kwargs)\\n            self.num_batches_tracked = torch.tensor(\\n                0,\\n                dtype=torch.long,\\n                **{k: v for k, v in factory_kwargs.items() if k != \\\"dtype\\\"},\\n            )\\n\\n    def reset_parameters(self) -> None:\\n        if not self.has_uninitialized_params() and self.num_features != 0:\\n            super().reset_parameters()\\n\\n    def initialize_parameters(self, input) -> None:  # type: ignore[override]\\n        if self.has_uninitialized_params():\\n            self.num_features = input.shape[1]\\n            if self.affine:\\n                assert isinstance(self.weight, UninitializedParameter)\\n                assert isinstance(self.bias, UninitializedParameter)\\n                self.weight.materialize((self.num_features,))\\n                self.bias.materialize((self.num_features,))\\n            if self.track_running_stats:\\n                self.running_mean.materialize(  # type:ignore[union-attr]\\n                    (self.num_features,)\\n                )\\n                self.running_var.materialize(  # type:ignore[union-attr]\\n                    (self.num_features,)\\n                )\\n            self.reset_parameters()\\n\\n\\nclass BatchNorm1d(_BatchNorm):\\n    r\\\"\\\"\\\"Applies Batch Normalization over a 2D or 3D input.\\n\\n    Method described in the paper\\n    `Batch Normalization: Accelerating Deep Network Training by Reducing\\n    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{\\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension over\\n    the mini-batches and :math:`\\\\gamma` and :math:`\\\\beta` are learnable parameter vectors\\n    of size `C` (where `C` is the number of features or channels of the input). By default, the\\n    elements of :math:`\\\\gamma` are set to 1 and the elements of :math:`\\\\beta` are set to 0.\\n    At train time in the forward pass, the standard-deviation is calculated via the biased estimator,\\n    equivalent to ``torch.var(input, unbiased=False)``. However, the value stored in the\\n    moving average of the standard-deviation is calculated via the unbiased  estimator, equivalent to\\n    ``torch.var(input, unbiased=True)``.\\n\\n    Also by default, during training this layer keeps running estimates of its\\n    computed mean and variance, which are then used for normalization during\\n    evaluation. The running estimates are kept with a default :attr:`momentum`\\n    of 0.1.\\n\\n    If :attr:`track_running_stats` is set to ``False``, this layer then does not\\n    keep running estimates, and batch statistics are instead used during\\n    evaluation time as well.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    Because the Batch Normalization is done over the `C` dimension, computing statistics\\n    on `(N, L)` slices, it's common terminology to call this Temporal Batch Normalization.\\n\\n    Args:\\n        num_features: number of features or channels :math:`C` of the input\\n        eps: a value added to the denominator for numerical stability.\\n            Default: 1e-5\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n\\n    Shape:\\n        - Input: :math:`(N, C)` or :math:`(N, C, L)`, where :math:`N` is the batch size,\\n          :math:`C` is the number of features or channels, and :math:`L` is the sequence length\\n        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # With Learnable Parameters\\n        >>> m = nn.BatchNorm1d(100)\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.BatchNorm1d(100, affine=False)\\n        >>> input = torch.randn(20, 100)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 2 and input.dim() != 3:\\n            raise ValueError(f\\\"expected 2D or 3D input (got {input.dim()}D input)\\\")\\n\\n\\nclass LazyBatchNorm1d(_LazyNormBase, _BatchNorm):\\n    r\\\"\\\"\\\"A :class:`torch.nn.BatchNorm1d` module with lazy initialization.\\n\\n    Lazy initialization based on the ``num_features`` argument of the :class:`BatchNorm1d` that is inferred\\n    from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight`, `bias`,\\n    `running_mean` and `running_var`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        eps: a value added to the denominator for numerical stability.\\n            Default: 1e-5\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = BatchNorm1d  # type: ignore[assignment]\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 2 and input.dim() != 3:\\n            raise ValueError(f\\\"expected 2D or 3D input (got {input.dim()}D input)\\\")\\n\\n\\nclass BatchNorm2d(_BatchNorm):\\n    r\\\"\\\"\\\"Applies Batch Normalization over a 4D input.\\n\\n    4D is a mini-batch of 2D inputs\\n    with additional channel dimension. Method described in the paper\\n    `Batch Normalization: Accelerating Deep Network Training by Reducing\\n    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension over\\n    the mini-batches and :math:`\\\\gamma` and :math:`\\\\beta` are learnable parameter vectors\\n    of size `C` (where `C` is the input size). By default, the elements of :math:`\\\\gamma` are set\\n    to 1 and the elements of :math:`\\\\beta` are set to 0. At train time in the forward pass, the\\n    standard-deviation is calculated via the biased estimator, equivalent to\\n    ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the\\n    standard-deviation is calculated via the unbiased  estimator, equivalent to\\n    ``torch.var(input, unbiased=True)``.\\n\\n    Also by default, during training this layer keeps running estimates of its\\n    computed mean and variance, which are then used for normalization during\\n    evaluation. The running estimates are kept with a default :attr:`momentum`\\n    of 0.1.\\n\\n    If :attr:`track_running_stats` is set to ``False``, this layer then does not\\n    keep running estimates, and batch statistics are instead used during\\n    evaluation time as well.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    Because the Batch Normalization is done over the `C` dimension, computing statistics\\n    on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, H, W)`\\n        eps: a value added to the denominator for numerical stability.\\n            Default: 1e-5\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n\\n    Shape:\\n        - Input: :math:`(N, C, H, W)`\\n        - Output: :math:`(N, C, H, W)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # With Learnable Parameters\\n        >>> m = nn.BatchNorm2d(100)\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.BatchNorm2d(100, affine=False)\\n        >>> input = torch.randn(20, 100, 35, 45)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 4:\\n            raise ValueError(f\\\"expected 4D input (got {input.dim()}D input)\\\")\\n\\n\\nclass LazyBatchNorm2d(_LazyNormBase, _BatchNorm):\\n    r\\\"\\\"\\\"A :class:`torch.nn.BatchNorm2d` module with lazy initialization.\\n\\n    Lazy initialization is done for the ``num_features`` argument of the :class:`BatchNorm2d` that is inferred\\n    from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight`, `bias`,\\n    `running_mean` and `running_var`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        eps: a value added to the denominator for numerical stability.\\n            Default: 1e-5\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = BatchNorm2d  # type: ignore[assignment]\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 4:\\n            raise ValueError(f\\\"expected 4D input (got {input.dim()}D input)\\\")\\n\\n\\nclass BatchNorm3d(_BatchNorm):\\n    r\\\"\\\"\\\"Applies Batch Normalization over a 5D input.\\n\\n    5D is a mini-batch of 3D inputs with additional channel dimension as described in the paper\\n    `Batch Normalization: Accelerating Deep Network Training by Reducing\\n    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension over\\n    the mini-batches and :math:`\\\\gamma` and :math:`\\\\beta` are learnable parameter vectors\\n    of size `C` (where `C` is the input size). By default, the elements of :math:`\\\\gamma` are set\\n    to 1 and the elements of :math:`\\\\beta` are set to 0. At train time in the forward pass, the\\n    standard-deviation is calculated via the biased estimator, equivalent to\\n    ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the\\n    standard-deviation is calculated via the unbiased  estimator, equivalent to\\n    ``torch.var(input, unbiased=True)``.\\n\\n    Also by default, during training this layer keeps running estimates of its\\n    computed mean and variance, which are then used for normalization during\\n    evaluation. The running estimates are kept with a default :attr:`momentum`\\n    of 0.1.\\n\\n    If :attr:`track_running_stats` is set to ``False``, this layer then does not\\n    keep running estimates, and batch statistics are instead used during\\n    evaluation time as well.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    Because the Batch Normalization is done over the `C` dimension, computing statistics\\n    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric Batch Normalization\\n    or Spatio-temporal Batch Normalization.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, D, H, W)`\\n        eps: a value added to the denominator for numerical stability.\\n            Default: 1e-5\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n\\n    Shape:\\n        - Input: :math:`(N, C, D, H, W)`\\n        - Output: :math:`(N, C, D, H, W)` (same shape as input)\\n\\n    Examples::\\n\\n        >>> # With Learnable Parameters\\n        >>> m = nn.BatchNorm3d(100)\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.BatchNorm3d(100, affine=False)\\n        >>> input = torch.randn(20, 100, 35, 45, 10)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 5:\\n            raise ValueError(f\\\"expected 5D input (got {input.dim()}D input)\\\")\\n\\n\\nclass LazyBatchNorm3d(_LazyNormBase, _BatchNorm):\\n    r\\\"\\\"\\\"A :class:`torch.nn.BatchNorm3d` module with lazy initialization.\\n\\n    Lazy initialization is done for the ``num_features`` argument of the :class:`BatchNorm3d` that is inferred\\n    from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight`, `bias`,\\n    `running_mean` and `running_var`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        eps: a value added to the denominator for numerical stability.\\n            Default: 1e-5\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n    \\\"\\\"\\\"\\n\\n    cls_to_become = BatchNorm3d  # type: ignore[assignment]\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 5:\\n            raise ValueError(f\\\"expected 5D input (got {input.dim()}D input)\\\")\\n\\n\\nclass SyncBatchNorm(_BatchNorm):\\n    r\\\"\\\"\\\"Applies Batch Normalization over a N-Dimensional input.\\n\\n    The N-D input is a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paper\\n    `Batch Normalization: Accelerating Deep Network Training by Reducing\\n    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .\\n\\n    .. math::\\n\\n        y = \\\\frac{x - \\\\mathrm{E}[x]}{ \\\\sqrt{\\\\mathrm{Var}[x] + \\\\epsilon}} * \\\\gamma + \\\\beta\\n\\n    The mean and standard-deviation are calculated per-dimension over all\\n    mini-batches of the same process groups. :math:`\\\\gamma` and :math:`\\\\beta`\\n    are learnable parameter vectors of size `C` (where `C` is the input size).\\n    By default, the elements of :math:`\\\\gamma` are sampled from\\n    :math:`\\\\mathcal{U}(0, 1)` and the elements of :math:`\\\\beta` are set to 0.\\n    The standard-deviation is calculated via the biased estimator, equivalent to\\n    `torch.var(input, unbiased=False)`.\\n\\n    Also by default, during training this layer keeps running estimates of its\\n    computed mean and variance, which are then used for normalization during\\n    evaluation. The running estimates are kept with a default :attr:`momentum`\\n    of 0.1.\\n\\n    If :attr:`track_running_stats` is set to ``False``, this layer then does not\\n    keep running estimates, and batch statistics are instead used during\\n    evaluation time as well.\\n\\n    .. note::\\n        This :attr:`momentum` argument is different from one used in optimizer\\n        classes and the conventional notion of momentum. Mathematically, the\\n        update rule for running statistics here is\\n        :math:`\\\\hat{x}_\\\\text{new} = (1 - \\\\text{momentum}) \\\\times \\\\hat{x} + \\\\text{momentum} \\\\times x_t`,\\n        where :math:`\\\\hat{x}` is the estimated statistic and :math:`x_t` is the\\n        new observed value.\\n\\n    Because the Batch Normalization is done for each channel in the ``C`` dimension, computing\\n    statistics on ``(N, +)`` slices, it's common terminology to call this Volumetric Batch\\n    Normalization or Spatio-temporal Batch Normalization.\\n\\n    Currently :class:`SyncBatchNorm` only supports\\n    :class:`~torch.nn.DistributedDataParallel` (DDP) with single GPU per process. Use\\n    :meth:`torch.nn.SyncBatchNorm.convert_sync_batchnorm()` to convert\\n    :attr:`BatchNorm*D` layer to :class:`SyncBatchNorm` before wrapping\\n    Network with DDP.\\n\\n    Args:\\n        num_features: :math:`C` from an expected input of size\\n            :math:`(N, C, +)`\\n        eps: a value added to the denominator for numerical stability.\\n            Default: ``1e-5``\\n        momentum: the value used for the running_mean and running_var\\n            computation. Can be set to ``None`` for cumulative moving average\\n            (i.e. simple average). Default: 0.1\\n        affine: a boolean value that when set to ``True``, this module has\\n            learnable affine parameters. Default: ``True``\\n        track_running_stats: a boolean value that when set to ``True``, this\\n            module tracks the running mean and variance, and when set to ``False``,\\n            this module does not track such statistics, and initializes statistics\\n            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.\\n            When these buffers are ``None``, this module always uses batch statistics.\\n            in both training and eval modes. Default: ``True``\\n        process_group: synchronization of stats happen within each process group\\n            individually. Default behavior is synchronization across the whole\\n            world\\n\\n    Shape:\\n        - Input: :math:`(N, C, +)`\\n        - Output: :math:`(N, C, +)` (same shape as input)\\n\\n    .. note::\\n        Synchronization of batchnorm statistics occurs only while training, i.e.\\n        synchronization is disabled when ``model.eval()`` is set or if\\n        ``self.training`` is otherwise ``False``.\\n\\n    Examples::\\n\\n        >>> # xdoctest: +SKIP\\n        >>> # With Learnable Parameters\\n        >>> m = nn.SyncBatchNorm(100)\\n        >>> # creating process group (optional)\\n        >>> # ranks is a list of int identifying rank ids.\\n        >>> ranks = list(range(8))\\n        >>> r1, r2 = ranks[:4], ranks[4:]\\n        >>> # Note: every rank calls into new_group for every\\n        >>> # process group created, even if that rank is not\\n        >>> # part of the group.\\n        >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]\\n        >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]\\n        >>> # Without Learnable Parameters\\n        >>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group)\\n        >>> input = torch.randn(20, 100, 35, 45, 10)\\n        >>> output = m(input)\\n\\n        >>> # network is nn.BatchNorm layer\\n        >>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group)\\n        >>> # only single gpu per process is currently supported\\n        >>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(\\n        >>>                         sync_bn_network,\\n        >>>                         device_ids=[args.local_rank],\\n        >>>                         output_device=args.local_rank)\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        num_features: int,\\n        eps: float = 1e-5,\\n        momentum: Optional[float] = 0.1,\\n        affine: bool = True,\\n        track_running_stats: bool = True,\\n        process_group: Optional[Any] = None,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs\\n        )\\n        self.process_group = process_group\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() < 2:\\n            raise ValueError(f\\\"expected at least 2D input (got {input.dim()}D input)\\\")\\n\\n    def _check_non_zero_input_channels(self, input):\\n        if input.size(1) == 0:\\n            raise ValueError(\\n                \\\"SyncBatchNorm number of input channels should be non-zero\\\"\\n            )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        self._check_input_dim(input)\\n        self._check_non_zero_input_channels(input)\\n\\n        # exponential_average_factor is set to self.momentum\\n        # (when it is available) only so that it gets updated\\n        # in ONNX graph when this node is exported to ONNX.\\n        if self.momentum is None:\\n            exponential_average_factor = 0.0\\n        else:\\n            exponential_average_factor = self.momentum\\n\\n        if self.training and self.track_running_stats:\\n            assert self.num_batches_tracked is not None\\n            self.num_batches_tracked.add_(1)\\n            if self.momentum is None:  # use cumulative moving average\\n                exponential_average_factor = 1.0 / self.num_batches_tracked.item()\\n            else:  # use exponential moving average\\n                exponential_average_factor = self.momentum\\n\\n        r\\\"\\\"\\\"\\n        Decide whether the mini-batch stats should be used for normalization rather than the buffers.\\n        Mini-batch stats are used in training mode, and in eval mode when buffers are None.\\n        \\\"\\\"\\\"\\n        if self.training:\\n            bn_training = True\\n        else:\\n            bn_training = (self.running_mean is None) and (self.running_var is None)\\n\\n        r\\\"\\\"\\\"\\n        Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be\\n        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are\\n        used for normalization (i.e. in eval mode when buffers are not None).\\n        \\\"\\\"\\\"\\n        # If buffers are not to be tracked, ensure that they won't be updated\\n        running_mean = (\\n            self.running_mean if not self.training or self.track_running_stats else None\\n        )\\n        running_var = (\\n            self.running_var if not self.training or self.track_running_stats else None\\n        )\\n\\n        # Don't sync batchnorm stats in inference mode (model.eval()).\\n        need_sync = (\\n            bn_training\\n            and self.training\\n            and torch.distributed.is_available()\\n            and torch.distributed.is_initialized()\\n        )\\n        if need_sync:\\n            # currently only GPU/PrivateUse1 input is supported\\n            if input.device.type not in [\\n                \\\"cuda\\\",\\n                torch._C._get_privateuse1_backend_name(),\\n            ]:\\n                raise ValueError(\\n                    \\\"SyncBatchNorm expected input tensor to be on GPU or \\\"\\n                    f\\\"{torch._C._get_privateuse1_backend_name()}\\\"\\n                )\\n\\n            process_group = torch.distributed.group.WORLD\\n            if self.process_group:\\n                process_group = self.process_group\\n            world_size = torch.distributed.get_world_size(process_group)\\n            need_sync = world_size > 1\\n\\n        # fallback to framework BN when synchronization is not necessary\\n        if not need_sync:\\n            return F.batch_norm(\\n                input,\\n                running_mean,\\n                running_var,\\n                self.weight,\\n                self.bias,\\n                bn_training,\\n                exponential_average_factor,\\n                self.eps,\\n            )\\n        else:\\n            assert bn_training\\n            return sync_batch_norm.apply(\\n                input,\\n                self.weight,\\n                self.bias,\\n                running_mean,\\n                running_var,\\n                self.eps,\\n                exponential_average_factor,\\n                process_group,  # type: ignore[possibly-undefined]\\n                world_size,  # type: ignore[possibly-undefined]\\n            )\\n\\n    @classmethod\\n    def convert_sync_batchnorm(cls, module, process_group=None):\\n        r\\\"\\\"\\\"Converts all :attr:`BatchNorm*D` layers in the model to :class:`torch.nn.SyncBatchNorm` layers.\\n\\n        Args:\\n            module (nn.Module): module containing one or more :attr:`BatchNorm*D` layers\\n            process_group (optional): process group to scope synchronization,\\n                default is the whole world\\n\\n        Returns:\\n            The original :attr:`module` with the converted :class:`torch.nn.SyncBatchNorm`\\n            layers. If the original :attr:`module` is a :attr:`BatchNorm*D` layer,\\n            a new :class:`torch.nn.SyncBatchNorm` layer object will be returned\\n            instead.\\n\\n        Example::\\n\\n            >>> # Network with nn.BatchNorm layer\\n            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)\\n            >>> module = torch.nn.Sequential(\\n            >>>            torch.nn.Linear(20, 100),\\n            >>>            torch.nn.BatchNorm1d(100),\\n            >>>          ).cuda()\\n            >>> # creating process group (optional)\\n            >>> # ranks is a list of int identifying rank ids.\\n            >>> ranks = list(range(8))\\n            >>> r1, r2 = ranks[:4], ranks[4:]\\n            >>> # Note: every rank calls into new_group for every\\n            >>> # process group created, even if that rank is not\\n            >>> # part of the group.\\n            >>> # xdoctest: +SKIP(\\\"distributed\\\")\\n            >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]\\n            >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]\\n            >>> sync_bn_module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)\\n\\n        \\\"\\\"\\\"\\n        module_output = module\\n        if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):\\n            module_output = torch.nn.SyncBatchNorm(\\n                module.num_features,\\n                module.eps,\\n                module.momentum,\\n                module.affine,\\n                module.track_running_stats,\\n                process_group,\\n            )\\n            if module.affine:\\n                with torch.no_grad():\\n                    module_output.weight = module.weight\\n                    module_output.bias = module.bias\\n            module_output.running_mean = module.running_mean\\n            module_output.running_var = module.running_var\\n            module_output.num_batches_tracked = module.num_batches_tracked\\n            module_output.training = module.training\\n            if hasattr(module, \\\"qconfig\\\"):\\n                module_output.qconfig = module.qconfig\\n        for name, child in module.named_children():\\n            module_output.add_module(\\n                name, cls.convert_sync_batchnorm(child, process_group)\\n            )\\n        del module\\n        return module_output\\n\\n\\n# mypy: allow-untyped-defs\\nimport warnings\\nfrom typing import Optional, Tuple\\n\\nimport torch\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.init import constant_, xavier_normal_, xavier_uniform_\\nfrom torch.nn.parameter import Parameter\\n\\nfrom .linear import NonDynamicallyQuantizableLinear\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"Threshold\\\",\\n    \\\"ReLU\\\",\\n    \\\"RReLU\\\",\\n    \\\"Hardtanh\\\",\\n    \\\"ReLU6\\\",\\n    \\\"Sigmoid\\\",\\n    \\\"Hardsigmoid\\\",\\n    \\\"Tanh\\\",\\n    \\\"SiLU\\\",\\n    \\\"Mish\\\",\\n    \\\"Hardswish\\\",\\n    \\\"ELU\\\",\\n    \\\"CELU\\\",\\n    \\\"SELU\\\",\\n    \\\"GLU\\\",\\n    \\\"GELU\\\",\\n    \\\"Hardshrink\\\",\\n    \\\"LeakyReLU\\\",\\n    \\\"LogSigmoid\\\",\\n    \\\"Softplus\\\",\\n    \\\"Softshrink\\\",\\n    \\\"MultiheadAttention\\\",\\n    \\\"PReLU\\\",\\n    \\\"Softsign\\\",\\n    \\\"Tanhshrink\\\",\\n    \\\"Softmin\\\",\\n    \\\"Softmax\\\",\\n    \\\"Softmax2d\\\",\\n    \\\"LogSoftmax\\\",\\n]\\n\\n\\nclass Threshold(Module):\\n    r\\\"\\\"\\\"Thresholds each element of the input Tensor.\\n\\n    Threshold is defined as:\\n\\n    .. math::\\n        y =\\n        \\\\begin{cases}\\n        x, &\\\\text{ if } x > \\\\text{threshold} \\\\\\\\\\n        \\\\text{value}, &\\\\text{ otherwise }\\n        \\\\end{cases}\\n\\n    Args:\\n        threshold: The value to threshold at\\n        value: The value to replace with\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    Examples::\\n\\n        >>> m = nn.Threshold(0.1, 20)\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"threshold\\\", \\\"value\\\", \\\"inplace\\\"]\\n\\n    threshold: float\\n    value: float\\n    inplace: bool\\n\\n    def __init__(self, threshold: float, value: float, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.threshold = threshold\\n        self.value = value\\n        self.inplace = inplace\\n        # TODO: check in THNN (if inplace == True, then assert value <= threshold)\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.threshold(input, self.threshold, self.value, self.inplace)\\n\\n    def extra_repr(self):\\n        inplace_str = \\\", inplace=True\\\" if self.inplace else \\\"\\\"\\n        return f\\\"threshold={self.threshold}, value={self.value}{inplace_str}\\\"\\n\\n\\nclass ReLU(Module):\\n    r\\\"\\\"\\\"Applies the rectified linear unit function element-wise.\\n\\n    :math:`\\\\text{ReLU}(x) = (x)^+ = \\\\max(0, x)`\\n\\n    Args:\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/ReLU.png\\n\\n    Examples::\\n\\n        >>> m = nn.ReLU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n\\n\\n      An implementation of CReLU - https://arxiv.org/abs/1603.05201\\n\\n        >>> m = nn.ReLU()\\n        >>> input = torch.randn(2).unsqueeze(0)\\n        >>> output = torch.cat((m(input), m(-input)))\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\"]\\n    inplace: bool\\n\\n    def __init__(self, inplace: bool = False):\\n        super().__init__()\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.relu(input, inplace=self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\"inplace=True\\\" if self.inplace else \\\"\\\"\\n        return inplace_str\\n\\n\\nclass RReLU(Module):\\n    r\\\"\\\"\\\"Applies the randomized leaky rectified linear unit function, element-wise.\\n\\n    Method described in the paper:\\n    `Empirical Evaluation of Rectified Activations in Convolutional Network <https://arxiv.org/abs/1505.00853>`_.\\n\\n    The function is defined as:\\n\\n    .. math::\\n        \\\\text{RReLU}(x) =\\n        \\\\begin{cases}\\n            x & \\\\text{if } x \\\\geq 0 \\\\\\\\\\n            ax & \\\\text{ otherwise }\\n        \\\\end{cases}\\n\\n    where :math:`a` is randomly sampled from uniform distribution\\n    :math:`\\\\mathcal{U}(\\\\text{lower}, \\\\text{upper})` during training while during\\n    evaluation :math:`a` is fixed with :math:`a = \\\\frac{\\\\text{lower} + \\\\text{upper}}{2}`.\\n\\n    Args:\\n        lower: lower bound of the uniform distribution. Default: :math:`\\\\frac{1}{8}`\\n        upper: upper bound of the uniform distribution. Default: :math:`\\\\frac{1}{3}`\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/RReLU.png\\n\\n    Examples::\\n\\n        >>> m = nn.RReLU(0.1, 0.3)\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"lower\\\", \\\"upper\\\", \\\"inplace\\\"]\\n\\n    lower: float\\n    upper: float\\n    inplace: bool\\n\\n    def __init__(\\n        self, lower: float = 1.0 / 8, upper: float = 1.0 / 3, inplace: bool = False\\n    ):\\n        super().__init__()\\n        self.lower = lower\\n        self.upper = upper\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.rrelu(input, self.lower, self.upper, self.training, self.inplace)\\n\\n    def extra_repr(self):\\n        inplace_str = \\\", inplace=True\\\" if self.inplace else \\\"\\\"\\n        return f\\\"lower={self.lower}, upper={self.upper}{inplace_str}\\\"\\n\\n\\nclass Hardtanh(Module):\\n    r\\\"\\\"\\\"Applies the HardTanh function element-wise.\\n\\n    HardTanh is defined as:\\n\\n    .. math::\\n        \\\\text{HardTanh}(x) = \\\\begin{cases}\\n            \\\\text{max\\\\_val} & \\\\text{ if } x > \\\\text{ max\\\\_val } \\\\\\\\\\n            \\\\text{min\\\\_val} & \\\\text{ if } x < \\\\text{ min\\\\_val } \\\\\\\\\\n            x & \\\\text{ otherwise } \\\\\\\\\\n        \\\\end{cases}\\n\\n    Args:\\n        min_val: minimum value of the linear region range. Default: -1\\n        max_val: maximum value of the linear region range. Default: 1\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Keyword arguments :attr:`min_value` and :attr:`max_value`\\n    have been deprecated in favor of :attr:`min_val` and :attr:`max_val`.\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Hardtanh.png\\n\\n    Examples::\\n\\n        >>> m = nn.Hardtanh(-2, 2)\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"min_val\\\", \\\"max_val\\\", \\\"inplace\\\"]\\n\\n    min_val: float\\n    max_val: float\\n    inplace: bool\\n\\n    def __init__(\\n        self,\\n        min_val: float = -1.0,\\n        max_val: float = 1.0,\\n        inplace: bool = False,\\n        min_value: Optional[float] = None,\\n        max_value: Optional[float] = None,\\n    ) -> None:\\n        super().__init__()\\n        if min_value is not None:\\n            warnings.warn(\\n                \\\"keyword argument `min_value` is deprecated and rename to `min_val`\\\",\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n            min_val = min_value\\n        if max_value is not None:\\n            warnings.warn(\\n                \\\"keyword argument `max_value` is deprecated and rename to `max_val`\\\",\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n            max_val = max_value\\n\\n        self.min_val = min_val\\n        self.max_val = max_val\\n        self.inplace = inplace\\n        assert self.max_val > self.min_val\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.hardtanh(input, self.min_val, self.max_val, self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\", inplace=True\\\" if self.inplace else \\\"\\\"\\n        return f\\\"min_val={self.min_val}, max_val={self.max_val}{inplace_str}\\\"\\n\\n\\nclass ReLU6(Hardtanh):\\n    r\\\"\\\"\\\"Applies the ReLU6 function element-wise.\\n\\n    .. math::\\n        \\\\text{ReLU6}(x) = \\\\min(\\\\max(0,x), 6)\\n\\n    Args:\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/ReLU6.png\\n\\n    Examples::\\n\\n        >>> m = nn.ReLU6()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, inplace: bool = False):\\n        super().__init__(0.0, 6.0, inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\"inplace=True\\\" if self.inplace else \\\"\\\"\\n        return inplace_str\\n\\n\\nclass Sigmoid(Module):\\n    r\\\"\\\"\\\"Applies the Sigmoid function element-wise.\\n\\n    .. math::\\n        \\\\text{Sigmoid}(x) = \\\\sigma(x) = \\\\frac{1}{1 + \\\\exp(-x)}\\n\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Sigmoid.png\\n\\n    Examples::\\n\\n        >>> m = nn.Sigmoid()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return torch.sigmoid(input)\\n\\n\\nclass Hardsigmoid(Module):\\n    r\\\"\\\"\\\"Applies the Hardsigmoid function element-wise.\\n\\n    Hardsigmoid is defined as:\\n\\n    .. math::\\n        \\\\text{Hardsigmoid}(x) = \\\\begin{cases}\\n            0 & \\\\text{if~} x \\\\le -3, \\\\\\\\\\n            1 & \\\\text{if~} x \\\\ge +3, \\\\\\\\\\n            x / 6 + 1 / 2 & \\\\text{otherwise}\\n        \\\\end{cases}\\n\\n    Args:\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Hardsigmoid.png\\n\\n    Examples::\\n\\n        >>> m = nn.Hardsigmoid()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\"]\\n\\n    inplace: bool\\n\\n    def __init__(self, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.hardsigmoid(input, self.inplace)\\n\\n\\nclass Tanh(Module):\\n    r\\\"\\\"\\\"Applies the Hyperbolic Tangent (Tanh) function element-wise.\\n\\n    Tanh is defined as:\\n\\n    .. math::\\n        \\\\text{Tanh}(x) = \\\\tanh(x) = \\\\frac{\\\\exp(x) - \\\\exp(-x)} {\\\\exp(x) + \\\\exp(-x)}\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Tanh.png\\n\\n    Examples::\\n\\n        >>> m = nn.Tanh()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return torch.tanh(input)\\n\\n\\nclass SiLU(Module):\\n    r\\\"\\\"\\\"Applies the Sigmoid Linear Unit (SiLU) function, element-wise.\\n\\n    The SiLU function is also known as the swish function.\\n\\n    .. math::\\n        \\\\text{silu}(x) = x * \\\\sigma(x), \\\\text{where } \\\\sigma(x) \\\\text{ is the logistic sigmoid.}\\n\\n    .. note::\\n        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_\\n        where the SiLU (Sigmoid Linear Unit) was originally coined, and see\\n        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation\\n        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:\\n        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_\\n        where the SiLU was experimented with later.\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/SiLU.png\\n\\n    Examples::\\n\\n        >>> m = nn.SiLU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\"]\\n    inplace: bool\\n\\n    def __init__(self, inplace: bool = False):\\n        super().__init__()\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.silu(input, inplace=self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\"inplace=True\\\" if self.inplace else \\\"\\\"\\n        return inplace_str\\n\\n\\nclass Mish(Module):\\n    r\\\"\\\"\\\"Applies the Mish function, element-wise.\\n\\n    Mish: A Self Regularized Non-Monotonic Neural Activation Function.\\n\\n    .. math::\\n        \\\\text{Mish}(x) = x * \\\\text{Tanh}(\\\\text{Softplus}(x))\\n\\n    .. note::\\n        See `Mish: A Self Regularized Non-Monotonic Neural Activation Function <https://arxiv.org/abs/1908.08681>`_\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Mish.png\\n\\n    Examples::\\n\\n        >>> m = nn.Mish()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\"]\\n    inplace: bool\\n\\n    def __init__(self, inplace: bool = False):\\n        super().__init__()\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.mish(input, inplace=self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\"inplace=True\\\" if self.inplace else \\\"\\\"\\n        return inplace_str\\n\\n\\nclass Hardswish(Module):\\n    r\\\"\\\"\\\"Applies the Hardswish function, element-wise.\\n\\n    Method described in the paper: `Searching for MobileNetV3 <https://arxiv.org/abs/1905.02244>`_.\\n\\n    Hardswish is defined as:\\n\\n    .. math::\\n        \\\\text{Hardswish}(x) = \\\\begin{cases}\\n            0 & \\\\text{if~} x \\\\le -3, \\\\\\\\\\n            x & \\\\text{if~} x \\\\ge +3, \\\\\\\\\\n            x \\\\cdot (x + 3) /6 & \\\\text{otherwise}\\n        \\\\end{cases}\\n\\n    Args:\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Hardswish.png\\n\\n    Examples::\\n\\n        >>> m = nn.Hardswish()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\"]\\n\\n    inplace: bool\\n\\n    def __init__(self, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.hardswish(input, self.inplace)\\n\\n\\nclass ELU(Module):\\n    r\\\"\\\"\\\"Applies the Exponential Linear Unit (ELU) function, element-wise.\\n\\n    Method described in the paper: `Fast and Accurate Deep Network Learning by Exponential Linear\\n    Units (ELUs) <https://arxiv.org/abs/1511.07289>`__.\\n\\n    ELU is defined as:\\n\\n    .. math::\\n        \\\\text{ELU}(x) = \\\\begin{cases}\\n        x, & \\\\text{ if } x > 0\\\\\\\\\\n        \\\\alpha * (\\\\exp(x) - 1), & \\\\text{ if } x \\\\leq 0\\n        \\\\end{cases}\\n\\n    Args:\\n        alpha: the :math:`\\\\alpha` value for the ELU formulation. Default: 1.0\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/ELU.png\\n\\n    Examples::\\n\\n        >>> m = nn.ELU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"alpha\\\", \\\"inplace\\\"]\\n    alpha: float\\n    inplace: bool\\n\\n    def __init__(self, alpha: float = 1.0, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.alpha = alpha\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.elu(input, self.alpha, self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\", inplace=True\\\" if self.inplace else \\\"\\\"\\n        return f\\\"alpha={self.alpha}{inplace_str}\\\"\\n\\n\\nclass CELU(Module):\\n    r\\\"\\\"\\\"Applies the CELU function element-wise.\\n\\n    .. math::\\n        \\\\text{CELU}(x) = \\\\max(0,x) + \\\\min(0, \\\\alpha * (\\\\exp(x/\\\\alpha) - 1))\\n\\n    More details can be found in the paper `Continuously Differentiable Exponential Linear Units`_ .\\n\\n    Args:\\n        alpha: the :math:`\\\\alpha` value for the CELU formulation. Default: 1.0\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/CELU.png\\n\\n    Examples::\\n\\n        >>> m = nn.CELU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n\\n    .. _`Continuously Differentiable Exponential Linear Units`:\\n        https://arxiv.org/abs/1704.07483\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"alpha\\\", \\\"inplace\\\"]\\n    alpha: float\\n    inplace: bool\\n\\n    def __init__(self, alpha: float = 1.0, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.alpha = alpha\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.celu(input, self.alpha, self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\", inplace=True\\\" if self.inplace else \\\"\\\"\\n        return f\\\"alpha={self.alpha}{inplace_str}\\\"\\n\\n\\nclass SELU(Module):\\n    r\\\"\\\"\\\"Applies the SELU function element-wise.\\n\\n    .. math::\\n        \\\\text{SELU}(x) = \\\\text{scale} * (\\\\max(0,x) + \\\\min(0, \\\\alpha * (\\\\exp(x) - 1)))\\n\\n    with :math:`\\\\alpha = 1.6732632423543772848170429916717` and\\n    :math:`\\\\text{scale} = 1.0507009873554804934193349852946`.\\n\\n    .. warning::\\n        When using ``kaiming_normal`` or ``kaiming_normal_`` for initialisation,\\n        ``nonlinearity='linear'`` should be used instead of ``nonlinearity='selu'``\\n        in order to get `Self-Normalizing Neural Networks`_.\\n        See :func:`torch.nn.init.calculate_gain` for more information.\\n\\n    More details can be found in the paper `Self-Normalizing Neural Networks`_ .\\n\\n    Args:\\n        inplace (bool, optional): can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/SELU.png\\n\\n    Examples::\\n\\n        >>> m = nn.SELU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n\\n    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\"]\\n    inplace: bool\\n\\n    def __init__(self, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.selu(input, self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\"inplace=True\\\" if self.inplace else \\\"\\\"\\n        return inplace_str\\n\\n\\nclass GLU(Module):\\n    r\\\"\\\"\\\"Applies the gated linear unit function.\\n\\n    :math:`{GLU}(a, b)= a \\\\otimes \\\\sigma(b)` where :math:`a` is the first half\\n    of the input matrices and :math:`b` is the second half.\\n\\n    Args:\\n        dim (int): the dimension on which to split the input. Default: -1\\n\\n    Shape:\\n        - Input: :math:`(\\\\ast_1, N, \\\\ast_2)` where `*` means, any number of additional\\n          dimensions\\n        - Output: :math:`(\\\\ast_1, M, \\\\ast_2)` where :math:`M=N/2`\\n\\n    Examples::\\n\\n        >>> m = nn.GLU()\\n        >>> input = torch.randn(4, 2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"dim\\\"]\\n    dim: int\\n\\n    def __init__(self, dim: int = -1) -> None:\\n        super().__init__()\\n        self.dim = dim\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.glu(input, self.dim)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"dim={self.dim}\\\"\\n\\n\\nclass GELU(Module):\\n    r\\\"\\\"\\\"Applies the Gaussian Error Linear Units function.\\n\\n    .. math:: \\\\text{GELU}(x) = x * \\\\Phi(x)\\n\\n    where :math:`\\\\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.\\n\\n    When the approximate argument is 'tanh', Gelu is estimated with:\\n\\n    .. math:: \\\\text{GELU}(x) = 0.5 * x * (1 + \\\\text{Tanh}(\\\\sqrt{2 / \\\\pi} * (x + 0.044715 * x^3)))\\n\\n    Args:\\n        approximate (str, optional): the gelu approximation algorithm to use:\\n            ``'none'`` | ``'tanh'``. Default: ``'none'``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/GELU.png\\n\\n    Examples::\\n\\n        >>> m = nn.GELU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"approximate\\\"]\\n    approximate: str\\n\\n    def __init__(self, approximate: str = \\\"none\\\") -> None:\\n        super().__init__()\\n        self.approximate = approximate\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.gelu(input, approximate=self.approximate)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"approximate={repr(self.approximate)}\\\"\\n\\n\\nclass Hardshrink(Module):\\n    r\\\"\\\"\\\"Applies the Hard Shrinkage (Hardshrink) function element-wise.\\n\\n    Hardshrink is defined as:\\n\\n    .. math::\\n        \\\\text{HardShrink}(x) =\\n        \\\\begin{cases}\\n        x, & \\\\text{ if } x > \\\\lambda \\\\\\\\\\n        x, & \\\\text{ if } x < -\\\\lambda \\\\\\\\\\n        0, & \\\\text{ otherwise }\\n        \\\\end{cases}\\n\\n    Args:\\n        lambd: the :math:`\\\\lambda` value for the Hardshrink formulation. Default: 0.5\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Hardshrink.png\\n\\n    Examples::\\n\\n        >>> m = nn.Hardshrink()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"lambd\\\"]\\n    lambd: float\\n\\n    def __init__(self, lambd: float = 0.5) -> None:\\n        super().__init__()\\n        self.lambd = lambd\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.hardshrink(input, self.lambd)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.lambd}\\\"\\n\\n\\nclass LeakyReLU(Module):\\n    r\\\"\\\"\\\"Applies the LeakyReLU function element-wise.\\n\\n    .. math::\\n        \\\\text{LeakyReLU}(x) = \\\\max(0, x) + \\\\text{negative\\\\_slope} * \\\\min(0, x)\\n\\n\\n    or\\n\\n    .. math::\\n        \\\\text{LeakyReLU}(x) =\\n        \\\\begin{cases}\\n        x, & \\\\text{ if } x \\\\geq 0 \\\\\\\\\\n        \\\\text{negative\\\\_slope} \\\\times x, & \\\\text{ otherwise }\\n        \\\\end{cases}\\n\\n    Args:\\n        negative_slope: Controls the angle of the negative slope (which is used for\\n          negative input values). Default: 1e-2\\n        inplace: can optionally do the operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)` where `*` means, any number of additional\\n          dimensions\\n        - Output: :math:`(*)`, same shape as the input\\n\\n    .. image:: ../scripts/activation_images/LeakyReLU.png\\n\\n    Examples::\\n\\n        >>> m = nn.LeakyReLU(0.1)\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"inplace\\\", \\\"negative_slope\\\"]\\n    inplace: bool\\n    negative_slope: float\\n\\n    def __init__(self, negative_slope: float = 1e-2, inplace: bool = False) -> None:\\n        super().__init__()\\n        self.negative_slope = negative_slope\\n        self.inplace = inplace\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.leaky_relu(input, self.negative_slope, self.inplace)\\n\\n    def extra_repr(self) -> str:\\n        inplace_str = \\\", inplace=True\\\" if self.inplace else \\\"\\\"\\n        return f\\\"negative_slope={self.negative_slope}{inplace_str}\\\"\\n\\n\\nclass LogSigmoid(Module):\\n    r\\\"\\\"\\\"Applies the Logsigmoid function element-wise.\\n\\n    .. math::\\n        \\\\text{LogSigmoid}(x) = \\\\log\\\\left(\\\\frac{ 1 }{ 1 + \\\\exp(-x)}\\\\right)\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/LogSigmoid.png\\n\\n    Examples::\\n\\n        >>> m = nn.LogSigmoid()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.logsigmoid(input)\\n\\n\\nclass Softplus(Module):\\n    r\\\"\\\"\\\"Applies the Softplus function element-wise.\\n\\n    .. math::\\n        \\\\text{Softplus}(x) = \\\\frac{1}{\\\\beta} * \\\\log(1 + \\\\exp(\\\\beta * x))\\n\\n    SoftPlus is a smooth approximation to the ReLU function and can be used\\n    to constrain the output of a machine to always be positive.\\n\\n    For numerical stability the implementation reverts to the linear function\\n    when :math:`input \\\\times \\\\beta > threshold`.\\n\\n    Args:\\n        beta: the :math:`\\\\beta` value for the Softplus formulation. Default: 1\\n        threshold: values above this revert to a linear function. Default: 20\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Softplus.png\\n\\n    Examples::\\n\\n        >>> m = nn.Softplus()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"beta\\\", \\\"threshold\\\"]\\n    beta: float\\n    threshold: float\\n\\n    def __init__(self, beta: float = 1.0, threshold: float = 20.0) -> None:\\n        super().__init__()\\n        self.beta = beta\\n        self.threshold = threshold\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.softplus(input, self.beta, self.threshold)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"beta={self.beta}, threshold={self.threshold}\\\"\\n\\n\\nclass Softshrink(Module):\\n    r\\\"\\\"\\\"Applies the soft shrinkage function element-wise.\\n\\n    .. math::\\n        \\\\text{SoftShrinkage}(x) =\\n        \\\\begin{cases}\\n        x - \\\\lambda, & \\\\text{ if } x > \\\\lambda \\\\\\\\\\n        x + \\\\lambda, & \\\\text{ if } x < -\\\\lambda \\\\\\\\\\n        0, & \\\\text{ otherwise }\\n        \\\\end{cases}\\n\\n    Args:\\n        lambd: the :math:`\\\\lambda` (must be no less than zero) value for the Softshrink formulation. Default: 0.5\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Softshrink.png\\n\\n    Examples::\\n\\n        >>> m = nn.Softshrink()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"lambd\\\"]\\n    lambd: float\\n\\n    def __init__(self, lambd: float = 0.5) -> None:\\n        super().__init__()\\n        self.lambd = lambd\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.softshrink(input, self.lambd)\\n\\n    def extra_repr(self) -> str:\\n        return str(self.lambd)\\n\\n\\ndef _check_arg_device(x: Optional[torch.Tensor]) -> bool:\\n    if x is not None:\\n        return x.device.type in [\\n            \\\"cpu\\\",\\n            \\\"cuda\\\",\\n            torch.utils.backend_registration._privateuse1_backend_name,\\n        ]\\n    return True\\n\\n\\ndef _arg_requires_grad(x: Optional[torch.Tensor]) -> bool:\\n    if x is not None:\\n        return x.requires_grad\\n    return False\\n\\n\\ndef _is_make_fx_tracing():\\n    if not torch.jit.is_scripting():\\n        torch_dispatch_mode_stack = (\\n            torch.utils._python_dispatch._get_current_dispatch_mode_stack()\\n        )\\n        return any(\\n            type(x) == torch.fx.experimental.proxy_tensor.ProxyTorchDispatchMode\\n            for x in torch_dispatch_mode_stack\\n        )\\n    else:\\n        return False\\n\\n\\nclass MultiheadAttention(Module):\\n    r\\\"\\\"\\\"Allows the model to jointly attend to information from different representation subspaces.\\n\\n    Method described in the paper:\\n    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.\\n\\n    Multi-Head Attention is defined as:\\n\\n    .. math::\\n        \\\\text{MultiHead}(Q, K, V) = \\\\text{Concat}(head_1,\\\\dots,head_h)W^O\\n\\n    where :math:`head_i = \\\\text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.\\n\\n    ``nn.MultiHeadAttention`` will use the optimized implementations of\\n    ``scaled_dot_product_attention()`` when possible.\\n\\n    In addition to support for the new ``scaled_dot_product_attention()``\\n    function, for speeding up Inference, MHA will use\\n    fastpath inference with support for Nested Tensors, iff:\\n\\n    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor).\\n    - inputs are batched (3D) with ``batch_first==True``\\n    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``\\n    - training is disabled (using ``.eval()``)\\n    - ``add_bias_kv`` is ``False``\\n    - ``add_zero_attn`` is ``False``\\n    - ``kdim`` and ``vdim`` are equal to ``embed_dim``\\n    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``\\n      nor ``attn_mask`` is passed\\n    - autocast is disabled\\n\\n    If the optimized inference fastpath implementation is in use, a\\n    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for\\n    ``query``/``key``/``value`` to represent padding more efficiently than using a\\n    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_\\n    will be returned, and an additional speedup proportional to the fraction of the input\\n    that is padding can be expected.\\n\\n    Args:\\n        embed_dim: Total dimension of the model.\\n        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split\\n            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).\\n        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).\\n        bias: If specified, adds bias to input / output projection layers. Default: ``True``.\\n        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.\\n        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.\\n            Default: ``False``.\\n        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).\\n        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).\\n\\n    Examples::\\n\\n        >>> # xdoctest: +SKIP\\n        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)\\n        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)\\n\\n    .. _`FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`:\\n         https://arxiv.org/abs/2205.14135\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"batch_first\\\"]\\n    bias_k: Optional[torch.Tensor]\\n    bias_v: Optional[torch.Tensor]\\n\\n    def __init__(\\n        self,\\n        embed_dim,\\n        num_heads,\\n        dropout=0.0,\\n        bias=True,\\n        add_bias_kv=False,\\n        add_zero_attn=False,\\n        kdim=None,\\n        vdim=None,\\n        batch_first=False,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        if embed_dim <= 0 or num_heads <= 0:\\n            raise ValueError(\\n                f\\\"embed_dim and num_heads must be greater than 0,\\\"\\n                f\\\" got embed_dim={embed_dim} and num_heads={num_heads} instead\\\"\\n            )\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.embed_dim = embed_dim\\n        self.kdim = kdim if kdim is not None else embed_dim\\n        self.vdim = vdim if vdim is not None else embed_dim\\n        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim\\n\\n        self.num_heads = num_heads\\n        self.dropout = dropout\\n        self.batch_first = batch_first\\n        self.head_dim = embed_dim // num_heads\\n        assert (\\n            self.head_dim * num_heads == self.embed_dim\\n        ), \\\"embed_dim must be divisible by num_heads\\\"\\n\\n        if not self._qkv_same_embed_dim:\\n            self.q_proj_weight = Parameter(\\n                torch.empty((embed_dim, embed_dim), **factory_kwargs)\\n            )\\n            self.k_proj_weight = Parameter(\\n                torch.empty((embed_dim, self.kdim), **factory_kwargs)\\n            )\\n            self.v_proj_weight = Parameter(\\n                torch.empty((embed_dim, self.vdim), **factory_kwargs)\\n            )\\n            self.register_parameter(\\\"in_proj_weight\\\", None)\\n        else:\\n            self.in_proj_weight = Parameter(\\n                torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)\\n            )\\n            self.register_parameter(\\\"q_proj_weight\\\", None)\\n            self.register_parameter(\\\"k_proj_weight\\\", None)\\n            self.register_parameter(\\\"v_proj_weight\\\", None)\\n\\n        if bias:\\n            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim, **factory_kwargs))\\n        else:\\n            self.register_parameter(\\\"in_proj_bias\\\", None)\\n        self.out_proj = NonDynamicallyQuantizableLinear(\\n            embed_dim, embed_dim, bias=bias, **factory_kwargs\\n        )\\n\\n        if add_bias_kv:\\n            self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))\\n            self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))\\n        else:\\n            self.bias_k = self.bias_v = None\\n\\n        self.add_zero_attn = add_zero_attn\\n\\n        self._reset_parameters()\\n\\n    def _reset_parameters(self):\\n        if self._qkv_same_embed_dim:\\n            xavier_uniform_(self.in_proj_weight)\\n        else:\\n            xavier_uniform_(self.q_proj_weight)\\n            xavier_uniform_(self.k_proj_weight)\\n            xavier_uniform_(self.v_proj_weight)\\n\\n        if self.in_proj_bias is not None:\\n            constant_(self.in_proj_bias, 0.0)\\n            constant_(self.out_proj.bias, 0.0)\\n        if self.bias_k is not None:\\n            xavier_normal_(self.bias_k)\\n        if self.bias_v is not None:\\n            xavier_normal_(self.bias_v)\\n\\n    def __setstate__(self, state):\\n        # Support loading old MultiheadAttention checkpoints generated by v1.1.0\\n        if \\\"_qkv_same_embed_dim\\\" not in state:\\n            state[\\\"_qkv_same_embed_dim\\\"] = True\\n\\n        super().__setstate__(state)\\n\\n    def forward(\\n        self,\\n        query: Tensor,\\n        key: Tensor,\\n        value: Tensor,\\n        key_padding_mask: Optional[Tensor] = None,\\n        need_weights: bool = True,\\n        attn_mask: Optional[Tensor] = None,\\n        average_attn_weights: bool = True,\\n        is_causal: bool = False,\\n    ) -> Tuple[Tensor, Optional[Tensor]]:\\n        r\\\"\\\"\\\"Compute attention outputs using query, key, and value embeddings.\\n\\n            Supports optional parameters for padding, masks and attention weights.\\n\\n        Args:\\n            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``\\n                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,\\n                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.\\n                Queries are compared against key-value pairs to produce the output.\\n                See \\\"Attention Is All You Need\\\" for more details.\\n            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``\\n                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,\\n                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.\\n                See \\\"Attention Is All You Need\\\" for more details.\\n            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when\\n                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source\\n                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.\\n                See \\\"Attention Is All You Need\\\" for more details.\\n            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``\\n                to ignore for the purpose of attention (i.e. treat as \\\"padding\\\"). For unbatched `query`, shape should be :math:`(S)`.\\n                Binary and float masks are supported.\\n                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for\\n                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.\\n            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.\\n                Set ``need_weights=False`` to use the optimized ``scaled_dot_product_attention``\\n                and achieve the best performance for MHA.\\n                Default: ``True``.\\n            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape\\n                :math:`(L, S)` or :math:`(N\\\\cdot\\\\text{num\\\\_heads}, L, S)`, where :math:`N` is the batch size,\\n                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be\\n                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.\\n                Binary and float masks are supported. For a binary mask, a ``True`` value indicates that the\\n                corresponding position is not allowed to attend. For a float mask, the mask values will be added to\\n                the attention weight.\\n                If both attn_mask and key_padding_mask are supplied, their types should match.\\n            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across\\n                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an\\n                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)\\n            is_causal: If specified, applies a causal mask as attention mask.\\n                Default: ``False``.\\n                Warning:\\n                ``is_causal`` provides a hint that ``attn_mask`` is the\\n                causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n\\n        Outputs:\\n            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,\\n              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,\\n              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the\\n              embedding dimension ``embed_dim``.\\n            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,\\n              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or\\n              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and\\n              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per\\n              head of shape :math:`(\\\\text{num\\\\_heads}, L, S)` when input is unbatched or :math:`(N, \\\\text{num\\\\_heads}, L, S)`.\\n\\n            .. note::\\n                `batch_first` argument is ignored for unbatched inputs.\\n        \\\"\\\"\\\"  # noqa: B950\\n        why_not_fast_path = \\\"\\\"\\n        if (\\n            (attn_mask is not None and torch.is_floating_point(attn_mask))\\n            or (key_padding_mask is not None)\\n            and torch.is_floating_point(key_padding_mask)\\n        ):\\n            why_not_fast_path = \\\"floating-point masks are not supported for fast path.\\\"\\n\\n        is_batched = query.dim() == 3\\n\\n        key_padding_mask = F._canonical_mask(\\n            mask=key_padding_mask,\\n            mask_name=\\\"key_padding_mask\\\",\\n            other_type=F._none_or_dtype(attn_mask),\\n            other_name=\\\"attn_mask\\\",\\n            target_type=query.dtype,\\n        )\\n\\n        attn_mask = F._canonical_mask(\\n            mask=attn_mask,\\n            mask_name=\\\"attn_mask\\\",\\n            other_type=None,\\n            other_name=\\\"\\\",\\n            target_type=query.dtype,\\n            check_other=False,\\n        )\\n\\n        is_fastpath_enabled = torch.backends.mha.get_fastpath_enabled()\\n\\n        if not is_fastpath_enabled:\\n            why_not_fast_path = \\\"torch.backends.mha.get_fastpath_enabled() was not True\\\"\\n        elif not is_batched:\\n            why_not_fast_path = (\\n                f\\\"input not batched; expected query.dim() of 3 but got {query.dim()}\\\"\\n            )\\n        elif query is not key or key is not value:\\n            # When lifting this restriction, don't forget to either\\n            # enforce that the dtypes all match or test cases where\\n            # they don't!\\n            why_not_fast_path = \\\"non-self attention was used (query, key, and value are not the same Tensor)\\\"\\n        elif self.in_proj_bias is not None and query.dtype != self.in_proj_bias.dtype:\\n            why_not_fast_path = f\\\"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match\\\"\\n        elif self.in_proj_weight is None:\\n            why_not_fast_path = \\\"in_proj_weight was None\\\"\\n        elif query.dtype != self.in_proj_weight.dtype:\\n            # this case will fail anyway, but at least they'll get a useful error message.\\n            why_not_fast_path = f\\\"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match\\\"\\n        elif self.training:\\n            why_not_fast_path = \\\"training is enabled\\\"\\n        elif (self.num_heads % 2) != 0:\\n            why_not_fast_path = \\\"self.num_heads is not even\\\"\\n        elif not self.batch_first:\\n            why_not_fast_path = \\\"batch_first was not True\\\"\\n        elif self.bias_k is not None:\\n            why_not_fast_path = \\\"self.bias_k was not None\\\"\\n        elif self.bias_v is not None:\\n            why_not_fast_path = \\\"self.bias_v was not None\\\"\\n        elif self.add_zero_attn:\\n            why_not_fast_path = \\\"add_zero_attn was enabled\\\"\\n        elif not self._qkv_same_embed_dim:\\n            why_not_fast_path = \\\"_qkv_same_embed_dim was not True\\\"\\n        elif query.is_nested and (\\n            key_padding_mask is not None or attn_mask is not None\\n        ):\\n            why_not_fast_path = \\\"supplying both src_key_padding_mask and src_mask at the same time \\\\\\n                                 is not supported with NestedTensor input\\\"\\n        elif torch.is_autocast_enabled():\\n            why_not_fast_path = \\\"autocast is enabled\\\"\\n\\n        if not why_not_fast_path:\\n            tensor_args = (\\n                query,\\n                key,\\n                value,\\n                self.in_proj_weight,\\n                self.in_proj_bias,\\n                self.out_proj.weight,\\n                self.out_proj.bias,\\n            )\\n            # We have to use list comprehensions below because TorchScript does not support\\n            # generator expressions.\\n            if torch.overrides.has_torch_function(tensor_args):\\n                why_not_fast_path = \\\"some Tensor argument has_torch_function\\\"\\n            elif _is_make_fx_tracing():\\n                why_not_fast_path = \\\"we are running make_fx tracing\\\"\\n            elif not all(_check_arg_device(x) for x in tensor_args):\\n                why_not_fast_path = (\\n                    \\\"some Tensor argument's device is neither one of \\\"\\n                    f\\\"cpu, cuda or {torch.utils.backend_registration._privateuse1_backend_name}\\\"\\n                )\\n            elif torch.is_grad_enabled() and any(\\n                _arg_requires_grad(x) for x in tensor_args\\n            ):\\n                why_not_fast_path = (\\n                    \\\"grad is enabled and at least one of query or the \\\"\\n                    \\\"input/output projection weights or biases requires_grad\\\"\\n                )\\n            if not why_not_fast_path:\\n                merged_mask, mask_type = self.merge_masks(\\n                    attn_mask, key_padding_mask, query\\n                )\\n\\n                if self.in_proj_bias is not None and self.in_proj_weight is not None:\\n                    return torch._native_multi_head_attention(\\n                        query,\\n                        key,\\n                        value,\\n                        self.embed_dim,\\n                        self.num_heads,\\n                        self.in_proj_weight,\\n                        self.in_proj_bias,\\n                        self.out_proj.weight,\\n                        self.out_proj.bias,\\n                        merged_mask,\\n                        need_weights,\\n                        average_attn_weights,\\n                        mask_type,\\n                    )\\n\\n        any_nested = query.is_nested or key.is_nested or value.is_nested\\n        assert not any_nested, (\\n            \\\"MultiheadAttention does not support NestedTensor outside of its fast path. \\\"\\n            + f\\\"The fast path was not hit because {why_not_fast_path}\\\"\\n        )\\n\\n        if self.batch_first and is_batched:\\n            # make sure that the transpose op does not affect the \\\"is\\\" property\\n            if key is value:\\n                if query is key:\\n                    query = key = value = query.transpose(1, 0)\\n                else:\\n                    query, key = (x.transpose(1, 0) for x in (query, key))\\n                    value = key\\n            else:\\n                query, key, value = (x.transpose(1, 0) for x in (query, key, value))\\n\\n        if not self._qkv_same_embed_dim:\\n            attn_output, attn_output_weights = F.multi_head_attention_forward(\\n                query,\\n                key,\\n                value,\\n                self.embed_dim,\\n                self.num_heads,\\n                self.in_proj_weight,\\n                self.in_proj_bias,\\n                self.bias_k,\\n                self.bias_v,\\n                self.add_zero_attn,\\n                self.dropout,\\n                self.out_proj.weight,\\n                self.out_proj.bias,\\n                training=self.training,\\n                key_padding_mask=key_padding_mask,\\n                need_weights=need_weights,\\n                attn_mask=attn_mask,\\n                use_separate_proj_weight=True,\\n                q_proj_weight=self.q_proj_weight,\\n                k_proj_weight=self.k_proj_weight,\\n                v_proj_weight=self.v_proj_weight,\\n                average_attn_weights=average_attn_weights,\\n                is_causal=is_causal,\\n            )\\n        else:\\n            attn_output, attn_output_weights = F.multi_head_attention_forward(\\n                query,\\n                key,\\n                value,\\n                self.embed_dim,\\n                self.num_heads,\\n                self.in_proj_weight,\\n                self.in_proj_bias,\\n                self.bias_k,\\n                self.bias_v,\\n                self.add_zero_attn,\\n                self.dropout,\\n                self.out_proj.weight,\\n                self.out_proj.bias,\\n                training=self.training,\\n                key_padding_mask=key_padding_mask,\\n                need_weights=need_weights,\\n                attn_mask=attn_mask,\\n                average_attn_weights=average_attn_weights,\\n                is_causal=is_causal,\\n            )\\n        if self.batch_first and is_batched:\\n            return attn_output.transpose(1, 0), attn_output_weights\\n        else:\\n            return attn_output, attn_output_weights\\n\\n    def merge_masks(\\n        self,\\n        attn_mask: Optional[Tensor],\\n        key_padding_mask: Optional[Tensor],\\n        query: Tensor,\\n    ) -> Tuple[Optional[Tensor], Optional[int]]:\\n        r\\\"\\\"\\\"Determine mask type and combine masks if necessary.\\n\\n        If only one mask is provided, that mask\\n        and the corresponding mask type will be returned. If both masks are provided, they will be both\\n        expanded to shape ``(batch_size, num_heads, seq_len, seq_len)``, combined with logical ``or``\\n        and mask type 2 will be returned\\n        Args:\\n            attn_mask: attention mask of shape ``(seq_len, seq_len)``, mask type 0\\n            key_padding_mask: padding mask of shape ``(batch_size, seq_len)``, mask type 1\\n            query: query embeddings of shape ``(batch_size, seq_len, embed_dim)``\\n        Returns:\\n            merged_mask: merged mask\\n            mask_type: merged mask type (0, 1, or 2)\\n        \\\"\\\"\\\"\\n        mask_type: Optional[int] = None\\n        merged_mask: Optional[Tensor] = None\\n\\n        if key_padding_mask is not None:\\n            mask_type = 1\\n            merged_mask = key_padding_mask\\n\\n        if attn_mask is not None:\\n            # In this branch query can't be a nested tensor, so it has a shape\\n            batch_size, seq_len, _ = query.shape\\n            mask_type = 2\\n\\n            # Always expands attn_mask to 4D\\n            if attn_mask.dim() == 3:\\n                attn_mask_expanded = attn_mask.view(batch_size, -1, seq_len, seq_len)\\n            else:  # attn_mask.dim() == 2:\\n                attn_mask_expanded = attn_mask.view(1, 1, seq_len, seq_len).expand(\\n                    batch_size, self.num_heads, -1, -1\\n                )\\n            merged_mask = attn_mask_expanded\\n\\n            if key_padding_mask is not None:\\n                key_padding_mask_expanded = key_padding_mask.view(\\n                    batch_size, 1, 1, seq_len\\n                ).expand(-1, self.num_heads, -1, -1)\\n                merged_mask = attn_mask_expanded + key_padding_mask_expanded\\n\\n        # no attn_mask and no key_padding_mask, returns None, None\\n        return merged_mask, mask_type\\n\\n\\nclass PReLU(Module):\\n    r\\\"\\\"\\\"Applies the element-wise PReLU function.\\n\\n    .. math::\\n        \\\\text{PReLU}(x) = \\\\max(0,x) + a * \\\\min(0,x)\\n\\n    or\\n\\n    .. math::\\n        \\\\text{PReLU}(x) =\\n        \\\\begin{cases}\\n        x, & \\\\text{ if } x \\\\ge 0 \\\\\\\\\\n        ax, & \\\\text{ otherwise }\\n        \\\\end{cases}\\n\\n    Here :math:`a` is a learnable parameter. When called without arguments, `nn.PReLU()` uses a single\\n    parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,\\n    a separate :math:`a` is used for each input channel.\\n\\n\\n    .. note::\\n        weight decay should not be used when learning :math:`a` for good performance.\\n\\n    .. note::\\n        Channel dim is the 2nd dim of input. When input has dims < 2, then there is\\n        no channel dim and the number of channels = 1.\\n\\n    Args:\\n        num_parameters (int): number of :math:`a` to learn.\\n            Although it takes an int as input, there is only two values are legitimate:\\n            1, or the number of channels at input. Default: 1\\n        init (float): the initial value of :math:`a`. Default: 0.25\\n\\n    Shape:\\n        - Input: :math:`( *)` where `*` means, any number of additional\\n          dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of shape (:attr:`num_parameters`).\\n\\n    .. image:: ../scripts/activation_images/PReLU.png\\n\\n    Examples::\\n\\n        >>> m = nn.PReLU()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"num_parameters\\\"]\\n    num_parameters: int\\n\\n    def __init__(\\n        self, num_parameters: int = 1, init: float = 0.25, device=None, dtype=None\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        self.num_parameters = num_parameters\\n        super().__init__()\\n        self.init = init\\n        self.weight = Parameter(torch.empty(num_parameters, **factory_kwargs))\\n        self.reset_parameters()\\n\\n    def reset_parameters(self):\\n        torch.nn.init.constant_(self.weight, self.init)\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.prelu(input, self.weight)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"num_parameters={self.num_parameters}\\\"\\n\\n\\nclass Softsign(Module):\\n    r\\\"\\\"\\\"Applies the element-wise Softsign function.\\n\\n    .. math::\\n        \\\\text{SoftSign}(x) = \\\\frac{x}{ 1 + |x|}\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Softsign.png\\n\\n    Examples::\\n\\n        >>> m = nn.Softsign()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.softsign(input)\\n\\n\\nclass Tanhshrink(Module):\\n    r\\\"\\\"\\\"Applies the element-wise Tanhshrink function.\\n\\n    .. math::\\n        \\\\text{Tanhshrink}(x) = x - \\\\tanh(x)\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Output: :math:`(*)`, same shape as the input.\\n\\n    .. image:: ../scripts/activation_images/Tanhshrink.png\\n\\n    Examples::\\n\\n        >>> m = nn.Tanhshrink()\\n        >>> input = torch.randn(2)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.tanhshrink(input)\\n\\n\\nclass Softmin(Module):\\n    r\\\"\\\"\\\"Applies the Softmin function to an n-dimensional input Tensor.\\n\\n    Rescales them so that the elements of the n-dimensional output Tensor\\n    lie in the range `[0, 1]` and sum to 1.\\n\\n    Softmin is defined as:\\n\\n    .. math::\\n        \\\\text{Softmin}(x_{i}) = \\\\frac{\\\\exp(-x_i)}{\\\\sum_j \\\\exp(-x_j)}\\n\\n    Shape:\\n        - Input: :math:`(*)` where `*` means, any number of additional\\n          dimensions\\n        - Output: :math:`(*)`, same shape as the input\\n\\n    Args:\\n        dim (int): A dimension along which Softmin will be computed (so every slice\\n            along dim will sum to 1).\\n\\n    Returns:\\n        a Tensor of the same dimension and shape as the input, with\\n        values in the range [0, 1]\\n\\n    Examples::\\n\\n        >>> m = nn.Softmin(dim=1)\\n        >>> input = torch.randn(2, 3)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"dim\\\"]\\n    dim: Optional[int]\\n\\n    def __init__(self, dim: Optional[int] = None) -> None:\\n        super().__init__()\\n        self.dim = dim\\n\\n    def __setstate__(self, state):\\n        super().__setstate__(state)\\n        if not hasattr(self, \\\"dim\\\"):\\n            self.dim = None\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.softmin(input, self.dim, _stacklevel=5)\\n\\n    def extra_repr(self):\\n        return f\\\"dim={self.dim}\\\"\\n\\n\\nclass Softmax(Module):\\n    r\\\"\\\"\\\"Applies the Softmax function to an n-dimensional input Tensor.\\n\\n    Rescales them so that the elements of the n-dimensional output Tensor\\n    lie in the range [0,1] and sum to 1.\\n\\n    Softmax is defined as:\\n\\n    .. math::\\n        \\\\text{Softmax}(x_{i}) = \\\\frac{\\\\exp(x_i)}{\\\\sum_j \\\\exp(x_j)}\\n\\n    When the input Tensor is a sparse tensor then the unspecified\\n    values are treated as ``-inf``.\\n\\n    Shape:\\n        - Input: :math:`(*)` where `*` means, any number of additional\\n          dimensions\\n        - Output: :math:`(*)`, same shape as the input\\n\\n    Returns:\\n        a Tensor of the same dimension and shape as the input with\\n        values in the range [0, 1]\\n\\n    Args:\\n        dim (int): A dimension along which Softmax will be computed (so every slice\\n            along dim will sum to 1).\\n\\n    .. note::\\n        This module doesn't work directly with NLLLoss,\\n        which expects the Log to be computed between the Softmax and itself.\\n        Use `LogSoftmax` instead (it's faster and has better numerical properties).\\n\\n    Examples::\\n\\n        >>> m = nn.Softmax(dim=1)\\n        >>> input = torch.randn(2, 3)\\n        >>> output = m(input)\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"dim\\\"]\\n    dim: Optional[int]\\n\\n    def __init__(self, dim: Optional[int] = None) -> None:\\n        super().__init__()\\n        self.dim = dim\\n\\n    def __setstate__(self, state):\\n        super().__setstate__(state)\\n        if not hasattr(self, \\\"dim\\\"):\\n            self.dim = None\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.softmax(input, self.dim, _stacklevel=5)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"dim={self.dim}\\\"\\n\\n\\nclass Softmax2d(Module):\\n    r\\\"\\\"\\\"Applies SoftMax over features to each spatial location.\\n\\n    When given an image of ``Channels x Height x Width``, it will\\n    apply `Softmax` to each location :math:`(Channels, h_i, w_j)`\\n\\n    Shape:\\n        - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`.\\n        - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)\\n\\n    Returns:\\n        a Tensor of the same dimension and shape as the input with\\n        values in the range [0, 1]\\n\\n    Examples::\\n\\n        >>> m = nn.Softmax2d()\\n        >>> # you softmax over the 2nd dimension\\n        >>> input = torch.randn(2, 3, 12, 13)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        if input.dim() not in (3, 4):\\n            raise ValueError(\\n                f\\\"Softmax2d: expected input to be 3D or 4D, got {input.dim()}D instead\\\"\\n            )\\n        return F.softmax(input, -3, _stacklevel=5)\\n\\n\\nclass LogSoftmax(Module):\\n    r\\\"\\\"\\\"Applies the :math:`\\\\log(\\\\text{Softmax}(x))` function to an n-dimensional input Tensor.\\n\\n    The LogSoftmax formulation can be simplified as:\\n\\n    .. math::\\n        \\\\text{LogSoftmax}(x_{i}) = \\\\log\\\\left(\\\\frac{\\\\exp(x_i) }{ \\\\sum_j \\\\exp(x_j)} \\\\right)\\n\\n    Shape:\\n        - Input: :math:`(*)` where `*` means, any number of additional\\n          dimensions\\n        - Output: :math:`(*)`, same shape as the input\\n\\n    Args:\\n        dim (int): A dimension along which LogSoftmax will be computed.\\n\\n    Returns:\\n        a Tensor of the same dimension and shape as the input with\\n        values in the range [-inf, 0)\\n\\n    Examples::\\n\\n        >>> m = nn.LogSoftmax(dim=1)\\n        >>> input = torch.randn(2, 3)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"dim\\\"]\\n    dim: Optional[int]\\n\\n    def __init__(self, dim: Optional[int] = None) -> None:\\n        super().__init__()\\n        self.dim = dim\\n\\n    def __setstate__(self, state):\\n        super().__setstate__(state)\\n        if not hasattr(self, \\\"dim\\\"):\\n            self.dim = None\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.log_softmax(input, self.dim, _stacklevel=5)\\n\\n    def extra_repr(self):\\n        return f\\\"dim={self.dim}\\\"\\n\\n\\n# mypy: allow-untyped-defs\\nimport copy\\nimport warnings\\nfrom typing import Any, Callable, Optional, Union\\n\\nimport torch\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.init import xavier_uniform_\\n\\nfrom .activation import MultiheadAttention\\nfrom .container import ModuleList\\nfrom .dropout import Dropout\\nfrom .linear import Linear\\nfrom .module import Module\\nfrom .normalization import LayerNorm\\n\\n\\n__all__ = [\\n    \\\"Transformer\\\",\\n    \\\"TransformerEncoder\\\",\\n    \\\"TransformerDecoder\\\",\\n    \\\"TransformerEncoderLayer\\\",\\n    \\\"TransformerDecoderLayer\\\",\\n]\\n\\n\\ndef _generate_square_subsequent_mask(\\n    sz: int,\\n    device: Optional[torch.device] = None,\\n    dtype: Optional[torch.dtype] = None,\\n) -> Tensor:\\n    r\\\"\\\"\\\"Generate a square causal mask for the sequence.\\n\\n    The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0).\\n    \\\"\\\"\\\"\\n    if device is None:\\n        device = torch.device(\\\"cpu\\\")\\n    if dtype is None:\\n        dtype = torch.float32\\n    return torch.triu(\\n        torch.full((sz, sz), float(\\\"-inf\\\"), dtype=dtype, device=device),\\n        diagonal=1,\\n    )\\n\\n\\ndef _get_seq_len(src: Tensor, batch_first: bool) -> Optional[int]:\\n    if src.is_nested:\\n        return None\\n    else:\\n        src_size = src.size()\\n        if len(src_size) == 2:\\n            # unbatched: S, E\\n            return src_size[0]\\n        else:\\n            # batched: B, S, E if batch_first else S, B, E\\n            seq_len_pos = 1 if batch_first else 0\\n            return src_size[seq_len_pos]\\n\\n\\nclass Transformer(Module):\\n    r\\\"\\\"\\\"A transformer model.\\n\\n    User is able to modify the attributes as needed. The architecture\\n    is based on the paper \\\"Attention Is All You Need\\\". Ashish Vaswani, Noam Shazeer,\\n    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and\\n    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information\\n    Processing Systems, pages 6000-6010.\\n\\n    Args:\\n        d_model: the number of expected features in the encoder/decoder inputs (default=512).\\n        nhead: the number of heads in the multiheadattention models (default=8).\\n        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).\\n        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).\\n        dim_feedforward: the dimension of the feedforward network model (default=2048).\\n        dropout: the dropout value (default=0.1).\\n        activation: the activation function of encoder/decoder intermediate layer, can be a string\\n            (\\\"relu\\\" or \\\"gelu\\\") or a unary callable. Default: relu\\n        custom_encoder: custom encoder (default=None).\\n        custom_decoder: custom decoder (default=None).\\n        layer_norm_eps: the eps value in layer normalization components (default=1e-5).\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).\\n        norm_first: if ``True``, encoder and decoder layers will perform LayerNorms before\\n            other attention and feedforward operations, otherwise after. Default: ``False`` (after).\\n        bias: If set to ``False``, ``Linear`` and ``LayerNorm`` layers will not learn an additive\\n            bias. Default: ``True``.\\n\\n    Examples::\\n        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)\\n        >>> src = torch.rand((10, 32, 512))\\n        >>> tgt = torch.rand((20, 32, 512))\\n        >>> out = transformer_model(src, tgt)\\n\\n    Note: A full example to apply nn.Transformer module for the word language model is available in\\n    https://github.com/pytorch/examples/tree/master/word_language_model\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        d_model: int = 512,\\n        nhead: int = 8,\\n        num_encoder_layers: int = 6,\\n        num_decoder_layers: int = 6,\\n        dim_feedforward: int = 2048,\\n        dropout: float = 0.1,\\n        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,\\n        custom_encoder: Optional[Any] = None,\\n        custom_decoder: Optional[Any] = None,\\n        layer_norm_eps: float = 1e-5,\\n        batch_first: bool = False,\\n        norm_first: bool = False,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        torch._C._log_api_usage_once(f\\\"torch.nn.modules.{self.__class__.__name__}\\\")\\n\\n        if custom_encoder is not None:\\n            self.encoder = custom_encoder\\n        else:\\n            encoder_layer = TransformerEncoderLayer(\\n                d_model,\\n                nhead,\\n                dim_feedforward,\\n                dropout,\\n                activation,\\n                layer_norm_eps,\\n                batch_first,\\n                norm_first,\\n                bias,\\n                **factory_kwargs,\\n            )\\n            encoder_norm = LayerNorm(\\n                d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs\\n            )\\n            self.encoder = TransformerEncoder(\\n                encoder_layer, num_encoder_layers, encoder_norm\\n            )\\n\\n        if custom_decoder is not None:\\n            self.decoder = custom_decoder\\n        else:\\n            decoder_layer = TransformerDecoderLayer(\\n                d_model,\\n                nhead,\\n                dim_feedforward,\\n                dropout,\\n                activation,\\n                layer_norm_eps,\\n                batch_first,\\n                norm_first,\\n                bias,\\n                **factory_kwargs,\\n            )\\n            decoder_norm = LayerNorm(\\n                d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs\\n            )\\n            self.decoder = TransformerDecoder(\\n                decoder_layer, num_decoder_layers, decoder_norm\\n            )\\n\\n        self._reset_parameters()\\n\\n        self.d_model = d_model\\n        self.nhead = nhead\\n\\n        self.batch_first = batch_first\\n\\n    def forward(\\n        self,\\n        src: Tensor,\\n        tgt: Tensor,\\n        src_mask: Optional[Tensor] = None,\\n        tgt_mask: Optional[Tensor] = None,\\n        memory_mask: Optional[Tensor] = None,\\n        src_key_padding_mask: Optional[Tensor] = None,\\n        tgt_key_padding_mask: Optional[Tensor] = None,\\n        memory_key_padding_mask: Optional[Tensor] = None,\\n        src_is_causal: Optional[bool] = None,\\n        tgt_is_causal: Optional[bool] = None,\\n        memory_is_causal: bool = False,\\n    ) -> Tensor:\\n        r\\\"\\\"\\\"Take in and process masked source/target sequences.\\n\\n        .. note::\\n\\n            If a boolean tensor is provided for any of the [src/tgt/memory]_mask arguments, positions with a ``True`` value are\\n            not allowed to participate in the attention,\\n            which is the opposite of the definition for :attr:`attn_mask`\\n            in :func:`torch.nn.functional.scaled_dot_product_attention`.\\n\\n        Args:\\n            src: the sequence to the encoder (required).\\n            tgt: the sequence to the decoder (required).\\n            src_mask: the additive mask for the src sequence (optional).\\n            tgt_mask: the additive mask for the tgt sequence (optional).\\n            memory_mask: the additive mask for the encoder output (optional).\\n            src_key_padding_mask: the Tensor mask for src keys per batch (optional).\\n            tgt_key_padding_mask: the Tensor mask for tgt keys per batch (optional).\\n            memory_key_padding_mask: the Tensor mask for memory keys per batch (optional).\\n            src_is_causal: If specified, applies a causal mask as ``src_mask``.\\n                Default: ``None``; try to detect a causal mask.\\n                Warning:\\n                ``src_is_causal`` provides a hint that ``src_mask`` is\\n                the causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n            tgt_is_causal: If specified, applies a causal mask as ``tgt_mask``.\\n                Default: ``None``; try to detect a causal mask.\\n                Warning:\\n                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is\\n                the causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n            memory_is_causal: If specified, applies a causal mask as\\n                ``memory_mask``.\\n                Default: ``False``.\\n                Warning:\\n                ``memory_is_causal`` provides a hint that\\n                ``memory_mask`` is the causal mask. Providing incorrect\\n                hints can result in incorrect execution, including\\n                forward and backward compatibility.\\n\\n        Shape:\\n            - src: :math:`(S, E)` for unbatched input, :math:`(S, N, E)` if `batch_first=False` or\\n              `(N, S, E)` if `batch_first=True`.\\n            - tgt: :math:`(T, E)` for unbatched input, :math:`(T, N, E)` if `batch_first=False` or\\n              `(N, T, E)` if `batch_first=True`.\\n            - src_mask: :math:`(S, S)` or :math:`(N\\\\cdot\\\\text{num\\\\_heads}, S, S)`.\\n            - tgt_mask: :math:`(T, T)` or :math:`(N\\\\cdot\\\\text{num\\\\_heads}, T, T)`.\\n            - memory_mask: :math:`(T, S)`.\\n            - src_key_padding_mask: :math:`(S)` for unbatched input otherwise :math:`(N, S)`.\\n            - tgt_key_padding_mask: :math:`(T)` for unbatched input otherwise :math:`(N, T)`.\\n            - memory_key_padding_mask: :math:`(S)` for unbatched input otherwise :math:`(N, S)`.\\n\\n            Note: [src/tgt/memory]_mask ensures that position :math:`i` is allowed to attend the unmasked\\n            positions. If a BoolTensor is provided, positions with ``True``\\n            are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor\\n            is provided, it will be added to the attention weight.\\n            [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by\\n            the attention. If a BoolTensor is provided, the positions with the\\n            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.\\n\\n            - output: :math:`(T, E)` for unbatched input, :math:`(T, N, E)` if `batch_first=False` or\\n              `(N, T, E)` if `batch_first=True`.\\n\\n            Note: Due to the multi-head attention architecture in the transformer model,\\n            the output sequence length of a transformer is same as the input sequence\\n            (i.e. target) length of the decoder.\\n\\n            where :math:`S` is the source sequence length, :math:`T` is the target sequence length, :math:`N` is the\\n            batch size, :math:`E` is the feature number\\n\\n        Examples:\\n            >>> # xdoctest: +SKIP\\n            >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)\\n        \\\"\\\"\\\"\\n        is_batched = src.dim() == 3\\n        if not self.batch_first and src.size(1) != tgt.size(1) and is_batched:\\n            raise RuntimeError(\\\"the batch number of src and tgt must be equal\\\")\\n        elif self.batch_first and src.size(0) != tgt.size(0) and is_batched:\\n            raise RuntimeError(\\\"the batch number of src and tgt must be equal\\\")\\n\\n        if src.size(-1) != self.d_model or tgt.size(-1) != self.d_model:\\n            raise RuntimeError(\\n                \\\"the feature number of src and tgt must be equal to d_model\\\"\\n            )\\n\\n        memory = self.encoder(\\n            src,\\n            mask=src_mask,\\n            src_key_padding_mask=src_key_padding_mask,\\n            is_causal=src_is_causal,\\n        )\\n        output = self.decoder(\\n            tgt,\\n            memory,\\n            tgt_mask=tgt_mask,\\n            memory_mask=memory_mask,\\n            tgt_key_padding_mask=tgt_key_padding_mask,\\n            memory_key_padding_mask=memory_key_padding_mask,\\n            tgt_is_causal=tgt_is_causal,\\n            memory_is_causal=memory_is_causal,\\n        )\\n        return output\\n\\n    @staticmethod\\n    def generate_square_subsequent_mask(\\n        sz: int,\\n        device: Optional[torch.device] = None,\\n        dtype: Optional[torch.dtype] = None,\\n    ) -> Tensor:\\n        r\\\"\\\"\\\"Generate a square causal mask for the sequence.\\n\\n        The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0).\\n        \\\"\\\"\\\"\\n        return _generate_square_subsequent_mask(sz, dtype=dtype, device=device)\\n\\n    def _reset_parameters(self):\\n        r\\\"\\\"\\\"Initiate parameters in the transformer model.\\\"\\\"\\\"\\n        for p in self.parameters():\\n            if p.dim() > 1:\\n                xavier_uniform_(p)\\n\\n\\nclass TransformerEncoder(Module):\\n    r\\\"\\\"\\\"TransformerEncoder is a stack of N encoder layers.\\n\\n    Users can build the BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.\\n\\n    Args:\\n        encoder_layer: an instance of the TransformerEncoderLayer() class (required).\\n        num_layers: the number of sub-encoder-layers in the encoder (required).\\n        norm: the layer normalization component (optional).\\n        enable_nested_tensor: if True, input will automatically convert to nested tensor\\n            (and convert back on output). This will improve the overall performance of\\n            TransformerEncoder when padding rate is high. Default: ``True`` (enabled).\\n\\n    Examples::\\n        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)\\n        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)\\n        >>> src = torch.rand(10, 32, 512)\\n        >>> out = transformer_encoder(src)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"norm\\\"]\\n\\n    def __init__(\\n        self,\\n        encoder_layer: \\\"TransformerEncoderLayer\\\",\\n        num_layers: int,\\n        norm: Optional[Module] = None,\\n        enable_nested_tensor: bool = True,\\n        mask_check: bool = True,\\n    ) -> None:\\n        super().__init__()\\n        torch._C._log_api_usage_once(f\\\"torch.nn.modules.{self.__class__.__name__}\\\")\\n        self.layers = _get_clones(encoder_layer, num_layers)\\n        self.num_layers = num_layers\\n        self.norm = norm\\n        # this attribute saves the value providedat object construction\\n        self.enable_nested_tensor = enable_nested_tensor\\n        # this attribute controls whether nested tensors are used\\n        self.use_nested_tensor = enable_nested_tensor\\n        self.mask_check = mask_check\\n\\n        enc_layer = \\\"encoder_layer\\\"\\n        why_not_sparsity_fast_path = \\\"\\\"\\n        if not isinstance(encoder_layer, torch.nn.TransformerEncoderLayer):\\n            why_not_sparsity_fast_path = f\\\"{enc_layer} was not TransformerEncoderLayer\\\"\\n        elif encoder_layer.norm_first:\\n            why_not_sparsity_fast_path = f\\\"{enc_layer}.norm_first was True\\\"\\n        elif not encoder_layer.self_attn.batch_first:\\n            why_not_sparsity_fast_path = (\\n                f\\\"{enc_layer}.self_attn.batch_first was not True\\\"\\n                + \\\"(use batch_first for better inference performance)\\\"\\n            )\\n        elif not encoder_layer.self_attn._qkv_same_embed_dim:\\n            why_not_sparsity_fast_path = (\\n                f\\\"{enc_layer}.self_attn._qkv_same_embed_dim was not True\\\"\\n            )\\n        elif encoder_layer.self_attn.in_proj_bias is None:\\n            why_not_sparsity_fast_path = f\\\"{enc_layer}.self_attn was passed bias=False\\\"\\n        elif not encoder_layer.activation_relu_or_gelu:\\n            why_not_sparsity_fast_path = (\\n                f\\\"{enc_layer}.activation_relu_or_gelu was not True\\\"\\n            )\\n        elif not (encoder_layer.norm1.eps == encoder_layer.norm2.eps):\\n            why_not_sparsity_fast_path = (\\n                f\\\"{enc_layer}.norm1.eps was not equal to {enc_layer}.norm2.eps\\\"\\n            )\\n        elif encoder_layer.self_attn.num_heads % 2 == 1:\\n            why_not_sparsity_fast_path = f\\\"{enc_layer}.self_attn.num_heads is odd\\\"\\n\\n        if enable_nested_tensor and why_not_sparsity_fast_path:\\n            warnings.warn(\\n                f\\\"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}\\\"\\n            )\\n            self.use_nested_tensor = False\\n\\n    def forward(\\n        self,\\n        src: Tensor,\\n        mask: Optional[Tensor] = None,\\n        src_key_padding_mask: Optional[Tensor] = None,\\n        is_causal: Optional[bool] = None,\\n    ) -> Tensor:\\n        r\\\"\\\"\\\"Pass the input through the encoder layers in turn.\\n\\n        Args:\\n            src: the sequence to the encoder (required).\\n            mask: the mask for the src sequence (optional).\\n            src_key_padding_mask: the mask for the src keys per batch (optional).\\n            is_causal: If specified, applies a causal mask as ``mask``.\\n                Default: ``None``; try to detect a causal mask.\\n                Warning:\\n                ``is_causal`` provides a hint that ``mask`` is the\\n                causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n\\n        Shape:\\n            see the docs in :class:`~torch.nn.Transformer`.\\n        \\\"\\\"\\\"\\n        src_key_padding_mask = F._canonical_mask(\\n            mask=src_key_padding_mask,\\n            mask_name=\\\"src_key_padding_mask\\\",\\n            other_type=F._none_or_dtype(mask),\\n            other_name=\\\"mask\\\",\\n            target_type=src.dtype,\\n        )\\n\\n        mask = F._canonical_mask(\\n            mask=mask,\\n            mask_name=\\\"mask\\\",\\n            other_type=None,\\n            other_name=\\\"\\\",\\n            target_type=src.dtype,\\n            check_other=False,\\n        )\\n\\n        output = src\\n        convert_to_nested = False\\n        first_layer = self.layers[0]\\n        src_key_padding_mask_for_layers = src_key_padding_mask\\n        why_not_sparsity_fast_path = \\\"\\\"\\n        str_first_layer = \\\"self.layers[0]\\\"\\n        batch_first = first_layer.self_attn.batch_first\\n        is_fastpath_enabled = torch.backends.mha.get_fastpath_enabled()\\n\\n        if not is_fastpath_enabled:\\n            why_not_sparsity_fast_path = (\\n                \\\"torch.backends.mha.get_fastpath_enabled() was not True\\\"\\n            )\\n        elif not hasattr(self, \\\"use_nested_tensor\\\"):\\n            why_not_sparsity_fast_path = \\\"use_nested_tensor attribute not present\\\"\\n        elif not self.use_nested_tensor:\\n            why_not_sparsity_fast_path = (\\n                \\\"self.use_nested_tensor (set in init) was not True\\\"\\n            )\\n        elif first_layer.training:\\n            why_not_sparsity_fast_path = f\\\"{str_first_layer} was in training mode\\\"\\n        elif not src.dim() == 3:\\n            why_not_sparsity_fast_path = (\\n                f\\\"input not batched; expected src.dim() of 3 but got {src.dim()}\\\"\\n            )\\n        elif src_key_padding_mask is None:\\n            why_not_sparsity_fast_path = \\\"src_key_padding_mask was None\\\"\\n        elif (\\n            (not hasattr(self, \\\"mask_check\\\")) or self.mask_check\\n        ) and not torch._nested_tensor_from_mask_left_aligned(\\n            src, src_key_padding_mask.logical_not()\\n        ):\\n            why_not_sparsity_fast_path = \\\"mask_check enabled, and src and src_key_padding_mask was not left aligned\\\"\\n        elif output.is_nested:\\n            why_not_sparsity_fast_path = \\\"NestedTensor input is not supported\\\"\\n        elif mask is not None:\\n            why_not_sparsity_fast_path = (\\n                \\\"src_key_padding_mask and mask were both supplied\\\"\\n            )\\n        elif torch.is_autocast_enabled():\\n            why_not_sparsity_fast_path = \\\"autocast is enabled\\\"\\n\\n        if not why_not_sparsity_fast_path:\\n            tensor_args = (\\n                src,\\n                first_layer.self_attn.in_proj_weight,\\n                first_layer.self_attn.in_proj_bias,\\n                first_layer.self_attn.out_proj.weight,\\n                first_layer.self_attn.out_proj.bias,\\n                first_layer.norm1.weight,\\n                first_layer.norm1.bias,\\n                first_layer.norm2.weight,\\n                first_layer.norm2.bias,\\n                first_layer.linear1.weight,\\n                first_layer.linear1.bias,\\n                first_layer.linear2.weight,\\n                first_layer.linear2.bias,\\n            )\\n            _supported_device_type = [\\n                \\\"cpu\\\",\\n                \\\"cuda\\\",\\n                torch.utils.backend_registration._privateuse1_backend_name,\\n            ]\\n            if torch.overrides.has_torch_function(tensor_args):\\n                why_not_sparsity_fast_path = \\\"some Tensor argument has_torch_function\\\"\\n            elif src.device.type not in _supported_device_type:\\n                why_not_sparsity_fast_path = (\\n                    f\\\"src device is neither one of {_supported_device_type}\\\"\\n                )\\n            elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args):\\n                why_not_sparsity_fast_path = (\\n                    \\\"grad is enabled and at least one of query or the \\\"\\n                    \\\"input/output projection weights or biases requires_grad\\\"\\n                )\\n\\n            if (not why_not_sparsity_fast_path) and (src_key_padding_mask is not None):\\n                convert_to_nested = True\\n                output = torch._nested_tensor_from_mask(\\n                    output, src_key_padding_mask.logical_not(), mask_check=False\\n                )\\n                src_key_padding_mask_for_layers = None\\n\\n        seq_len = _get_seq_len(src, batch_first)\\n        is_causal = _detect_is_causal_mask(mask, is_causal, seq_len)\\n\\n        for mod in self.layers:\\n            output = mod(\\n                output,\\n                src_mask=mask,\\n                is_causal=is_causal,\\n                src_key_padding_mask=src_key_padding_mask_for_layers,\\n            )\\n\\n        if convert_to_nested:\\n            output = output.to_padded_tensor(0.0, src.size())\\n\\n        if self.norm is not None:\\n            output = self.norm(output)\\n\\n        return output\\n\\n\\nclass TransformerDecoder(Module):\\n    r\\\"\\\"\\\"TransformerDecoder is a stack of N decoder layers.\\n\\n    Args:\\n        decoder_layer: an instance of the TransformerDecoderLayer() class (required).\\n        num_layers: the number of sub-decoder-layers in the decoder (required).\\n        norm: the layer normalization component (optional).\\n\\n    Examples::\\n        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)\\n        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)\\n        >>> memory = torch.rand(10, 32, 512)\\n        >>> tgt = torch.rand(20, 32, 512)\\n        >>> out = transformer_decoder(tgt, memory)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"norm\\\"]\\n\\n    def __init__(\\n        self,\\n        decoder_layer: \\\"TransformerDecoderLayer\\\",\\n        num_layers: int,\\n        norm: Optional[Module] = None,\\n    ) -> None:\\n        super().__init__()\\n        torch._C._log_api_usage_once(f\\\"torch.nn.modules.{self.__class__.__name__}\\\")\\n        self.layers = _get_clones(decoder_layer, num_layers)\\n        self.num_layers = num_layers\\n        self.norm = norm\\n\\n    def forward(\\n        self,\\n        tgt: Tensor,\\n        memory: Tensor,\\n        tgt_mask: Optional[Tensor] = None,\\n        memory_mask: Optional[Tensor] = None,\\n        tgt_key_padding_mask: Optional[Tensor] = None,\\n        memory_key_padding_mask: Optional[Tensor] = None,\\n        tgt_is_causal: Optional[bool] = None,\\n        memory_is_causal: bool = False,\\n    ) -> Tensor:\\n        r\\\"\\\"\\\"Pass the inputs (and mask) through the decoder layer in turn.\\n\\n        Args:\\n            tgt: the sequence to the decoder (required).\\n            memory: the sequence from the last layer of the encoder (required).\\n            tgt_mask: the mask for the tgt sequence (optional).\\n            memory_mask: the mask for the memory sequence (optional).\\n            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).\\n            memory_key_padding_mask: the mask for the memory keys per batch (optional).\\n            tgt_is_causal: If specified, applies a causal mask as ``tgt mask``.\\n                Default: ``None``; try to detect a causal mask.\\n                Warning:\\n                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is\\n                the causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n            memory_is_causal: If specified, applies a causal mask as\\n                ``memory mask``.\\n                Default: ``False``.\\n                Warning:\\n                ``memory_is_causal`` provides a hint that\\n                ``memory_mask`` is the causal mask. Providing incorrect\\n                hints can result in incorrect execution, including\\n                forward and backward compatibility.\\n\\n        Shape:\\n            see the docs in :class:`~torch.nn.Transformer`.\\n        \\\"\\\"\\\"\\n        output = tgt\\n\\n        seq_len = _get_seq_len(tgt, self.layers[0].self_attn.batch_first)\\n        tgt_is_causal = _detect_is_causal_mask(tgt_mask, tgt_is_causal, seq_len)\\n\\n        for mod in self.layers:\\n            output = mod(\\n                output,\\n                memory,\\n                tgt_mask=tgt_mask,\\n                memory_mask=memory_mask,\\n                tgt_key_padding_mask=tgt_key_padding_mask,\\n                memory_key_padding_mask=memory_key_padding_mask,\\n                tgt_is_causal=tgt_is_causal,\\n                memory_is_causal=memory_is_causal,\\n            )\\n\\n        if self.norm is not None:\\n            output = self.norm(output)\\n\\n        return output\\n\\n\\nclass TransformerEncoderLayer(Module):\\n    r\\\"\\\"\\\"TransformerEncoderLayer is made up of self-attn and feedforward network.\\n\\n    This standard encoder layer is based on the paper \\\"Attention Is All You Need\\\".\\n    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,\\n    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in\\n    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement\\n    in a different way during application.\\n\\n    TransformerEncoderLayer can handle either traditional torch.tensor inputs,\\n    or Nested Tensor inputs.  Derived classes are expected to similarly accept\\n    both input formats.  (Not all combinations of inputs are currently\\n    supported by TransformerEncoderLayer while Nested Tensor is in prototype\\n    state.)\\n\\n    If you are implementing a custom layer, you may derive it either from\\n    the Module or TransformerEncoderLayer class.  If your custom layer\\n    supports both torch.Tensors and Nested Tensors inputs, make its\\n    implementation a derived class of TransformerEncoderLayer. If your custom\\n    Layer supports only torch.Tensor inputs, derive its implementation from\\n    Module.\\n\\n    Args:\\n        d_model: the number of expected features in the input (required).\\n        nhead: the number of heads in the multiheadattention models (required).\\n        dim_feedforward: the dimension of the feedforward network model (default=2048).\\n        dropout: the dropout value (default=0.1).\\n        activation: the activation function of the intermediate layer, can be a string\\n            (\\\"relu\\\" or \\\"gelu\\\") or a unary callable. Default: relu\\n        layer_norm_eps: the eps value in layer normalization components (default=1e-5).\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).\\n        norm_first: if ``True``, layer norm is done prior to attention and feedforward\\n            operations, respectively. Otherwise it's done after. Default: ``False`` (after).\\n        bias: If set to ``False``, ``Linear`` and ``LayerNorm`` layers will not learn an additive\\n            bias. Default: ``True``.\\n\\n    Examples::\\n        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)\\n        >>> src = torch.rand(10, 32, 512)\\n        >>> out = encoder_layer(src)\\n\\n    Alternatively, when ``batch_first`` is ``True``:\\n        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)\\n        >>> src = torch.rand(32, 10, 512)\\n        >>> out = encoder_layer(src)\\n\\n    Fast path:\\n        forward() will use a special optimized implementation described in\\n        `FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`_ if all of the following\\n        conditions are met:\\n\\n        - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor\\n          argument ``requires_grad``\\n        - training is disabled (using ``.eval()``)\\n        - batch_first is ``True`` and the input is batched (i.e., ``src.dim() == 3``)\\n        - activation is one of: ``\\\"relu\\\"``, ``\\\"gelu\\\"``, ``torch.functional.relu``, or ``torch.functional.gelu``\\n        - at most one of ``src_mask`` and ``src_key_padding_mask`` is passed\\n        - if src is a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_, neither ``src_mask``\\n          nor ``src_key_padding_mask`` is passed\\n        - the two ``LayerNorm`` instances have a consistent ``eps`` value (this will naturally be the case\\n          unless the caller has manually modified one without modifying the other)\\n\\n        If the optimized implementation is in use, a\\n        `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be\\n        passed for ``src`` to represent padding more efficiently than using a padding\\n        mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ will be\\n        returned, and an additional speedup proportional to the fraction of the input that\\n        is padding can be expected.\\n\\n        .. _`FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`:\\n         https://arxiv.org/abs/2205.14135\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"norm_first\\\"]\\n\\n    def __init__(\\n        self,\\n        d_model: int,\\n        nhead: int,\\n        dim_feedforward: int = 2048,\\n        dropout: float = 0.1,\\n        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,\\n        layer_norm_eps: float = 1e-5,\\n        batch_first: bool = False,\\n        norm_first: bool = False,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.self_attn = MultiheadAttention(\\n            d_model,\\n            nhead,\\n            dropout=dropout,\\n            bias=bias,\\n            batch_first=batch_first,\\n            **factory_kwargs,\\n        )\\n        # Implementation of Feedforward model\\n        self.linear1 = Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs)\\n        self.dropout = Dropout(dropout)\\n        self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)\\n\\n        self.norm_first = norm_first\\n        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)\\n        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)\\n        self.dropout1 = Dropout(dropout)\\n        self.dropout2 = Dropout(dropout)\\n\\n        # Legacy string support for activation function.\\n        if isinstance(activation, str):\\n            activation = _get_activation_fn(activation)\\n\\n        # We can't test self.activation in forward() in TorchScript,\\n        # so stash some information about it instead.\\n        if activation is F.relu or isinstance(activation, torch.nn.ReLU):\\n            self.activation_relu_or_gelu = 1\\n        elif activation is F.gelu or isinstance(activation, torch.nn.GELU):\\n            self.activation_relu_or_gelu = 2\\n        else:\\n            self.activation_relu_or_gelu = 0\\n        self.activation = activation\\n\\n    def __setstate__(self, state):\\n        super().__setstate__(state)\\n        if not hasattr(self, \\\"activation\\\"):\\n            self.activation = F.relu\\n\\n    def forward(\\n        self,\\n        src: Tensor,\\n        src_mask: Optional[Tensor] = None,\\n        src_key_padding_mask: Optional[Tensor] = None,\\n        is_causal: bool = False,\\n    ) -> Tensor:\\n        r\\\"\\\"\\\"Pass the input through the encoder layer.\\n\\n        Args:\\n            src: the sequence to the encoder layer (required).\\n            src_mask: the mask for the src sequence (optional).\\n            src_key_padding_mask: the mask for the src keys per batch (optional).\\n            is_causal: If specified, applies a causal mask as ``src mask``.\\n                Default: ``False``.\\n                Warning:\\n                ``is_causal`` provides a hint that ``src_mask`` is the\\n                causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n\\n        Shape:\\n            see the docs in :class:`~torch.nn.Transformer`.\\n        \\\"\\\"\\\"\\n        src_key_padding_mask = F._canonical_mask(\\n            mask=src_key_padding_mask,\\n            mask_name=\\\"src_key_padding_mask\\\",\\n            other_type=F._none_or_dtype(src_mask),\\n            other_name=\\\"src_mask\\\",\\n            target_type=src.dtype,\\n        )\\n\\n        src_mask = F._canonical_mask(\\n            mask=src_mask,\\n            mask_name=\\\"src_mask\\\",\\n            other_type=None,\\n            other_name=\\\"\\\",\\n            target_type=src.dtype,\\n            check_other=False,\\n        )\\n\\n        is_fastpath_enabled = torch.backends.mha.get_fastpath_enabled()\\n\\n        why_not_sparsity_fast_path = \\\"\\\"\\n        if not is_fastpath_enabled:\\n            why_not_sparsity_fast_path = (\\n                \\\"torch.backends.mha.get_fastpath_enabled() was not True\\\"\\n            )\\n        elif not src.dim() == 3:\\n            why_not_sparsity_fast_path = (\\n                f\\\"input not batched; expected src.dim() of 3 but got {src.dim()}\\\"\\n            )\\n        elif self.training:\\n            why_not_sparsity_fast_path = \\\"training is enabled\\\"\\n        elif not self.self_attn.batch_first:\\n            why_not_sparsity_fast_path = \\\"self_attn.batch_first was not True\\\"\\n        elif self.self_attn.in_proj_bias is None:\\n            why_not_sparsity_fast_path = \\\"self_attn was passed bias=False\\\"\\n        elif not self.self_attn._qkv_same_embed_dim:\\n            why_not_sparsity_fast_path = \\\"self_attn._qkv_same_embed_dim was not True\\\"\\n        elif not self.activation_relu_or_gelu:\\n            why_not_sparsity_fast_path = \\\"activation_relu_or_gelu was not True\\\"\\n        elif not (self.norm1.eps == self.norm2.eps):\\n            why_not_sparsity_fast_path = \\\"norm1.eps is not equal to norm2.eps\\\"\\n        elif src.is_nested and (\\n            src_key_padding_mask is not None or src_mask is not None\\n        ):\\n            why_not_sparsity_fast_path = \\\"neither src_key_padding_mask nor src_mask are not supported with NestedTensor input\\\"\\n        elif self.self_attn.num_heads % 2 == 1:\\n            why_not_sparsity_fast_path = \\\"num_head is odd\\\"\\n        elif torch.is_autocast_enabled():\\n            why_not_sparsity_fast_path = \\\"autocast is enabled\\\"\\n        elif any(\\n            len(getattr(m, \\\"_forward_hooks\\\", {}))\\n            + len(getattr(m, \\\"_forward_pre_hooks\\\", {}))\\n            for m in self.modules()\\n        ):\\n            why_not_sparsity_fast_path = \\\"forward pre-/hooks are attached to the module\\\"\\n        if not why_not_sparsity_fast_path:\\n            tensor_args = (\\n                src,\\n                self.self_attn.in_proj_weight,\\n                self.self_attn.in_proj_bias,\\n                self.self_attn.out_proj.weight,\\n                self.self_attn.out_proj.bias,\\n                self.norm1.weight,\\n                self.norm1.bias,\\n                self.norm2.weight,\\n                self.norm2.bias,\\n                self.linear1.weight,\\n                self.linear1.bias,\\n                self.linear2.weight,\\n                self.linear2.bias,\\n            )\\n\\n            # We have to use list comprehensions below because TorchScript does not support\\n            # generator expressions.\\n            _supported_device_type = [\\n                \\\"cpu\\\",\\n                \\\"cuda\\\",\\n                torch.utils.backend_registration._privateuse1_backend_name,\\n            ]\\n            if torch.overrides.has_torch_function(tensor_args):\\n                why_not_sparsity_fast_path = \\\"some Tensor argument has_torch_function\\\"\\n            elif not all(\\n                (x.device.type in _supported_device_type) for x in tensor_args\\n            ):\\n                why_not_sparsity_fast_path = (\\n                    \\\"some Tensor argument's device is neither one of \\\"\\n                    f\\\"{_supported_device_type}\\\"\\n                )\\n            elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args):\\n                why_not_sparsity_fast_path = (\\n                    \\\"grad is enabled and at least one of query or the \\\"\\n                    \\\"input/output projection weights or biases requires_grad\\\"\\n                )\\n\\n            if not why_not_sparsity_fast_path:\\n                merged_mask, mask_type = self.self_attn.merge_masks(\\n                    src_mask, src_key_padding_mask, src\\n                )\\n                return torch._transformer_encoder_layer_fwd(\\n                    src,\\n                    self.self_attn.embed_dim,\\n                    self.self_attn.num_heads,\\n                    self.self_attn.in_proj_weight,\\n                    self.self_attn.in_proj_bias,\\n                    self.self_attn.out_proj.weight,\\n                    self.self_attn.out_proj.bias,\\n                    self.activation_relu_or_gelu == 2,\\n                    self.norm_first,\\n                    self.norm1.eps,\\n                    self.norm1.weight,\\n                    self.norm1.bias,\\n                    self.norm2.weight,\\n                    self.norm2.bias,\\n                    self.linear1.weight,\\n                    self.linear1.bias,\\n                    self.linear2.weight,\\n                    self.linear2.bias,\\n                    merged_mask,\\n                    mask_type,\\n                )\\n\\n        # see Fig. 1 of https://arxiv.org/pdf/2002.04745v1.pdf\\n        x = src\\n        if self.norm_first:\\n            x = x + self._sa_block(\\n                self.norm1(x), src_mask, src_key_padding_mask, is_causal=is_causal\\n            )\\n            x = x + self._ff_block(self.norm2(x))\\n        else:\\n            x = self.norm1(\\n                x\\n                + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal)\\n            )\\n            x = self.norm2(x + self._ff_block(x))\\n\\n        return x\\n\\n    # self-attention block\\n    def _sa_block(\\n        self,\\n        x: Tensor,\\n        attn_mask: Optional[Tensor],\\n        key_padding_mask: Optional[Tensor],\\n        is_causal: bool = False,\\n    ) -> Tensor:\\n        x = self.self_attn(\\n            x,\\n            x,\\n            x,\\n            attn_mask=attn_mask,\\n            key_padding_mask=key_padding_mask,\\n            need_weights=False,\\n            is_causal=is_causal,\\n        )[0]\\n        return self.dropout1(x)\\n\\n    # feed forward block\\n    def _ff_block(self, x: Tensor) -> Tensor:\\n        x = self.linear2(self.dropout(self.activation(self.linear1(x))))\\n        return self.dropout2(x)\\n\\n\\nclass TransformerDecoderLayer(Module):\\n    r\\\"\\\"\\\"TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.\\n\\n    This standard decoder layer is based on the paper \\\"Attention Is All You Need\\\".\\n    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,\\n    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in\\n    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement\\n    in a different way during application.\\n\\n    Args:\\n        d_model: the number of expected features in the input (required).\\n        nhead: the number of heads in the multiheadattention models (required).\\n        dim_feedforward: the dimension of the feedforward network model (default=2048).\\n        dropout: the dropout value (default=0.1).\\n        activation: the activation function of the intermediate layer, can be a string\\n            (\\\"relu\\\" or \\\"gelu\\\") or a unary callable. Default: relu\\n        layer_norm_eps: the eps value in layer normalization components (default=1e-5).\\n        batch_first: If ``True``, then the input and output tensors are provided\\n            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).\\n        norm_first: if ``True``, layer norm is done prior to self attention, multihead\\n            attention and feedforward operations, respectively. Otherwise it's done after.\\n            Default: ``False`` (after).\\n        bias: If set to ``False``, ``Linear`` and ``LayerNorm`` layers will not learn an additive\\n            bias. Default: ``True``.\\n\\n    Examples::\\n        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)\\n        >>> memory = torch.rand(10, 32, 512)\\n        >>> tgt = torch.rand(20, 32, 512)\\n        >>> out = decoder_layer(tgt, memory)\\n\\n    Alternatively, when ``batch_first`` is ``True``:\\n        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)\\n        >>> memory = torch.rand(32, 10, 512)\\n        >>> tgt = torch.rand(32, 20, 512)\\n        >>> out = decoder_layer(tgt, memory)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"norm_first\\\"]\\n\\n    def __init__(\\n        self,\\n        d_model: int,\\n        nhead: int,\\n        dim_feedforward: int = 2048,\\n        dropout: float = 0.1,\\n        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,\\n        layer_norm_eps: float = 1e-5,\\n        batch_first: bool = False,\\n        norm_first: bool = False,\\n        bias: bool = True,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.self_attn = MultiheadAttention(\\n            d_model,\\n            nhead,\\n            dropout=dropout,\\n            batch_first=batch_first,\\n            bias=bias,\\n            **factory_kwargs,\\n        )\\n        self.multihead_attn = MultiheadAttention(\\n            d_model,\\n            nhead,\\n            dropout=dropout,\\n            batch_first=batch_first,\\n            bias=bias,\\n            **factory_kwargs,\\n        )\\n        # Implementation of Feedforward model\\n        self.linear1 = Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs)\\n        self.dropout = Dropout(dropout)\\n        self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)\\n\\n        self.norm_first = norm_first\\n        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)\\n        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)\\n        self.norm3 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)\\n        self.dropout1 = Dropout(dropout)\\n        self.dropout2 = Dropout(dropout)\\n        self.dropout3 = Dropout(dropout)\\n\\n        # Legacy string support for activation function.\\n        if isinstance(activation, str):\\n            self.activation = _get_activation_fn(activation)\\n        else:\\n            self.activation = activation\\n\\n    def __setstate__(self, state):\\n        if \\\"activation\\\" not in state:\\n            state[\\\"activation\\\"] = F.relu\\n        super().__setstate__(state)\\n\\n    def forward(\\n        self,\\n        tgt: Tensor,\\n        memory: Tensor,\\n        tgt_mask: Optional[Tensor] = None,\\n        memory_mask: Optional[Tensor] = None,\\n        tgt_key_padding_mask: Optional[Tensor] = None,\\n        memory_key_padding_mask: Optional[Tensor] = None,\\n        tgt_is_causal: bool = False,\\n        memory_is_causal: bool = False,\\n    ) -> Tensor:\\n        r\\\"\\\"\\\"Pass the inputs (and mask) through the decoder layer.\\n\\n        Args:\\n            tgt: the sequence to the decoder layer (required).\\n            memory: the sequence from the last layer of the encoder (required).\\n            tgt_mask: the mask for the tgt sequence (optional).\\n            memory_mask: the mask for the memory sequence (optional).\\n            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).\\n            memory_key_padding_mask: the mask for the memory keys per batch (optional).\\n            tgt_is_causal: If specified, applies a causal mask as ``tgt mask``.\\n                Default: ``False``.\\n                Warning:\\n                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is\\n                the causal mask. Providing incorrect hints can result in\\n                incorrect execution, including forward and backward\\n                compatibility.\\n            memory_is_causal: If specified, applies a causal mask as\\n                ``memory mask``.\\n                Default: ``False``.\\n                Warning:\\n                ``memory_is_causal`` provides a hint that\\n                ``memory_mask`` is the causal mask. Providing incorrect\\n                hints can result in incorrect execution, including\\n                forward and backward compatibility.\\n\\n        Shape:\\n            see the docs in :class:`~torch.nn.Transformer`.\\n        \\\"\\\"\\\"\\n        # see Fig. 1 of https://arxiv.org/pdf/2002.04745v1.pdf\\n\\n        x = tgt\\n        if self.norm_first:\\n            x = x + self._sa_block(\\n                self.norm1(x), tgt_mask, tgt_key_padding_mask, tgt_is_causal\\n            )\\n            x = x + self._mha_block(\\n                self.norm2(x),\\n                memory,\\n                memory_mask,\\n                memory_key_padding_mask,\\n                memory_is_causal,\\n            )\\n            x = x + self._ff_block(self.norm3(x))\\n        else:\\n            x = self.norm1(\\n                x + self._sa_block(x, tgt_mask, tgt_key_padding_mask, tgt_is_causal)\\n            )\\n            x = self.norm2(\\n                x\\n                + self._mha_block(\\n                    x, memory, memory_mask, memory_key_padding_mask, memory_is_causal\\n                )\\n            )\\n            x = self.norm3(x + self._ff_block(x))\\n\\n        return x\\n\\n    # self-attention block\\n    def _sa_block(\\n        self,\\n        x: Tensor,\\n        attn_mask: Optional[Tensor],\\n        key_padding_mask: Optional[Tensor],\\n        is_causal: bool = False,\\n    ) -> Tensor:\\n        x = self.self_attn(\\n            x,\\n            x,\\n            x,\\n            attn_mask=attn_mask,\\n            key_padding_mask=key_padding_mask,\\n            is_causal=is_causal,\\n            need_weights=False,\\n        )[0]\\n        return self.dropout1(x)\\n\\n    # multihead attention block\\n    def _mha_block(\\n        self,\\n        x: Tensor,\\n        mem: Tensor,\\n        attn_mask: Optional[Tensor],\\n        key_padding_mask: Optional[Tensor],\\n        is_causal: bool = False,\\n    ) -> Tensor:\\n        x = self.multihead_attn(\\n            x,\\n            mem,\\n            mem,\\n            attn_mask=attn_mask,\\n            key_padding_mask=key_padding_mask,\\n            is_causal=is_causal,\\n            need_weights=False,\\n        )[0]\\n        return self.dropout2(x)\\n\\n    # feed forward block\\n    def _ff_block(self, x: Tensor) -> Tensor:\\n        x = self.linear2(self.dropout(self.activation(self.linear1(x))))\\n        return self.dropout3(x)\\n\\n\\ndef _get_clones(module, N):\\n    # FIXME: copy.deepcopy() is not defined on nn.module\\n    return ModuleList([copy.deepcopy(module) for i in range(N)])\\n\\n\\ndef _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:\\n    if activation == \\\"relu\\\":\\n        return F.relu\\n    elif activation == \\\"gelu\\\":\\n        return F.gelu\\n\\n    raise RuntimeError(f\\\"activation should be relu/gelu, not {activation}\\\")\\n\\n\\ndef _detect_is_causal_mask(\\n    mask: Optional[Tensor],\\n    is_causal: Optional[bool] = None,\\n    size: Optional[int] = None,\\n) -> bool:\\n    \\\"\\\"\\\"Return whether the given attention mask is causal.\\n\\n    Warning:\\n    If ``is_causal`` is not ``None``, its value will be returned as is.  If a\\n    user supplies an incorrect ``is_causal`` hint,\\n\\n    ``is_causal=False`` when the mask is in fact a causal attention.mask\\n       may lead to reduced performance relative to what would be achievable\\n       with ``is_causal=True``;\\n    ``is_causal=True`` when the mask is in fact not a causal attention.mask\\n       may lead to incorrect and unpredictable execution - in some scenarios,\\n       a causal mask may be applied based on the hint, in other execution\\n       scenarios the specified mask may be used.  The choice may not appear\\n       to be deterministic, in that a number of factors like alignment,\\n       hardware SKU, etc influence the decision whether to use a mask or\\n       rely on the hint.\\n    ``size`` if not None, check whether the mask is a causal mask of the provided size\\n       Otherwise, checks for any causal mask.\\n    \\\"\\\"\\\"\\n    # Prevent type refinement\\n    make_causal = is_causal is True\\n\\n    if is_causal is None and mask is not None:\\n        sz = size if size is not None else mask.size(-2)\\n        causal_comparison = _generate_square_subsequent_mask(\\n            sz, device=mask.device, dtype=mask.dtype\\n        )\\n\\n        # Do not use `torch.equal` so we handle batched masks by\\n        # broadcasting the comparison.\\n        if mask.size() == causal_comparison.size():\\n            make_causal = bool((mask == causal_comparison).all())\\n        else:\\n            make_causal = False\\n\\n    return make_causal\\n\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.common_types import _size_any_t\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"Fold\\\", \\\"Unfold\\\"]\\n\\n\\nclass Fold(Module):\\n    r\\\"\\\"\\\"Combines an array of sliding local blocks into a large containing tensor.\\n\\n    Consider a batched :attr:`input` tensor containing sliding local blocks,\\n    e.g., patches of images, of shape :math:`(N, C \\\\times  \\\\prod(\\\\text{kernel\\\\_size}), L)`,\\n    where :math:`N` is batch dimension, :math:`C \\\\times \\\\prod(\\\\text{kernel\\\\_size})`\\n    is the number of values within a block (a block has :math:`\\\\prod(\\\\text{kernel\\\\_size})`\\n    spatial locations each containing a :math:`C`-channeled vector), and\\n    :math:`L` is the total number of blocks. (This is exactly the\\n    same specification as the output shape of :class:`~torch.nn.Unfold`.) This\\n    operation combines these local blocks into the large :attr:`output` tensor\\n    of shape :math:`(N, C, \\\\text{output\\\\_size}[0], \\\\text{output\\\\_size}[1], \\\\dots)`\\n    by summing the overlapping values. Similar to :class:`~torch.nn.Unfold`, the\\n    arguments must satisfy\\n\\n    .. math::\\n        L = \\\\prod_d \\\\left\\\\lfloor\\\\frac{\\\\text{output\\\\_size}[d] + 2 \\\\times \\\\text{padding}[d] %\\n            - \\\\text{dilation}[d] \\\\times (\\\\text{kernel\\\\_size}[d] - 1) - 1}{\\\\text{stride}[d]} + 1\\\\right\\\\rfloor,\\n\\n    where :math:`d` is over all spatial dimensions.\\n\\n    * :attr:`output_size` describes the spatial shape of the large containing\\n      tensor of the sliding local blocks. It is useful to resolve the ambiguity\\n      when multiple input shapes map to same number of sliding blocks, e.g.,\\n      with ``stride > 0``.\\n\\n    The :attr:`padding`, :attr:`stride` and :attr:`dilation` arguments specify\\n    how the sliding blocks are retrieved.\\n\\n    * :attr:`stride` controls the stride for the sliding blocks.\\n\\n    * :attr:`padding` controls the amount of implicit zero-paddings on both\\n      sides for :attr:`padding` number of points for each dimension before\\n      reshaping.\\n\\\"\\\"\\\" \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also known as the \\\\u00e0 trous algorithm.\\n      It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\" r\\\"\\\"\\\"\\n    Args:\\n        output_size (int or tuple): the shape of the spatial dimensions of the\\n                                    output (i.e., ``output.sizes()[2:]``)\\n        kernel_size (int or tuple): the size of the sliding blocks\\n        dilation (int or tuple, optional): a parameter that controls the\\n                                           stride of elements within the\\n                                           neighborhood. Default: 1\\n        padding (int or tuple, optional): implicit zero padding to be added on\\n                                          both sides of input. Default: 0\\n        stride (int or tuple): the stride of the sliding blocks in the input\\n                               spatial dimensions. Default: 1\\n\\n    * If :attr:`output_size`, :attr:`kernel_size`, :attr:`dilation`,\\n      :attr:`padding` or :attr:`stride` is an int or a tuple of length 1 then\\n      their values will be replicated across all spatial dimensions.\\n\\n    * For the case of two output spatial dimensions this operation is sometimes\\n      called ``col2im``.\\n\\n    .. note::\\n        :class:`~torch.nn.Fold` calculates each combined value in the resulting\\n        large tensor by summing all values from all containing blocks.\\n        :class:`~torch.nn.Unfold` extracts the values in the local blocks by\\n        copying from the large tensor. So, if the blocks overlap, they are not\\n        inverses of each other.\\n\\n        In general, folding and unfolding operations are related as\\n        follows. Consider :class:`~torch.nn.Fold` and\\n        :class:`~torch.nn.Unfold` instances created with the same\\n        parameters:\\n\\n        >>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)\\n        >>> fold = nn.Fold(output_size=..., **fold_params)\\n        >>> unfold = nn.Unfold(**fold_params)\\n\\n        Then for any (supported) ``input`` tensor the following\\n        equality holds:\\n\\n        ::\\n\\n            fold(unfold(input)) == divisor * input\\n\\n        where ``divisor`` is a tensor that depends only on the shape\\n        and dtype of the ``input``:\\n\\n        >>> # xdoctest: +SKIP\\n        >>> input_ones = torch.ones(input.shape, dtype=input.dtype)\\n        >>> divisor = fold(unfold(input_ones))\\n\\n        When the ``divisor`` tensor contains no zero elements, then\\n        ``fold`` and ``unfold`` operations are inverses of each\\n        other (up to constant divisor).\\n\\n    .. warning::\\n        Currently, only unbatched (3D) or batched (4D) image-like output tensors are supported.\\n\\n    Shape:\\n        - Input: :math:`(N, C \\\\times \\\\prod(\\\\text{kernel\\\\_size}), L)` or :math:`(C \\\\times \\\\prod(\\\\text{kernel\\\\_size}), L)`\\n        - Output: :math:`(N, C, \\\\text{output\\\\_size}[0], \\\\text{output\\\\_size}[1], \\\\dots)`\\n          or :math:`(C, \\\\text{output\\\\_size}[0], \\\\text{output\\\\_size}[1], \\\\dots)` as described above\\n\\n    Examples::\\n\\n        >>> fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))\\n        >>> input = torch.randn(1, 3 * 2 * 2, 12)\\n        >>> output = fold(input)\\n        >>> output.size()\\n        torch.Size([1, 3, 4, 5])\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"output_size\\\", \\\"kernel_size\\\", \\\"dilation\\\", \\\"padding\\\", \\\"stride\\\"]\\n    output_size: _size_any_t\\n    kernel_size: _size_any_t\\n    dilation: _size_any_t\\n    padding: _size_any_t\\n    stride: _size_any_t\\n\\n    def __init__(\\n        self,\\n        output_size: _size_any_t,\\n        kernel_size: _size_any_t,\\n        dilation: _size_any_t = 1,\\n        padding: _size_any_t = 0,\\n        stride: _size_any_t = 1,\\n    ) -> None:\\n        super().__init__()\\n        self.output_size = output_size\\n        self.kernel_size = kernel_size\\n        self.dilation = dilation\\n        self.padding = padding\\n        self.stride = stride\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.fold(\\n            input,\\n            self.output_size,\\n            self.kernel_size,\\n            self.dilation,\\n            self.padding,\\n            self.stride,\\n        )\\n\\n    def extra_repr(self) -> str:\\n        return (\\n            \\\"output_size={output_size}, kernel_size={kernel_size}, \\\"\\n            \\\"dilation={dilation}, padding={padding}, stride={stride}\\\".format(\\n                **self.__dict__\\n            )\\n        )\\n\\n\\nclass Unfold(Module):\\n    r\\\"\\\"\\\"Extracts sliding local blocks from a batched input tensor.\\n\\n    Consider a batched :attr:`input` tensor of shape :math:`(N, C, *)`,\\n    where :math:`N` is the batch dimension, :math:`C` is the channel dimension,\\n    and :math:`*` represent arbitrary spatial dimensions. This operation flattens\\n    each sliding :attr:`kernel_size`-sized block within the spatial dimensions\\n    of :attr:`input` into a column (i.e., last dimension) of a 3-D :attr:`output`\\n    tensor of shape :math:`(N, C \\\\times \\\\prod(\\\\text{kernel\\\\_size}), L)`, where\\n    :math:`C \\\\times \\\\prod(\\\\text{kernel\\\\_size})` is the total number of values\\n    within each block (a block has :math:`\\\\prod(\\\\text{kernel\\\\_size})` spatial\\n    locations each containing a :math:`C`-channeled vector), and :math:`L` is\\n    the total number of such blocks:\\n\\n    .. math::\\n        L = \\\\prod_d \\\\left\\\\lfloor\\\\frac{\\\\text{spatial\\\\_size}[d] + 2 \\\\times \\\\text{padding}[d] %\\n            - \\\\text{dilation}[d] \\\\times (\\\\text{kernel\\\\_size}[d] - 1) - 1}{\\\\text{stride}[d]} + 1\\\\right\\\\rfloor,\\n\\n    where :math:`\\\\text{spatial\\\\_size}` is formed by the spatial dimensions\\n    of :attr:`input` (:math:`*` above), and :math:`d` is over all spatial\\n    dimensions.\\n\\n    Therefore, indexing :attr:`output` at the last dimension (column dimension)\\n    gives all values within a certain block.\\n\\n    The :attr:`padding`, :attr:`stride` and :attr:`dilation` arguments specify\\n    how the sliding blocks are retrieved.\\n\\n    * :attr:`stride` controls the stride for the sliding blocks.\\n\\n    * :attr:`padding` controls the amount of implicit zero-paddings on both\\n      sides for :attr:`padding` number of points for each dimension before\\n      reshaping.\\n\\\"\\\"\\\" \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also known as the \\\\u00e0 trous algorithm.\\n      It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\" r\\\"\\\"\\\"\\n    Args:\\n        kernel_size (int or tuple): the size of the sliding blocks\\n        dilation (int or tuple, optional): a parameter that controls the\\n                                           stride of elements within the\\n                                           neighborhood. Default: 1\\n        padding (int or tuple, optional): implicit zero padding to be added on\\n                                          both sides of input. Default: 0\\n        stride (int or tuple, optional): the stride of the sliding blocks in the input\\n                                         spatial dimensions. Default: 1\\n\\n    * If :attr:`kernel_size`, :attr:`dilation`, :attr:`padding` or\\n      :attr:`stride` is an int or a tuple of length 1, their values will be\\n      replicated across all spatial dimensions.\\n\\n    * For the case of two input spatial dimensions this operation is sometimes\\n      called ``im2col``.\\n\\n    .. note::\\n        :class:`~torch.nn.Fold` calculates each combined value in the resulting\\n        large tensor by summing all values from all containing blocks.\\n        :class:`~torch.nn.Unfold` extracts the values in the local blocks by\\n        copying from the large tensor. So, if the blocks overlap, they are not\\n        inverses of each other.\\n\\n        In general, folding and unfolding operations are related as\\n        follows. Consider :class:`~torch.nn.Fold` and\\n        :class:`~torch.nn.Unfold` instances created with the same\\n        parameters:\\n\\n        >>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)\\n        >>> fold = nn.Fold(output_size=..., **fold_params)\\n        >>> unfold = nn.Unfold(**fold_params)\\n\\n        Then for any (supported) ``input`` tensor the following\\n        equality holds:\\n\\n        ::\\n\\n            fold(unfold(input)) == divisor * input\\n\\n        where ``divisor`` is a tensor that depends only on the shape\\n        and dtype of the ``input``:\\n\\n        >>> # xdoctest: +SKIP\\n        >>> input_ones = torch.ones(input.shape, dtype=input.dtype)\\n        >>> divisor = fold(unfold(input_ones))\\n\\n        When the ``divisor`` tensor contains no zero elements, then\\n        ``fold`` and ``unfold`` operations are inverses of each\\n        other (up to constant divisor).\\n\\n    .. warning::\\n        Currently, only 4-D input tensors (batched image-like tensors) are\\n        supported.\\n\\n    Shape:\\n        - Input: :math:`(N, C, *)`\\n        - Output: :math:`(N, C \\\\times \\\\prod(\\\\text{kernel\\\\_size}), L)` as described above\\n\\n    Examples::\\n\\n        >>> unfold = nn.Unfold(kernel_size=(2, 3))\\n        >>> input = torch.randn(2, 5, 3, 4)\\n        >>> output = unfold(input)\\n        >>> # each patch contains 30 values (2x3=6 vectors, each of 5 channels)\\n        >>> # 4 blocks (2x3 kernels) in total in the 3x4 input\\n        >>> output.size()\\n        torch.Size([2, 30, 4])\\n\\n        >>> # xdoctest: +IGNORE_WANT\\n        >>> # Convolution is equivalent with Unfold + Matrix Multiplication + Fold (or view to output shape)\\n        >>> inp = torch.randn(1, 3, 10, 12)\\n        >>> w = torch.randn(2, 3, 4, 5)\\n        >>> inp_unf = torch.nn.functional.unfold(inp, (4, 5))\\n        >>> out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)\\n        >>> out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))\\n        >>> # or equivalently (and avoiding a copy),\\n        >>> # out = out_unf.view(1, 2, 7, 8)\\n        >>> (torch.nn.functional.conv2d(inp, w) - out).abs().max()\\n        tensor(1.9073e-06)\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"kernel_size\\\", \\\"dilation\\\", \\\"padding\\\", \\\"stride\\\"]\\n    kernel_size: _size_any_t\\n    dilation: _size_any_t\\n    padding: _size_any_t\\n    stride: _size_any_t\\n\\n    def __init__(\\n        self,\\n        kernel_size: _size_any_t,\\n        dilation: _size_any_t = 1,\\n        padding: _size_any_t = 0,\\n        stride: _size_any_t = 1,\\n    ) -> None:\\n        super().__init__()\\n        self.kernel_size = kernel_size\\n        self.dilation = dilation\\n        self.padding = padding\\n        self.stride = stride\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.unfold(\\n            input, self.kernel_size, self.dilation, self.padding, self.stride\\n        )\\n\\n    def extra_repr(self) -> str:\\n        return (\\n            \\\"kernel_size={kernel_size}, dilation={dilation}, padding={padding},\\\"\\n            \\\" stride={stride}\\\".format(**self.__dict__)\\n        )\\n\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"PixelShuffle\\\", \\\"PixelUnshuffle\\\"]\\n\\n\\nclass PixelShuffle(Module):\\n    r\\\"\\\"\\\"Rearrange elements in a tensor according to an upscaling factor.\\n\\n    Rearranges elements in a tensor of shape :math:`(*, C \\\\times r^2, H, W)`\\n    to a tensor of shape :math:`(*, C, H \\\\times r, W \\\\times r)`, where r is an upscale factor.\\n\\n    This is useful for implementing efficient sub-pixel convolution\\n    with a stride of :math:`1/r`.\\n\\n    See the paper:\\n    `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_\\n    by Shi et al. (2016) for more details.\\n\\n    Args:\\n        upscale_factor (int): factor to increase spatial resolution by\\n\\n    Shape:\\n        - Input: :math:`(*, C_{in}, H_{in}, W_{in})`, where * is zero or more batch dimensions\\n        - Output: :math:`(*, C_{out}, H_{out}, W_{out})`, where\\n\\n    .. math::\\n        C_{out} = C_{in} \\\\div \\\\text{upscale\\\\_factor}^2\\n\\n    .. math::\\n        H_{out} = H_{in} \\\\times \\\\text{upscale\\\\_factor}\\n\\n    .. math::\\n        W_{out} = W_{in} \\\\times \\\\text{upscale\\\\_factor}\\n\\n    Examples::\\n\\n        >>> pixel_shuffle = nn.PixelShuffle(3)\\n        >>> input = torch.randn(1, 9, 4, 4)\\n        >>> output = pixel_shuffle(input)\\n        >>> print(output.size())\\n        torch.Size([1, 1, 12, 12])\\n\\n    .. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:\\n        https://arxiv.org/abs/1609.05158\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"upscale_factor\\\"]\\n    upscale_factor: int\\n\\n    def __init__(self, upscale_factor: int) -> None:\\n        super().__init__()\\n        self.upscale_factor = upscale_factor\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.pixel_shuffle(input, self.upscale_factor)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"upscale_factor={self.upscale_factor}\\\"\\n\\n\\nclass PixelUnshuffle(Module):\\n    r\\\"\\\"\\\"Reverse the PixelShuffle operation.\\n\\n    Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements\\n    in a tensor of shape :math:`(*, C, H \\\\times r, W \\\\times r)` to a tensor of shape\\n    :math:`(*, C \\\\times r^2, H, W)`, where r is a downscale factor.\\n\\n    See the paper:\\n    `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_\\n    by Shi et al. (2016) for more details.\\n\\n    Args:\\n        downscale_factor (int): factor to decrease spatial resolution by\\n\\n    Shape:\\n        - Input: :math:`(*, C_{in}, H_{in}, W_{in})`, where * is zero or more batch dimensions\\n        - Output: :math:`(*, C_{out}, H_{out}, W_{out})`, where\\n\\n    .. math::\\n        C_{out} = C_{in} \\\\times \\\\text{downscale\\\\_factor}^2\\n\\n    .. math::\\n        H_{out} = H_{in} \\\\div \\\\text{downscale\\\\_factor}\\n\\n    .. math::\\n        W_{out} = W_{in} \\\\div \\\\text{downscale\\\\_factor}\\n\\n    Examples::\\n\\n        >>> pixel_unshuffle = nn.PixelUnshuffle(3)\\n        >>> input = torch.randn(1, 1, 12, 12)\\n        >>> output = pixel_unshuffle(input)\\n        >>> print(output.size())\\n        torch.Size([1, 9, 4, 4])\\n\\n    .. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:\\n        https://arxiv.org/abs/1609.05158\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"downscale_factor\\\"]\\n    downscale_factor: int\\n\\n    def __init__(self, downscale_factor: int) -> None:\\n        super().__init__()\\n        self.downscale_factor = downscale_factor\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.pixel_unshuffle(input, self.downscale_factor)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"downscale_factor={self.downscale_factor}\\\"\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Callable, Optional\\nfrom typing_extensions import deprecated\\n\\nfrom torch import Tensor\\nfrom torch.nn import _reduction as _Reduction, functional as F\\n\\nfrom .distance import PairwiseDistance\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"L1Loss\\\",\\n    \\\"NLLLoss\\\",\\n    \\\"NLLLoss2d\\\",\\n    \\\"PoissonNLLLoss\\\",\\n    \\\"GaussianNLLLoss\\\",\\n    \\\"KLDivLoss\\\",\\n    \\\"MSELoss\\\",\\n    \\\"BCELoss\\\",\\n    \\\"BCEWithLogitsLoss\\\",\\n    \\\"HingeEmbeddingLoss\\\",\\n    \\\"MultiLabelMarginLoss\\\",\\n    \\\"SmoothL1Loss\\\",\\n    \\\"HuberLoss\\\",\\n    \\\"SoftMarginLoss\\\",\\n    \\\"CrossEntropyLoss\\\",\\n    \\\"MultiLabelSoftMarginLoss\\\",\\n    \\\"CosineEmbeddingLoss\\\",\\n    \\\"MarginRankingLoss\\\",\\n    \\\"MultiMarginLoss\\\",\\n    \\\"TripletMarginLoss\\\",\\n    \\\"TripletMarginWithDistanceLoss\\\",\\n    \\\"CTCLoss\\\",\\n]\\n\\n\\nclass _Loss(Module):\\n    reduction: str\\n\\n    def __init__(self, size_average=None, reduce=None, reduction: str = \\\"mean\\\") -> None:\\n        super().__init__()\\n        if size_average is not None or reduce is not None:\\n            self.reduction: str = _Reduction.legacy_get_string(size_average, reduce)\\n        else:\\n            self.reduction = reduction\\n\\n\\nclass _WeightedLoss(_Loss):\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.register_buffer(\\\"weight\\\", weight)\\n        self.weight: Optional[Tensor]\\n\\n\\nclass L1Loss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the mean absolute error (MAE) between each element in\\n    the input :math:`x` and target :math:`y`.\\n\\n    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n        l_n = \\\\left| x_n - y_n \\\\right|,\\n\\n    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``\\n    (default ``'mean'``), then:\\n\\n    .. math::\\n        \\\\ell(x, y) =\\n        \\\\begin{cases}\\n            \\\\operatorname{mean}(L), & \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  & \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total\\n    of :math:`N` elements each.\\n\\n    The sum operation still operates over all the elements, and divides by :math:`N`.\\n\\n    The division by :math:`N` can be avoided if one sets ``reduction = 'sum'``.\\n\\n    Supports real-valued and complex-valued inputs.\\n\\n    Args:\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then\\n          :math:`(*)`, same shape as the input.\\n\\n    Examples::\\n\\n        >>> loss = nn.L1Loss()\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.randn(3, 5)\\n        >>> output = loss(input, target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(self, size_average=None, reduce=None, reduction: str = \\\"mean\\\") -> None:\\n        super().__init__(size_average, reduce, reduction)\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.l1_loss(input, target, reduction=self.reduction)\\n\\n\\nclass NLLLoss(_WeightedLoss):\\n    r\\\"\\\"\\\"The negative log likelihood loss. It is useful to train a classification\\n    problem with `C` classes.\\n\\n    If provided, the optional argument :attr:`weight` should be a 1D Tensor assigning\\n    weight to each of the classes. This is particularly useful when you have an\\n    unbalanced training set.\\n\\n    The `input` given through a forward call is expected to contain\\n    log-probabilities of each class. `input` has to be a Tensor of size either\\n    :math:`(minibatch, C)` or :math:`(minibatch, C, d_1, d_2, ..., d_K)`\\n    with :math:`K \\\\geq 1` for the `K`-dimensional case. The latter is useful for\\n    higher dimension inputs, such as computing NLL loss per-pixel for 2D images.\\n\\n    Obtaining log-probabilities in a neural network is easily achieved by\\n    adding a  `LogSoftmax`  layer in the last layer of your network.\\n    You may use `CrossEntropyLoss` instead, if you prefer not to add an extra\\n    layer.\\n\\n    The `target` that this loss expects should be a class index in the range :math:`[0, C-1]`\\n    where `C = number of classes`; if `ignore_index` is specified, this loss also accepts\\n    this class index (this index may not necessarily be in the class range).\\n\\n    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n        l_n = - w_{y_n} x_{n,y_n}, \\\\quad\\n        w_{c} = \\\\text{weight}[c] \\\\cdot \\\\mathbb{1}\\\\{c \\\\not= \\\\text{ignore\\\\_index}\\\\},\\n\\n    where :math:`x` is the input, :math:`y` is the target, :math:`w` is the weight, and\\n    :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``\\n    (default ``'mean'``), then\\n\\n    .. math::\\n        \\\\ell(x, y) = \\\\begin{cases}\\n            \\\\sum_{n=1}^N \\\\frac{1}{\\\\sum_{n=1}^N w_{y_n}} l_n, &\\n            \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\sum_{n=1}^N l_n,  &\\n            \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    Args:\\n        weight (Tensor, optional): a manual rescaling weight given to each\\n            class. If given, it has to be a Tensor of size `C`. Otherwise, it is\\n            treated as if having all ones.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``None``\\n        ignore_index (int, optional): Specifies a target value that is ignored\\n            and does not contribute to the input gradient. When\\n            :attr:`size_average` is ``True``, the loss is averaged over\\n            non-ignored targets.\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``None``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will\\n            be applied, ``'mean'``: the weighted mean of the output is taken,\\n            ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in\\n            the meantime, specifying either of those two args will override\\n            :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape::\\n        - Input: :math:`(N, C)` or :math:`(C)`, where `C = number of classes`, `N = batch size`, or\\n          :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1`\\n          in the case of `K`-dimensional loss.\\n        - Target: :math:`(N)` or :math:`()`, where each value is\\n          :math:`0 \\\\leq \\\\text{targets}[i] \\\\leq C-1`, or\\n          :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1` in the case of\\n          K-dimensional loss.\\n        - Output: If :attr:`reduction` is ``'none'``, shape :math:`(N)` or\\n          :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1` in the case of K-dimensional loss.\\n          Otherwise, scalar.\\n\\n    Examples::\\n\\n        >>> log_softmax = nn.LogSoftmax(dim=1)\\n        >>> loss_fn = nn.NLLLoss()\\n        >>> # input to NLLLoss is of size N x C = 3 x 5\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> # each element in target must have 0 <= value < C\\n        >>> target = torch.tensor([1, 0, 4])\\n        >>> loss = loss_fn(log_softmax(input), target)\\n        >>> loss.backward()\\n        >>>\\n        >>>\\n        >>> # 2D loss example (used, for example, with image inputs)\\n        >>> N, C = 5, 4\\n        >>> loss_fn = nn.NLLLoss()\\n        >>> data = torch.randn(N, 16, 10, 10)\\n        >>> conv = nn.Conv2d(16, C, (3, 3))\\n        >>> log_softmax = nn.LogSoftmax(dim=1)\\n        >>> # output of conv forward is of shape [N, C, 8, 8]\\n        >>> output = log_softmax(conv(data))\\n        >>> # each element in target must have 0 <= value < C\\n        >>> target = torch.empty(N, 8, 8, dtype=torch.long).random_(0, C)\\n        >>> # input to NLLLoss is of size N x C x height (8) x width (8)\\n        >>> loss = loss_fn(output, target)\\n        >>> loss.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"ignore_index\\\", \\\"reduction\\\"]\\n    ignore_index: int\\n\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        ignore_index: int = -100,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(weight, size_average, reduce, reduction)\\n        self.ignore_index = ignore_index\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.nll_loss(\\n            input,\\n            target,\\n            weight=self.weight,\\n            ignore_index=self.ignore_index,\\n            reduction=self.reduction,\\n        )\\n\\n\\n@deprecated(\\n    \\\"`NLLLoss2d` has been deprecated. \\\"\\n    \\\"Please use `NLLLoss` instead as a drop-in replacement and see \\\"\\n    \\\"https://pytorch.org/docs/main/nn.html#torch.nn.NLLLoss for more details.\\\",\\n    category=FutureWarning,\\n)\\nclass NLLLoss2d(NLLLoss):\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        ignore_index: int = -100,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(weight, size_average, ignore_index, reduce, reduction)\\n\\n\\nclass PoissonNLLLoss(_Loss):\\n    r\\\"\\\"\\\"Negative log likelihood loss with Poisson distribution of target.\\n\\n    The loss can be described as:\\n\\n    .. math::\\n        \\\\text{target} \\\\sim \\\\mathrm{Poisson}(\\\\text{input})\\n\\n        \\\\text{loss}(\\\\text{input}, \\\\text{target}) = \\\\text{input} - \\\\text{target} * \\\\log(\\\\text{input})\\n                                    + \\\\log(\\\\text{target!})\\n\\n    The last term can be omitted or approximated with Stirling formula. The\\n    approximation is used for target values more than 1. For targets less or\\n    equal to 1 zeros are added to the loss.\\n\\n    Args:\\n        log_input (bool, optional): if ``True`` the loss is computed as\\n            :math:`\\\\exp(\\\\text{input}) - \\\\text{target}*\\\\text{input}`, if ``False`` the loss is\\n            :math:`\\\\text{input} - \\\\text{target}*\\\\log(\\\\text{input}+\\\\text{eps})`.\\n        full (bool, optional): whether to compute full loss, i. e. to add the\\n            Stirling approximation term\\n\\n            .. math::\\n                \\\\text{target}*\\\\log(\\\\text{target}) - \\\\text{target} + 0.5 * \\\\log(2\\\\pi\\\\text{target}).\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        eps (float, optional): Small value to avoid evaluation of :math:`\\\\log(0)` when\\n            :attr:`log_input = False`. Default: 1e-8\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Examples::\\n\\n        >>> loss = nn.PoissonNLLLoss()\\n        >>> log_input = torch.randn(5, 2, requires_grad=True)\\n        >>> target = torch.randn(5, 2)\\n        >>> output = loss(log_input, target)\\n        >>> output.backward()\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar by default. If :attr:`reduction` is ``'none'``, then :math:`(*)`,\\n          the same shape as the input.\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"log_input\\\", \\\"full\\\", \\\"eps\\\", \\\"reduction\\\"]\\n    log_input: bool\\n    full: bool\\n    eps: float\\n\\n    def __init__(\\n        self,\\n        log_input: bool = True,\\n        full: bool = False,\\n        size_average=None,\\n        eps: float = 1e-8,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.log_input = log_input\\n        self.full = full\\n        self.eps = eps\\n\\n    def forward(self, log_input: Tensor, target: Tensor) -> Tensor:\\n        return F.poisson_nll_loss(\\n            log_input,\\n            target,\\n            log_input=self.log_input,\\n            full=self.full,\\n            eps=self.eps,\\n            reduction=self.reduction,\\n        )\\n\\n\\nclass GaussianNLLLoss(_Loss):\\n    r\\\"\\\"\\\"Gaussian negative log likelihood loss.\\n\\n    The targets are treated as samples from Gaussian distributions with\\n    expectations and variances predicted by the neural network. For a\\n    ``target`` tensor modelled as having Gaussian distribution with a tensor\\n    of expectations ``input`` and a tensor of positive variances ``var`` the loss is:\\n\\n    .. math::\\n        \\\\text{loss} = \\\\frac{1}{2}\\\\left(\\\\log\\\\left(\\\\text{max}\\\\left(\\\\text{var},\\n        \\\\ \\\\text{eps}\\\\right)\\\\right) + \\\\frac{\\\\left(\\\\text{input} - \\\\text{target}\\\\right)^2}\\n        {\\\\text{max}\\\\left(\\\\text{var}, \\\\ \\\\text{eps}\\\\right)}\\\\right) + \\\\text{const.}\\n\\n    where :attr:`eps` is used for stability. By default, the constant term of\\n    the loss function is omitted unless :attr:`full` is ``True``. If ``var`` is not the same\\n    size as ``input`` (due to a homoscedastic assumption), it must either have a final dimension\\n    of 1 or have one fewer dimension (with all other sizes being the same) for correct broadcasting.\\n\\n    Args:\\n        full (bool, optional): include the constant term in the loss\\n            calculation. Default: ``False``.\\n        eps (float, optional): value used to clamp ``var`` (see note below), for\\n            stability. Default: 1e-6.\\n        reduction (str, optional): specifies the reduction to apply to the\\n            output:``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction\\n            will be applied, ``'mean'``: the output is the average of all batch\\n            member losses, ``'sum'``: the output is the sum of all batch member\\n            losses. Default: ``'mean'``.\\n\\n    Shape:\\n        - Input: :math:`(N, *)` or :math:`(*)` where :math:`*` means any number of additional\\n          dimensions\\n        - Target: :math:`(N, *)` or :math:`(*)`, same shape as the input, or same shape as the input\\n          but with one dimension equal to 1 (to allow for broadcasting)\\n        - Var: :math:`(N, *)` or :math:`(*)`, same shape as the input, or same shape as the input but\\n          with one dimension equal to 1, or same shape as the input but with one fewer\\n          dimension (to allow for broadcasting)\\n        - Output: scalar if :attr:`reduction` is ``'mean'`` (default) or\\n          ``'sum'``. If :attr:`reduction` is ``'none'``, then :math:`(N, *)`, same\\n          shape as the input\\n\\n    Examples::\\n        >>> loss = nn.GaussianNLLLoss()\\n        >>> input = torch.randn(5, 2, requires_grad=True)\\n        >>> target = torch.randn(5, 2)\\n        >>> var = torch.ones(5, 2, requires_grad=True)  # heteroscedastic\\n        >>> output = loss(input, target, var)\\n        >>> output.backward()\\n\\n        >>> loss = nn.GaussianNLLLoss()\\n        >>> input = torch.randn(5, 2, requires_grad=True)\\n        >>> target = torch.randn(5, 2)\\n        >>> var = torch.ones(5, 1, requires_grad=True)  # homoscedastic\\n        >>> output = loss(input, target, var)\\n        >>> output.backward()\\n\\n    Note:\\n        The clamping of ``var`` is ignored with respect to autograd, and so the\\n        gradients are unaffected by it.\\n\\n    Reference:\\n        Nix, D. A. and Weigend, A. S., \\\"Estimating the mean and variance of the\\n        target probability distribution\\\", Proceedings of 1994 IEEE International\\n        Conference on Neural Networks (ICNN'94), Orlando, FL, USA, 1994, pp. 55-60\\n        vol.1, doi: 10.1109/ICNN.1994.374138.\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"full\\\", \\\"eps\\\", \\\"reduction\\\"]\\n    full: bool\\n    eps: float\\n\\n    def __init__(\\n        self, *, full: bool = False, eps: float = 1e-6, reduction: str = \\\"mean\\\"\\n    ) -> None:\\n        super().__init__(None, None, reduction)\\n        self.full = full\\n        self.eps = eps\\n\\n    def forward(self, input: Tensor, target: Tensor, var: Tensor) -> Tensor:\\n        return F.gaussian_nll_loss(\\n            input, target, var, full=self.full, eps=self.eps, reduction=self.reduction\\n        )\\n\\n\\nclass KLDivLoss(_Loss):\\n    r\\\"\\\"\\\"The Kullback-Leibler divergence loss.\\n\\n    For tensors of the same shape :math:`y_{\\\\text{pred}},\\\\ y_{\\\\text{true}}`,\\n    where :math:`y_{\\\\text{pred}}` is the :attr:`input` and :math:`y_{\\\\text{true}}` is the\\n    :attr:`target`, we define the **pointwise KL-divergence** as\\n\\n    .. math::\\n\\n        L(y_{\\\\text{pred}},\\\\ y_{\\\\text{true}})\\n            = y_{\\\\text{true}} \\\\cdot \\\\log \\\\frac{y_{\\\\text{true}}}{y_{\\\\text{pred}}}\\n            = y_{\\\\text{true}} \\\\cdot (\\\\log y_{\\\\text{true}} - \\\\log y_{\\\\text{pred}})\\n\\n    To avoid underflow issues when computing this quantity, this loss expects the argument\\n    :attr:`input` in the log-space. The argument :attr:`target` may also be provided in the\\n    log-space if :attr:`log_target`\\\\ `= True`.\\n\\n    To summarise, this function is roughly equivalent to computing\\n\\n    .. code-block:: python\\n\\n        if not log_target: # default\\n            loss_pointwise = target * (target.log() - input)\\n        else:\\n            loss_pointwise = target.exp() * (target - input)\\n\\n    and then reducing this result depending on the argument :attr:`reduction` as\\n\\n    .. code-block:: python\\n\\n        if reduction == \\\"mean\\\":  # default\\n            loss = loss_pointwise.mean()\\n        elif reduction == \\\"batchmean\\\":  # mathematically correct\\n            loss = loss_pointwise.sum() / input.size(0)\\n        elif reduction == \\\"sum\\\":\\n            loss = loss_pointwise.sum()\\n        else:  # reduction == \\\"none\\\"\\n            loss = loss_pointwise\\n\\n    .. note::\\n        As all the other losses in PyTorch, this function expects the first argument,\\n        :attr:`input`, to be the output of the model (e.g. the neural network)\\n        and the second, :attr:`target`, to be the observations in the dataset.\\n        This differs from the standard mathematical notation :math:`KL(P\\\\ ||\\\\ Q)` where\\n        :math:`P` denotes the distribution of the observations and :math:`Q` denotes the model.\\n\\n    .. warning::\\n        :attr:`reduction`\\\\ `= \\\"mean\\\"` doesn't return the true KL divergence value, please use\\n        :attr:`reduction`\\\\ `= \\\"batchmean\\\"` which aligns with the mathematical definition.\\n\\n    Args:\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to `False`, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is `False`. Default: `True`\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is `False`, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: `True`\\n        reduction (str, optional): Specifies the reduction to apply to the output. Default: `\\\"mean\\\"`\\n        log_target (bool, optional): Specifies whether `target` is the log space. Default: `False`\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar by default. If :attr:`reduction` is `'none'`, then :math:`(*)`,\\n          same shape as the input.\\n\\n    Examples::\\n        >>> kl_loss = nn.KLDivLoss(reduction=\\\"batchmean\\\")\\n        >>> # input should be a distribution in the log space\\n        >>> input = F.log_softmax(torch.randn(3, 5, requires_grad=True), dim=1)\\n        >>> # Sample a batch of distributions. Usually this would come from the dataset\\n        >>> target = F.softmax(torch.rand(3, 5), dim=1)\\n        >>> output = kl_loss(input, target)\\n\\n        >>> kl_loss = nn.KLDivLoss(reduction=\\\"batchmean\\\", log_target=True)\\n        >>> log_target = F.log_softmax(torch.rand(3, 5), dim=1)\\n        >>> output = kl_loss(input, log_target)\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(\\n        self,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n        log_target: bool = False,\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.log_target = log_target\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.kl_div(\\n            input, target, reduction=self.reduction, log_target=self.log_target\\n        )\\n\\n\\nclass MSELoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the mean squared error (squared L2 norm) between\\n    each element in the input :math:`x` and target :math:`y`.\\n\\n    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n        l_n = \\\\left( x_n - y_n \\\\right)^2,\\n\\n    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``\\n    (default ``'mean'``), then:\\n\\n    .. math::\\n        \\\\ell(x, y) =\\n        \\\\begin{cases}\\n            \\\\operatorname{mean}(L), &  \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  &  \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total\\n    of :math:`N` elements each.\\n\\n    The mean operation still operates over all the elements, and divides by :math:`N`.\\n\\n    The division by :math:`N` can be avoided if one sets ``reduction = 'sum'``.\\n\\n    Args:\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n\\n    Examples::\\n\\n        >>> loss = nn.MSELoss()\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.randn(3, 5)\\n        >>> output = loss(input, target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(self, size_average=None, reduce=None, reduction: str = \\\"mean\\\") -> None:\\n        super().__init__(size_average, reduce, reduction)\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.mse_loss(input, target, reduction=self.reduction)\\n\\n\\nclass BCELoss(_WeightedLoss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the Binary Cross Entropy between the target and\\n    the input probabilities:\\n\\n    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n        l_n = - w_n \\\\left[ y_n \\\\cdot \\\\log x_n + (1 - y_n) \\\\cdot \\\\log (1 - x_n) \\\\right],\\n\\n    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``\\n    (default ``'mean'``), then\\n\\n    .. math::\\n        \\\\ell(x, y) = \\\\begin{cases}\\n            \\\\operatorname{mean}(L), & \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  & \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    This is used for measuring the error of a reconstruction in for example\\n    an auto-encoder. Note that the targets :math:`y` should be numbers\\n    between 0 and 1.\\n\\n    Notice that if :math:`x_n` is either 0 or 1, one of the log terms would be\\n    mathematically undefined in the above loss equation. PyTorch chooses to set\\n    :math:`\\\\log (0) = -\\\\infty`, since :math:`\\\\lim_{x\\\\to 0} \\\\log (x) = -\\\\infty`.\\n    However, an infinite term in the loss equation is not desirable for several reasons.\\n\\n    For one, if either :math:`y_n = 0` or :math:`(1 - y_n) = 0`, then we would be\\n    multiplying 0 with infinity. Secondly, if we have an infinite loss value, then\\n    we would also have an infinite term in our gradient, since\\n    :math:`\\\\lim_{x\\\\to 0} \\\\frac{d}{dx} \\\\log (x) = \\\\infty`.\\n    This would make BCELoss's backward method nonlinear with respect to :math:`x_n`,\\n    and using it for things like linear regression would not be straight-forward.\\n\\n    Our solution is that BCELoss clamps its log function outputs to be greater than\\n    or equal to -100. This way, we can always have a finite loss value and a linear\\n    backward method.\\n\\n\\n    Args:\\n        weight (Tensor, optional): a manual rescaling weight given to the loss\\n            of each batch element. If given, has to be a Tensor of size `nbatch`.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same\\n          shape as input.\\n\\n    Examples::\\n\\n        >>> m = nn.Sigmoid()\\n        >>> loss = nn.BCELoss()\\n        >>> input = torch.randn(3, 2, requires_grad=True)\\n        >>> target = torch.rand(3, 2, requires_grad=False)\\n        >>> output = loss(m(input), target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(weight, size_average, reduce, reduction)\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.binary_cross_entropy(\\n            input, target, weight=self.weight, reduction=self.reduction\\n        )\\n\\n\\nclass BCEWithLogitsLoss(_Loss):\\n    r\\\"\\\"\\\"This loss combines a `Sigmoid` layer and the `BCELoss` in one single\\n    class. This version is more numerically stable than using a plain `Sigmoid`\\n    followed by a `BCELoss` as, by combining the operations into one layer,\\n    we take advantage of the log-sum-exp trick for numerical stability.\\n\\n    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n        l_n = - w_n \\\\left[ y_n \\\\cdot \\\\log \\\\sigma(x_n)\\n        + (1 - y_n) \\\\cdot \\\\log (1 - \\\\sigma(x_n)) \\\\right],\\n\\n    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``\\n    (default ``'mean'``), then\\n\\n    .. math::\\n        \\\\ell(x, y) = \\\\begin{cases}\\n            \\\\operatorname{mean}(L), & \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  & \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    This is used for measuring the error of a reconstruction in for example\\n    an auto-encoder. Note that the targets `t[i]` should be numbers\\n    between 0 and 1.\\n\\n    It's possible to trade off recall and precision by adding weights to positive examples.\\n    In the case of multi-label classification the loss can be described as:\\n\\n    .. math::\\n        \\\\ell_c(x, y) = L_c = \\\\{l_{1,c},\\\\dots,l_{N,c}\\\\}^\\\\top, \\\\quad\\n        l_{n,c} = - w_{n,c} \\\\left[ p_c y_{n,c} \\\\cdot \\\\log \\\\sigma(x_{n,c})\\n        + (1 - y_{n,c}) \\\\cdot \\\\log (1 - \\\\sigma(x_{n,c})) \\\\right],\\n\\n    where :math:`c` is the class number (:math:`c > 1` for multi-label binary classification,\\n    :math:`c = 1` for single-label binary classification),\\n    :math:`n` is the number of the sample in the batch and\\n    :math:`p_c` is the weight of the positive answer for the class :math:`c`.\\n\\n    :math:`p_c > 1` increases the recall, :math:`p_c < 1` increases the precision.\\n\\n    For example, if a dataset contains 100 positive and 300 negative examples of a single class,\\n    then ``pos_weight`` for the class should be equal to :math:`\\\\frac{300}{100}=3`.\\n    The loss would act as if the dataset contains :math:`3\\\\times 100=300` positive examples.\\n\\n    Examples::\\n\\n        >>> target = torch.ones([10, 64], dtype=torch.float32)  # 64 classes, batch size = 10\\n        >>> output = torch.full([10, 64], 1.5)  # A prediction (logit)\\n        >>> pos_weight = torch.ones([64])  # All weights are equal to 1\\n        >>> criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)\\n        >>> criterion(output, target)  # -log(sigmoid(1.5))\\n        tensor(0.20...)\\n\\n    In the above example, the ``pos_weight`` tensor's elements correspond to the 64 distinct classes\\n    in a multi-label binary classification scenario. Each element in ``pos_weight`` is designed to adjust the\\n    loss function based on the imbalance between negative and positive samples for the respective class.\\n    This approach is useful in datasets with varying levels of class imbalance, ensuring that the loss\\n    calculation accurately accounts for the distribution in each class.\\n\\n    Args:\\n        weight (Tensor, optional): a manual rescaling weight given to the loss\\n            of each batch element. If given, has to be a Tensor of size `nbatch`.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n        pos_weight (Tensor, optional): a weight of positive examples to be broadcasted with target.\\n            Must be a tensor with equal size along the class dimension to the number of classes.\\n            Pay close attention to PyTorch's broadcasting semantics in order to achieve the desired\\n            operations. For a target of size [B, C, H, W] (where B is batch size) pos_weight of\\n            size [B, C, H, W] will apply different pos_weights to each element of the batch or\\n            [C, H, W] the same pos_weights across the batch. To apply the same positive weight\\n            along all spacial dimensions for a 2D multi-class target [C, H, W] use: [C, 1, 1].\\n            Default: ``None``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same\\n          shape as input.\\n\\n     Examples::\\n\\n        >>> loss = nn.BCEWithLogitsLoss()\\n        >>> input = torch.randn(3, requires_grad=True)\\n        >>> target = torch.empty(3).random_(2)\\n        >>> output = loss(input, target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n        pos_weight: Optional[Tensor] = None,\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.register_buffer(\\\"weight\\\", weight)\\n        self.register_buffer(\\\"pos_weight\\\", pos_weight)\\n        self.weight: Optional[Tensor]\\n        self.pos_weight: Optional[Tensor]\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.binary_cross_entropy_with_logits(\\n            input,\\n            target,\\n            self.weight,\\n            pos_weight=self.pos_weight,\\n            reduction=self.reduction,\\n        )\\n\\n\\nclass HingeEmbeddingLoss(_Loss):\\n    r\\\"\\\"\\\"Measures the loss given an input tensor :math:`x` and a labels tensor :math:`y`\\n    (containing 1 or -1).\\n    This is usually used for measuring whether two inputs are similar or\\n    dissimilar, e.g. using the L1 pairwise distance as :math:`x`, and is typically\\n    used for learning nonlinear embeddings or semi-supervised learning.\\n\\n    The loss function for :math:`n`-th sample in the mini-batch is\\n\\n    .. math::\\n        l_n = \\\\begin{cases}\\n            x_n, & \\\\text{if}\\\\; y_n = 1,\\\\\\\\\\n            \\\\max \\\\{0, margin - x_n\\\\}, & \\\\text{if}\\\\; y_n = -1,\\n        \\\\end{cases}\\n\\n    and the total loss functions is\\n\\n    .. math::\\n        \\\\ell(x, y) = \\\\begin{cases}\\n            \\\\operatorname{mean}(L), & \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  & \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    where :math:`L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top`.\\n\\n    Args:\\n        margin (float, optional): Has a default value of `1`.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(*)` where :math:`*` means, any number of dimensions. The sum operation\\n          operates over all the elements.\\n        - Target: :math:`(*)`, same shape as the input\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then same shape as the input\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"margin\\\", \\\"reduction\\\"]\\n    margin: float\\n\\n    def __init__(\\n        self,\\n        margin: float = 1.0,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.margin = margin\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.hinge_embedding_loss(\\n            input, target, margin=self.margin, reduction=self.reduction\\n        )\\n\\n\\nclass MultiLabelMarginLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that optimizes a multi-class multi-classification\\n    hinge loss (margin-based loss) between input :math:`x` (a 2D mini-batch `Tensor`)\\n    and output :math:`y` (which is a 2D `Tensor` of target class indices).\\n    For each sample in the mini-batch:\\n\\n    .. math::\\n        \\\\text{loss}(x, y) = \\\\sum_{ij}\\\\frac{\\\\max(0, 1 - (x[y[j]] - x[i]))}{\\\\text{x.size}(0)}\\n\\n    where :math:`x \\\\in \\\\left\\\\{0, \\\\; \\\\cdots , \\\\; \\\\text{x.size}(0) - 1\\\\right\\\\}`, \\\\\\n    :math:`y \\\\in \\\\left\\\\{0, \\\\; \\\\cdots , \\\\; \\\\text{y.size}(0) - 1\\\\right\\\\}`, \\\\\\n    :math:`0 \\\\leq y[j] \\\\leq \\\\text{x.size}(0)-1`, \\\\\\n    and :math:`i \\\\neq y[j]` for all :math:`i` and :math:`j`.\\n\\n    :math:`y` and :math:`x` must have the same size.\\n\\n    The criterion only considers a contiguous block of non-negative targets that\\n    starts at the front.\\n\\n    This allows for different samples to have variable amounts of target classes.\\n\\n    Args:\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(C)` or :math:`(N, C)` where `N` is the batch size and `C`\\n          is the number of classes.\\n        - Target: :math:`(C)` or :math:`(N, C)`, label targets padded by -1 ensuring same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(N)`.\\n\\n    Examples::\\n\\n        >>> loss = nn.MultiLabelMarginLoss()\\n        >>> x = torch.FloatTensor([[0.1, 0.2, 0.4, 0.8]])\\n        >>> # for target y, only consider labels 3 and 0, not after label -1\\n        >>> y = torch.LongTensor([[3, 0, -1, 1]])\\n        >>> # 0.25 * ((1-(0.1-0.2)) + (1-(0.1-0.4)) + (1-(0.8-0.2)) + (1-(0.8-0.4)))\\n        >>> loss(x, y)\\n        tensor(0.85...)\\n\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(self, size_average=None, reduce=None, reduction: str = \\\"mean\\\") -> None:\\n        super().__init__(size_average, reduce, reduction)\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.multilabel_margin_loss(input, target, reduction=self.reduction)\\n\\n\\nclass SmoothL1Loss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that uses a squared term if the absolute\\n    element-wise error falls below beta and an L1 term otherwise.\\n    It is less sensitive to outliers than :class:`torch.nn.MSELoss` and in some cases\\n    prevents exploding gradients (e.g. see the paper `Fast R-CNN`_ by Ross Girshick).\\n\\n    For a batch of size :math:`N`, the unreduced loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1, ..., l_N\\\\}^T\\n\\n    with\\n\\n    .. math::\\n        l_n = \\\\begin{cases}\\n        0.5 (x_n - y_n)^2 / beta, & \\\\text{if } |x_n - y_n| < beta \\\\\\\\\\n        |x_n - y_n| - 0.5 * beta, & \\\\text{otherwise }\\n        \\\\end{cases}\\n\\n    If `reduction` is not `none`, then:\\n\\n    .. math::\\n        \\\\ell(x, y) =\\n        \\\\begin{cases}\\n            \\\\operatorname{mean}(L), &  \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  &  \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    .. note::\\n        Smooth L1 loss can be seen as exactly :class:`L1Loss`, but with the :math:`|x - y| < beta`\\n        portion replaced with a quadratic function such that its slope is 1 at :math:`|x - y| = beta`.\\n        The quadratic segment smooths the L1 loss near :math:`|x - y| = 0`.\\n\\n    .. note::\\n        Smooth L1 loss is closely related to :class:`HuberLoss`, being\\n        equivalent to :math:`huber(x, y) / beta` (note that Smooth L1's beta hyper-parameter is\\n        also known as delta for Huber). This leads to the following differences:\\n\\n        * As beta -> 0, Smooth L1 loss converges to :class:`L1Loss`, while :class:`HuberLoss`\\n          converges to a constant 0 loss. When beta is 0, Smooth L1 loss is equivalent to L1 loss.\\n        * As beta -> :math:`+\\\\infty`, Smooth L1 loss converges to a constant 0 loss, while\\n          :class:`HuberLoss` converges to :class:`MSELoss`.\\n        * For Smooth L1 loss, as beta varies, the L1 segment of the loss has a constant slope of 1.\\n          For :class:`HuberLoss`, the slope of the L1 segment is beta.\\n\\n    .. _`Fast R-CNN`: https://arxiv.org/abs/1504.08083\\n\\n    Args:\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n        beta (float, optional): Specifies the threshold at which to change between L1 and L2 loss.\\n            The value must be non-negative. Default: 1.0\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same shape as the input.\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(\\n        self, size_average=None, reduce=None, reduction: str = \\\"mean\\\", beta: float = 1.0\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.beta = beta\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.smooth_l1_loss(input, target, reduction=self.reduction, beta=self.beta)\\n\\n\\nclass HuberLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that uses a squared term if the absolute\\n    element-wise error falls below delta and a delta-scaled L1 term otherwise.\\n    This loss combines advantages of both :class:`L1Loss` and :class:`MSELoss`; the\\n    delta-scaled L1 region makes the loss less sensitive to outliers than :class:`MSELoss`,\\n    while the L2 region provides smoothness over :class:`L1Loss` near 0. See\\n    `Huber loss <https://en.wikipedia.org/wiki/Huber_loss>`_ for more information.\\n\\n    For a batch of size :math:`N`, the unreduced loss can be described as:\\n\\n    .. math::\\n        \\\\ell(x, y) = L = \\\\{l_1, ..., l_N\\\\}^T\\n\\n    with\\n\\n    .. math::\\n        l_n = \\\\begin{cases}\\n        0.5 (x_n - y_n)^2, & \\\\text{if } |x_n - y_n| < delta \\\\\\\\\\n        delta * (|x_n - y_n| - 0.5 * delta), & \\\\text{otherwise }\\n        \\\\end{cases}\\n\\n    If `reduction` is not `none`, then:\\n\\n    .. math::\\n        \\\\ell(x, y) =\\n        \\\\begin{cases}\\n            \\\\operatorname{mean}(L), &  \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  &  \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    .. note::\\n        When delta is set to 1, this loss is equivalent to :class:`SmoothL1Loss`.\\n        In general, this loss differs from :class:`SmoothL1Loss` by a factor of delta (AKA beta\\n        in Smooth L1).\\n        See :class:`SmoothL1Loss` for additional discussion on the differences in behavior\\n        between the two losses.\\n\\n    Args:\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'``\\n        delta (float, optional): Specifies the threshold at which to change between delta-scaled L1 and L2 loss.\\n            The value must be positive.  Default: 1.0\\n\\n    Shape:\\n        - Input: :math:`(*)` where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same shape as the input.\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\", \\\"delta\\\"]\\n\\n    def __init__(self, reduction: str = \\\"mean\\\", delta: float = 1.0) -> None:\\n        super().__init__(reduction=reduction)\\n        self.delta = delta\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.huber_loss(input, target, reduction=self.reduction, delta=self.delta)\\n\\n\\nclass SoftMarginLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that optimizes a two-class classification\\n    logistic loss between input tensor :math:`x` and target tensor :math:`y`\\n    (containing 1 or -1).\\n\\n    .. math::\\n        \\\\text{loss}(x, y) = \\\\sum_i \\\\frac{\\\\log(1 + \\\\exp(-y[i]*x[i]))}{\\\\text{x.nelement}()}\\n\\n    Args:\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.\\n        - Target: :math:`(*)`, same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same\\n          shape as input.\\n\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(self, size_average=None, reduce=None, reduction: str = \\\"mean\\\") -> None:\\n        super().__init__(size_average, reduce, reduction)\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.soft_margin_loss(input, target, reduction=self.reduction)\\n\\n\\nclass CrossEntropyLoss(_WeightedLoss):\\n    r\\\"\\\"\\\"This criterion computes the cross entropy loss between input logits\\n    and target.\\n\\n    It is useful when training a classification problem with `C` classes.\\n    If provided, the optional argument :attr:`weight` should be a 1D `Tensor`\\n    assigning weight to each of the classes.\\n    This is particularly useful when you have an unbalanced training set.\\n\\n    The `input` is expected to contain the unnormalized logits for each class (which do `not` need\\n    to be positive or sum to 1, in general).\\n    `input` has to be a Tensor of size :math:`(C)` for unbatched input,\\n    :math:`(minibatch, C)` or :math:`(minibatch, C, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1` for the\\n    `K`-dimensional case. The last being useful for higher dimension inputs, such\\n    as computing cross entropy loss per-pixel for 2D images.\\n\\n    The `target` that this criterion expects should contain either:\\n\\n    - Class indices in the range :math:`[0, C)` where :math:`C` is the number of classes; if\\n      `ignore_index` is specified, this loss also accepts this class index (this index\\n      may not necessarily be in the class range). The unreduced (i.e. with :attr:`reduction`\\n      set to ``'none'``) loss for this case can be described as:\\n\\n      .. math::\\n          \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n          l_n = - w_{y_n} \\\\log \\\\frac{\\\\exp(x_{n,y_n})}{\\\\sum_{c=1}^C \\\\exp(x_{n,c})}\\n          \\\\cdot \\\\mathbb{1}\\\\{y_n \\\\not= \\\\text{ignore\\\\_index}\\\\}\\n\\n      where :math:`x` is the input, :math:`y` is the target, :math:`w` is the weight,\\n      :math:`C` is the number of classes, and :math:`N` spans the minibatch dimension as well as\\n      :math:`d_1, ..., d_k` for the `K`-dimensional case. If\\n      :attr:`reduction` is not ``'none'`` (default ``'mean'``), then\\n\\n      .. math::\\n          \\\\ell(x, y) = \\\\begin{cases}\\n              \\\\sum_{n=1}^N \\\\frac{1}{\\\\sum_{n=1}^N w_{y_n} \\\\cdot \\\\mathbb{1}\\\\{y_n \\\\not= \\\\text{ignore\\\\_index}\\\\}} l_n, &\\n               \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n                \\\\sum_{n=1}^N l_n,  &\\n                \\\\text{if reduction} = \\\\text{`sum'.}\\n            \\\\end{cases}\\n\\n      Note that this case is equivalent to applying :class:`~torch.nn.LogSoftmax`\\n      on an input, followed by :class:`~torch.nn.NLLLoss`.\\n\\n    - Probabilities for each class; useful when labels beyond a single class per minibatch item\\n      are required, such as for blended labels, label smoothing, etc. The unreduced (i.e. with\\n      :attr:`reduction` set to ``'none'``) loss for this case can be described as:\\n\\n      .. math::\\n          \\\\ell(x, y) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n          l_n = - \\\\sum_{c=1}^C w_c \\\\log \\\\frac{\\\\exp(x_{n,c})}{\\\\sum_{i=1}^C \\\\exp(x_{n,i})} y_{n,c}\\n\\n      where :math:`x` is the input, :math:`y` is the target, :math:`w` is the weight,\\n      :math:`C` is the number of classes, and :math:`N` spans the minibatch dimension as well as\\n      :math:`d_1, ..., d_k` for the `K`-dimensional case. If\\n      :attr:`reduction` is not ``'none'`` (default ``'mean'``), then\\n\\n      .. math::\\n          \\\\ell(x, y) = \\\\begin{cases}\\n              \\\\frac{\\\\sum_{n=1}^N l_n}{N}, &\\n               \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n                \\\\sum_{n=1}^N l_n,  &\\n                \\\\text{if reduction} = \\\\text{`sum'.}\\n            \\\\end{cases}\\n\\n    .. note::\\n        The performance of this criterion is generally better when `target` contains class\\n        indices, as this allows for optimized computation. Consider providing `target` as\\n        class probabilities only when a single class label per minibatch item is too restrictive.\\n\\n    Args:\\n        weight (Tensor, optional): a manual rescaling weight given to each class.\\n            If given, has to be a Tensor of size `C` and floating point dtype\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        ignore_index (int, optional): Specifies a target value that is ignored\\n            and does not contribute to the input gradient. When :attr:`size_average` is\\n            ``True``, the loss is averaged over non-ignored targets. Note that\\n            :attr:`ignore_index` is only applicable when the target contains class indices.\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will\\n            be applied, ``'mean'``: the weighted mean of the output is taken,\\n            ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in\\n            the meantime, specifying either of those two args will override\\n            :attr:`reduction`. Default: ``'mean'``\\n        label_smoothing (float, optional): A float in [0.0, 1.0]. Specifies the amount\\n            of smoothing when computing the loss, where 0.0 means no smoothing. The targets\\n            become a mixture of the original ground truth and a uniform distribution as described in\\n            `Rethinking the Inception Architecture for Computer Vision <https://arxiv.org/abs/1512.00567>`__. Default: :math:`0.0`.\\n\\n    Shape:\\n        - Input: Shape :math:`(C)`, :math:`(N, C)` or :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1`\\n          in the case of `K`-dimensional loss.\\n        - Target: If containing class indices, shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with\\n          :math:`K \\\\geq 1` in the case of K-dimensional loss where each value should be between :math:`[0, C)`.\\n          If containing class probabilities, same shape as the input and each value should be between :math:`[0, 1]`.\\n        - Output: If reduction is 'none', shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \\\\geq 1`\\n          in the case of K-dimensional loss, depending on the shape of the input. Otherwise, scalar.\\n\\n\\n        where:\\n\\n        .. math::\\n            \\\\begin{aligned}\\n                C ={} & \\\\text{number of classes} \\\\\\\\\\n                N ={} & \\\\text{batch size} \\\\\\\\\\n            \\\\end{aligned}\\n\\n    Examples::\\n\\n        >>> # Example of target with class indices\\n        >>> loss = nn.CrossEntropyLoss()\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.empty(3, dtype=torch.long).random_(5)\\n        >>> output = loss(input, target)\\n        >>> output.backward()\\n        >>>\\n        >>> # Example of target with class probabilities\\n        >>> input = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.randn(3, 5).softmax(dim=1)\\n        >>> output = loss(input, target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"ignore_index\\\", \\\"reduction\\\", \\\"label_smoothing\\\"]\\n    ignore_index: int\\n    label_smoothing: float\\n\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        ignore_index: int = -100,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n        label_smoothing: float = 0.0,\\n    ) -> None:\\n        super().__init__(weight, size_average, reduce, reduction)\\n        self.ignore_index = ignore_index\\n        self.label_smoothing = label_smoothing\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.cross_entropy(\\n            input,\\n            target,\\n            weight=self.weight,\\n            ignore_index=self.ignore_index,\\n            reduction=self.reduction,\\n            label_smoothing=self.label_smoothing,\\n        )\\n\\n\\nclass MultiLabelSoftMarginLoss(_WeightedLoss):\\n    r\\\"\\\"\\\"Creates a criterion that optimizes a multi-label one-versus-all\\n    loss based on max-entropy, between input :math:`x` and target :math:`y` of size\\n    :math:`(N, C)`.\\n    For each sample in the minibatch:\\n\\n    .. math::\\n        loss(x, y) = - \\\\frac{1}{C} * \\\\sum_i y[i] * \\\\log((1 + \\\\exp(-x[i]))^{-1})\\n                         + (1-y[i]) * \\\\log\\\\left(\\\\frac{\\\\exp(-x[i])}{(1 + \\\\exp(-x[i]))}\\\\right)\\n\\n    where :math:`i \\\\in \\\\left\\\\{0, \\\\; \\\\cdots , \\\\; \\\\text{x.nElement}() - 1\\\\right\\\\}`,\\n    :math:`y[i] \\\\in \\\\left\\\\{0, \\\\; 1\\\\right\\\\}`.\\n\\n    Args:\\n        weight (Tensor, optional): a manual rescaling weight given to each\\n            class. If given, it has to be a Tensor of size `C`. Otherwise, it is\\n            treated as if having all ones.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(N, C)` where `N` is the batch size and `C` is the number of classes.\\n        - Target: :math:`(N, C)`, label targets must have the same shape as the input.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(N)`.\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"reduction\\\"]\\n\\n    def __init__(\\n        self,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(weight, size_average, reduce, reduction)\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.multilabel_soft_margin_loss(\\n            input, target, weight=self.weight, reduction=self.reduction\\n        )\\n\\n\\nclass CosineEmbeddingLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the loss given input tensors\\n    :math:`x_1`, :math:`x_2` and a `Tensor` label :math:`y` with values 1 or -1.\\n    Use (:math:`y=1`) to maximize the cosine similarity of two inputs, and (:math:`y=-1`) otherwise.\\n    This is typically used for learning nonlinear\\n    embeddings or semi-supervised learning.\\n\\n    The loss function for each sample is:\\n\\n    .. math::\\n        \\\\text{loss}(x, y) =\\n        \\\\begin{cases}\\n        1 - \\\\cos(x_1, x_2), & \\\\text{if } y = 1 \\\\\\\\\\n        \\\\max(0, \\\\cos(x_1, x_2) - \\\\text{margin}), & \\\\text{if } y = -1\\n        \\\\end{cases}\\n\\n    Args:\\n        margin (float, optional): Should be a number from :math:`-1` to :math:`1`,\\n            :math:`0` to :math:`0.5` is suggested. If :attr:`margin` is missing, the\\n            default value is :math:`0`.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input1: :math:`(N, D)` or :math:`(D)`, where `N` is the batch size and `D` is the embedding dimension.\\n        - Input2: :math:`(N, D)` or :math:`(D)`, same shape as Input1.\\n        - Target: :math:`(N)` or :math:`()`.\\n        - Output: If :attr:`reduction` is ``'none'``, then :math:`(N)`, otherwise scalar.\\n\\n    Examples::\\n\\n        >>> loss = nn.CosineEmbeddingLoss()\\n        >>> input1 = torch.randn(3, 5, requires_grad=True)\\n        >>> input2 = torch.randn(3, 5, requires_grad=True)\\n        >>> target = torch.ones(3)\\n        >>> output = loss(input1, input2, target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"margin\\\", \\\"reduction\\\"]\\n    margin: float\\n\\n    def __init__(\\n        self,\\n        margin: float = 0.0,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.margin = margin\\n\\n    def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor:\\n        return F.cosine_embedding_loss(\\n            input1, input2, target, margin=self.margin, reduction=self.reduction\\n        )\\n\\n\\nclass MarginRankingLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the loss given\\n    inputs :math:`x1`, :math:`x2`, two 1D mini-batch or 0D `Tensors`,\\n    and a label 1D mini-batch or 0D `Tensor` :math:`y` (containing 1 or -1).\\n\\n    If :math:`y = 1` then it assumed the first input should be ranked higher\\n    (have a larger value) than the second input, and vice-versa for :math:`y = -1`.\\n\\n    The loss function for each pair of samples in the mini-batch is:\\n\\n    .. math::\\n        \\\\text{loss}(x1, x2, y) = \\\\max(0, -y * (x1 - x2) + \\\\text{margin})\\n\\n    Args:\\n        margin (float, optional): Has a default value of :math:`0`.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input1: :math:`(N)` or :math:`()` where `N` is the batch size.\\n        - Input2: :math:`(N)` or :math:`()`, same shape as the Input1.\\n        - Target: :math:`(N)` or :math:`()`, same shape as the inputs.\\n        - Output: scalar. If :attr:`reduction` is ``'none'`` and Input size is not :math:`()`, then :math:`(N)`.\\n\\n    Examples::\\n\\n        >>> loss = nn.MarginRankingLoss()\\n        >>> input1 = torch.randn(3, requires_grad=True)\\n        >>> input2 = torch.randn(3, requires_grad=True)\\n        >>> target = torch.randn(3).sign()\\n        >>> output = loss(input1, input2, target)\\n        >>> output.backward()\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"margin\\\", \\\"reduction\\\"]\\n    margin: float\\n\\n    def __init__(\\n        self,\\n        margin: float = 0.0,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(size_average, reduce, reduction)\\n        self.margin = margin\\n\\n    def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor:\\n        return F.margin_ranking_loss(\\n            input1, input2, target, margin=self.margin, reduction=self.reduction\\n        )\\n\\n\\nclass MultiMarginLoss(_WeightedLoss):\\n    r\\\"\\\"\\\"Creates a criterion that optimizes a multi-class classification hinge\\n    loss (margin-based loss) between input :math:`x` (a 2D mini-batch `Tensor`) and\\n    output :math:`y` (which is a 1D tensor of target class indices,\\n    :math:`0 \\\\leq y \\\\leq \\\\text{x.size}(1)-1`):\\n\\n    For each mini-batch sample, the loss in terms of the 1D input :math:`x` and scalar\\n    output :math:`y` is:\\n\\n    .. math::\\n        \\\\text{loss}(x, y) = \\\\frac{\\\\sum_i \\\\max(0, \\\\text{margin} - x[y] + x[i])^p}{\\\\text{x.size}(0)}\\n\\n    where :math:`i \\\\in \\\\left\\\\{0, \\\\; \\\\cdots , \\\\; \\\\text{x.size}(0) - 1\\\\right\\\\}`\\n    and :math:`i \\\\neq y`.\\n\\n    Optionally, you can give non-equal weighting on the classes by passing\\n    a 1D :attr:`weight` tensor into the constructor.\\n\\n    The loss function then becomes:\\n\\n    .. math::\\n        \\\\text{loss}(x, y) = \\\\frac{\\\\sum_i w[y] * \\\\max(0, \\\\text{margin} - x[y] + x[i])^p}{\\\\text{x.size}(0)}\\n\\n    Args:\\n        p (int, optional): Has a default value of :math:`1`. :math:`1` and :math:`2`\\n            are the only supported values.\\n        margin (float, optional): Has a default value of :math:`1`.\\n        weight (Tensor, optional): a manual rescaling weight given to each\\n            class. If given, it has to be a Tensor of size `C`. Otherwise, it is\\n            treated as if having all ones.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(N, C)` or :math:`(C)`, where :math:`N` is the batch size and :math:`C` is the number of classes.\\n        - Target: :math:`(N)` or :math:`()`, where each value is :math:`0 \\\\leq \\\\text{targets}[i] \\\\leq C-1`.\\n        - Output: scalar. If :attr:`reduction` is ``'none'``, then same shape as the target.\\n\\n    Examples::\\n\\n        >>> loss = nn.MultiMarginLoss()\\n        >>> x = torch.tensor([[0.1, 0.2, 0.4, 0.8]])\\n        >>> y = torch.tensor([3])\\n        >>> # 0.25 * ((1-(0.8-0.1)) + (1-(0.8-0.2)) + (1-(0.8-0.4)))\\n        >>> loss(x, y)\\n        tensor(0.32...)\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"p\\\", \\\"margin\\\", \\\"reduction\\\"]\\n    margin: float\\n    p: int\\n\\n    def __init__(\\n        self,\\n        p: int = 1,\\n        margin: float = 1.0,\\n        weight: Optional[Tensor] = None,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ) -> None:\\n        super().__init__(weight, size_average, reduce, reduction)\\n        if p != 1 and p != 2:\\n            raise ValueError(\\\"only p == 1 and p == 2 supported\\\")\\n        if weight is not None and weight.dim() != 1:\\n            raise ValueError(\\n                f\\\"MultiMarginLoss: expected weight to be None or 1D tensor, got {weight.dim()}D instead\\\"\\n            )\\n        self.p = p\\n        self.margin = margin\\n\\n    def forward(self, input: Tensor, target: Tensor) -> Tensor:\\n        return F.multi_margin_loss(\\n            input,\\n            target,\\n            p=self.p,\\n            margin=self.margin,\\n            weight=self.weight,\\n            reduction=self.reduction,\\n        )\\n\\n\\nclass TripletMarginLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the triplet loss given an input\\n    tensors :math:`x1`, :math:`x2`, :math:`x3` and a margin with a value greater than :math:`0`.\\n    This is used for measuring a relative similarity between samples. A triplet\\n    is composed by `a`, `p` and `n` (i.e., `anchor`, `positive examples` and `negative\\n    examples` respectively). The shapes of all input tensors should be\\n    :math:`(N, D)`.\\n\\n    The distance swap is described in detail in the paper `Learning shallow\\n    convolutional feature descriptors with triplet losses`_ by\\n    V. Balntas, E. Riba et al.\\n\\n    The loss function for each sample in the mini-batch is:\\n\\n    .. math::\\n        L(a, p, n) = \\\\max \\\\{d(a_i, p_i) - d(a_i, n_i) + {\\\\rm margin}, 0\\\\}\\n\\n\\n    where\\n\\n    .. math::\\n        d(x_i, y_i) = \\\\left\\\\lVert {\\\\bf x}_i - {\\\\bf y}_i \\\\right\\\\rVert_p\\n\\n    The norm is calculated using the specified p value and a small constant :math:`\\\\varepsilon` is\\n    added for numerical stability.\\n\\n    See also :class:`~torch.nn.TripletMarginWithDistanceLoss`, which computes the\\n    triplet margin loss for input tensors using a custom distance function.\\n\\n    Args:\\n        margin (float, optional): Default: :math:`1`.\\n        p (int, optional): The norm degree for pairwise distance. Default: :math:`2`.\\n        eps (float, optional): Small constant for numerical stability. Default: :math:`1e-6`.\\n        swap (bool, optional): The distance swap is described in detail in the paper\\n            `Learning shallow convolutional feature descriptors with triplet losses` by\\n            V. Balntas, E. Riba et al. Default: ``False``.\\n        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,\\n            the losses are averaged over each loss element in the batch. Note that for\\n            some losses, there are multiple elements per sample. If the field :attr:`size_average`\\n            is set to ``False``, the losses are instead summed for each minibatch. Ignored\\n            when :attr:`reduce` is ``False``. Default: ``True``\\n        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the\\n            losses are averaged or summed over observations for each minibatch depending\\n            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per\\n            batch element instead and ignores :attr:`size_average`. Default: ``True``\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`\\n            and :attr:`reduce` are in the process of being deprecated, and in the meantime,\\n            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``\\n\\n    Shape:\\n        - Input: :math:`(N, D)` or :math:`(D)` where :math:`D` is the vector dimension.\\n        - Output: A Tensor of shape :math:`(N)` if :attr:`reduction` is ``'none'`` and\\n          input shape is :math:`(N, D)`; a scalar otherwise.\\n\\n    Examples::\\n\\n    >>> triplet_loss = nn.TripletMarginLoss(margin=1.0, p=2, eps=1e-7)\\n    >>> anchor = torch.randn(100, 128, requires_grad=True)\\n    >>> positive = torch.randn(100, 128, requires_grad=True)\\n    >>> negative = torch.randn(100, 128, requires_grad=True)\\n    >>> output = triplet_loss(anchor, positive, negative)\\n    >>> output.backward()\\n\\n    .. _Learning shallow convolutional feature descriptors with triplet losses:\\n        http://www.bmva.org/bmvc/2016/papers/paper119/index.html\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"margin\\\", \\\"p\\\", \\\"eps\\\", \\\"swap\\\", \\\"reduction\\\"]\\n    margin: float\\n    p: float\\n    eps: float\\n    swap: bool\\n\\n    def __init__(\\n        self,\\n        margin: float = 1.0,\\n        p: float = 2.0,\\n        eps: float = 1e-6,\\n        swap: bool = False,\\n        size_average=None,\\n        reduce=None,\\n        reduction: str = \\\"mean\\\",\\n    ):\\n        super().__init__(size_average, reduce, reduction)\\n        if margin <= 0:\\n            raise ValueError(\\n                f\\\"TripletMarginLoss: expected margin to be greater than 0, got {margin} instead\\\"\\n            )\\n        self.margin = margin\\n        self.p = p\\n        self.eps = eps\\n        self.swap = swap\\n\\n    def forward(self, anchor: Tensor, positive: Tensor, negative: Tensor) -> Tensor:\\n        return F.triplet_margin_loss(\\n            anchor,\\n            positive,\\n            negative,\\n            margin=self.margin,\\n            p=self.p,\\n            eps=self.eps,\\n            swap=self.swap,\\n            reduction=self.reduction,\\n        )\\n\\n\\nclass TripletMarginWithDistanceLoss(_Loss):\\n    r\\\"\\\"\\\"Creates a criterion that measures the triplet loss given input\\n    tensors :math:`a`, :math:`p`, and :math:`n` (representing anchor,\\n    positive, and negative examples, respectively), and a nonnegative,\\n    real-valued function (\\\"distance function\\\") used to compute the relationship\\n    between the anchor and positive example (\\\"positive distance\\\") and the\\n    anchor and negative example (\\\"negative distance\\\").\\n\\n    The unreduced loss (i.e., with :attr:`reduction` set to ``'none'``)\\n    can be described as:\\n\\n    .. math::\\n        \\\\ell(a, p, n) = L = \\\\{l_1,\\\\dots,l_N\\\\}^\\\\top, \\\\quad\\n        l_i = \\\\max \\\\{d(a_i, p_i) - d(a_i, n_i) + {\\\\rm margin}, 0\\\\}\\n\\n    where :math:`N` is the batch size; :math:`d` is a nonnegative, real-valued function\\n    quantifying the closeness of two tensors, referred to as the :attr:`distance_function`;\\n    and :math:`margin` is a nonnegative margin representing the minimum difference\\n    between the positive and negative distances that is required for the loss to\\n    be 0.  The input tensors have :math:`N` elements each and can be of any shape\\n    that the distance function can handle.\\n\\n    If :attr:`reduction` is not ``'none'``\\n    (default ``'mean'``), then:\\n\\n    .. math::\\n        \\\\ell(x, y) =\\n        \\\\begin{cases}\\n            \\\\operatorname{mean}(L), &  \\\\text{if reduction} = \\\\text{`mean';}\\\\\\\\\\n            \\\\operatorname{sum}(L),  &  \\\\text{if reduction} = \\\\text{`sum'.}\\n        \\\\end{cases}\\n\\n    See also :class:`~torch.nn.TripletMarginLoss`, which computes the triplet\\n    loss for input tensors using the :math:`l_p` distance as the distance function.\\n\\n    Args:\\n        distance_function (Callable, optional): A nonnegative, real-valued function that\\n            quantifies the closeness of two tensors. If not specified,\\n            `nn.PairwiseDistance` will be used.  Default: ``None``\\n        margin (float, optional): A nonnegative margin representing the minimum difference\\n            between the positive and negative distances required for the loss to be 0. Larger\\n            margins penalize cases where the negative examples are not distant enough from the\\n            anchors, relative to the positives. Default: :math:`1`.\\n        swap (bool, optional): Whether to use the distance swap described in the paper\\n            `Learning shallow convolutional feature descriptors with triplet losses` by\\n            V. Balntas, E. Riba et al. If True, and if the positive example is closer to the\\n            negative example than the anchor is, swaps the positive example and the anchor in\\n            the loss computation. Default: ``False``.\\n        reduction (str, optional): Specifies the (optional) reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the sum of the output will be divided by the number of\\n            elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'``\\n\\n\\n    Shape:\\n        - Input: :math:`(N, *)` where :math:`*` represents any number of additional dimensions\\n          as supported by the distance function.\\n        - Output: A Tensor of shape :math:`(N)` if :attr:`reduction` is ``'none'``, or a scalar\\n          otherwise.\\n\\n    Examples::\\n\\n    >>> # Initialize embeddings\\n    >>> embedding = nn.Embedding(1000, 128)\\n    >>> anchor_ids = torch.randint(0, 1000, (1,))\\n    >>> positive_ids = torch.randint(0, 1000, (1,))\\n    >>> negative_ids = torch.randint(0, 1000, (1,))\\n    >>> anchor = embedding(anchor_ids)\\n    >>> positive = embedding(positive_ids)\\n    >>> negative = embedding(negative_ids)\\n    >>>\\n    >>> # Built-in Distance Function\\n    >>> triplet_loss = \\\\\\n    >>>     nn.TripletMarginWithDistanceLoss(distance_function=nn.PairwiseDistance())\\n    >>> output = triplet_loss(anchor, positive, negative)\\n    >>> output.backward()\\n    >>>\\n    >>> # Custom Distance Function\\n    >>> def l_infinity(x1, x2):\\n    >>>     return torch.max(torch.abs(x1 - x2), dim=1).values\\n    >>>\\n    >>> # xdoctest: +SKIP(\\\"FIXME: Would call backwards a second time\\\")\\n    >>> triplet_loss = (\\n    >>>     nn.TripletMarginWithDistanceLoss(distance_function=l_infinity, margin=1.5))\\n    >>> output = triplet_loss(anchor, positive, negative)\\n    >>> output.backward()\\n    >>>\\n    >>> # Custom Distance Function (Lambda)\\n    >>> triplet_loss = (\\n    >>>     nn.TripletMarginWithDistanceLoss(\\n    >>>         distance_function=lambda x, y: 1.0 - F.cosine_similarity(x, y)))\\n    >>> output = triplet_loss(anchor, positive, negative)\\n    >>> output.backward()\\n\\n    Reference:\\n        V. Balntas, et al.: Learning shallow convolutional feature descriptors with triplet losses:\\n        http://www.bmva.org/bmvc/2016/papers/paper119/index.html\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"margin\\\", \\\"swap\\\", \\\"reduction\\\"]\\n    margin: float\\n    swap: bool\\n\\n    def __init__(\\n        self,\\n        *,\\n        distance_function: Optional[Callable[[Tensor, Tensor], Tensor]] = None,\\n        margin: float = 1.0,\\n        swap: bool = False,\\n        reduction: str = \\\"mean\\\",\\n    ):\\n        super().__init__(size_average=None, reduce=None, reduction=reduction)\\n        if margin <= 0:\\n            raise ValueError(\\n                f\\\"TripletMarginWithDistanceLoss: expected margin to be greater than 0, got {margin} instead\\\"\\n            )\\n        self.distance_function: Optional[Callable[[Tensor, Tensor], Tensor]] = (\\n            distance_function if distance_function is not None else PairwiseDistance()\\n        )\\n        self.margin = margin\\n        self.swap = swap\\n\\n    def forward(self, anchor: Tensor, positive: Tensor, negative: Tensor) -> Tensor:\\n        return F.triplet_margin_with_distance_loss(\\n            anchor,\\n            positive,\\n            negative,\\n            distance_function=self.distance_function,\\n            margin=self.margin,\\n            swap=self.swap,\\n            reduction=self.reduction,\\n        )\\n\\n\\nclass CTCLoss(_Loss):\\n    r\\\"\\\"\\\"The Connectionist Temporal Classification loss.\\n\\n    Calculates loss between a continuous (unsegmented) time series and a target sequence. CTCLoss sums over the\\n    probability of possible alignments of input to target, producing a loss value which is differentiable\\n    with respect to each input node. The alignment of input to target is assumed to be \\\"many-to-one\\\", which\\n    limits the length of the target sequence such that it must be :math:`\\\\leq` the input length.\\n\\n    Args:\\n        blank (int, optional): blank label. Default :math:`0`.\\n        reduction (str, optional): Specifies the reduction to apply to the output:\\n            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,\\n            ``'mean'``: the output losses will be divided by the target lengths and\\n            then the mean over the batch is taken, ``'sum'``: the output losses will be summed.\\n            Default: ``'mean'``\\n        zero_infinity (bool, optional):\\n            Whether to zero infinite losses and the associated gradients.\\n            Default: ``False``\\n            Infinite losses mainly occur when the inputs are too short\\n            to be aligned to the targets.\\n\\n    Shape:\\n        - Log_probs: Tensor of size :math:`(T, N, C)` or :math:`(T, C)`,\\n          where :math:`T = \\\\text{input length}`,\\n          :math:`N = \\\\text{batch size}`, and\\n          :math:`C = \\\\text{number of classes (including blank)}`.\\n          The logarithmized probabilities of the outputs (e.g. obtained with\\n          :func:`torch.nn.functional.log_softmax`).\\n        - Targets: Tensor of size :math:`(N, S)` or\\n          :math:`(\\\\operatorname{sum}(\\\\text{target\\\\_lengths}))`,\\n          where :math:`N = \\\\text{batch size}` and\\n          :math:`S = \\\\text{max target length, if shape is } (N, S)`.\\n          It represents the target sequences. Each element in the target\\n          sequence is a class index. And the target index cannot be blank (default=0).\\n          In the :math:`(N, S)` form, targets are padded to the\\n          length of the longest sequence, and stacked.\\n          In the :math:`(\\\\operatorname{sum}(\\\\text{target\\\\_lengths}))` form,\\n          the targets are assumed to be un-padded and\\n          concatenated within 1 dimension.\\n        - Input_lengths: Tuple or tensor of size :math:`(N)` or :math:`()`,\\n          where :math:`N = \\\\text{batch size}`. It represents the lengths of the\\n          inputs (must each be :math:`\\\\leq T`). And the lengths are specified\\n          for each sequence to achieve masking under the assumption that sequences\\n          are padded to equal lengths.\\n        - Target_lengths: Tuple or tensor of size :math:`(N)` or :math:`()`,\\n          where :math:`N = \\\\text{batch size}`. It represents lengths of the targets.\\n          Lengths are specified for each sequence to achieve masking under the\\n          assumption that sequences are padded to equal lengths. If target shape is\\n          :math:`(N,S)`, target_lengths are effectively the stop index\\n          :math:`s_n` for each target sequence, such that ``target_n = targets[n,0:s_n]`` for\\n          each target in a batch. Lengths must each be :math:`\\\\leq S`\\n          If the targets are given as a 1d tensor that is the concatenation of individual\\n          targets, the target_lengths must add up to the total length of the tensor.\\n        - Output: scalar if :attr:`reduction` is ``'mean'`` (default) or\\n          ``'sum'``. If :attr:`reduction` is ``'none'``, then :math:`(N)` if input is batched or\\n          :math:`()` if input is unbatched, where :math:`N = \\\\text{batch size}`.\\n\\n    Examples::\\n\\n        >>> # Target are to be padded\\n        >>> T = 50      # Input sequence length\\n        >>> C = 20      # Number of classes (including blank)\\n        >>> N = 16      # Batch size\\n        >>> S = 30      # Target sequence length of longest target in batch (padding length)\\n        >>> S_min = 10  # Minimum target length, for demonstration purposes\\n        >>>\\n        >>> # Initialize random batch of input vectors, for *size = (T,N,C)\\n        >>> input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()\\n        >>>\\n        >>> # Initialize random batch of targets (0 = blank, 1:C = classes)\\n        >>> target = torch.randint(low=1, high=C, size=(N, S), dtype=torch.long)\\n        >>>\\n        >>> input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)\\n        >>> target_lengths = torch.randint(low=S_min, high=S, size=(N,), dtype=torch.long)\\n        >>> ctc_loss = nn.CTCLoss()\\n        >>> loss = ctc_loss(input, target, input_lengths, target_lengths)\\n        >>> loss.backward()\\n        >>>\\n        >>>\\n        >>> # Target are to be un-padded\\n        >>> T = 50      # Input sequence length\\n        >>> C = 20      # Number of classes (including blank)\\n        >>> N = 16      # Batch size\\n        >>>\\n        >>> # Initialize random batch of input vectors, for *size = (T,N,C)\\n        >>> input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()\\n        >>> input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)\\n        >>>\\n        >>> # Initialize random batch of targets (0 = blank, 1:C = classes)\\n        >>> target_lengths = torch.randint(low=1, high=T, size=(N,), dtype=torch.long)\\n        >>> target = torch.randint(low=1, high=C, size=(sum(target_lengths),), dtype=torch.long)\\n        >>> ctc_loss = nn.CTCLoss()\\n        >>> loss = ctc_loss(input, target, input_lengths, target_lengths)\\n        >>> loss.backward()\\n        >>>\\n        >>>\\n        >>> # Target are to be un-padded and unbatched (effectively N=1)\\n        >>> T = 50      # Input sequence length\\n        >>> C = 20      # Number of classes (including blank)\\n        >>>\\n        >>> # Initialize random batch of input vectors, for *size = (T,C)\\n        >>> # xdoctest: +SKIP(\\\"FIXME: error in doctest\\\")\\n        >>> input = torch.randn(T, C).log_softmax(1).detach().requires_grad_()\\n        >>> input_lengths = torch.tensor(T, dtype=torch.long)\\n        >>>\\n        >>> # Initialize random batch of targets (0 = blank, 1:C = classes)\\n        >>> target_lengths = torch.randint(low=1, high=T, size=(), dtype=torch.long)\\n        >>> target = torch.randint(low=1, high=C, size=(target_lengths,), dtype=torch.long)\\n        >>> ctc_loss = nn.CTCLoss()\\n        >>> loss = ctc_loss(input, target, input_lengths, target_lengths)\\n        >>> loss.backward()\\n\\n    Reference:\\n        A. Graves et al.: Connectionist Temporal Classification:\\n        Labelling Unsegmented Sequence Data with Recurrent Neural Networks:\\n        https://www.cs.toronto.edu/~graves/icml_2006.pdf\\n\\n    Note:\\n        In order to use CuDNN, the following must be satisfied: :attr:`targets` must be\\n        in concatenated format, all :attr:`input_lengths` must be `T`.  :math:`blank=0`,\\n        :attr:`target_lengths` :math:`\\\\leq 256`, the integer arguments must be of\\n        dtype :attr:`torch.int32`.\\n\\n        The regular implementation uses the (more common in PyTorch) `torch.long` dtype.\\n\\n\\n    Note:\\n        In some circumstances when using the CUDA backend with CuDNN, this operator\\n        may select a nondeterministic algorithm to increase performance. If this is\\n        undesirable, you can try to make the operation deterministic (potentially at\\n        a performance cost) by setting ``torch.backends.cudnn.deterministic =\\n        True``.\\n        Please see the notes on :doc:`/notes/randomness` for background.\\n    \\\"\\\"\\\"\\n    __constants__ = [\\\"blank\\\", \\\"reduction\\\"]\\n    blank: int\\n    zero_infinity: bool\\n\\n    def __init__(\\n        self, blank: int = 0, reduction: str = \\\"mean\\\", zero_infinity: bool = False\\n    ):\\n        super().__init__(reduction=reduction)\\n        self.blank = blank\\n        self.zero_infinity = zero_infinity\\n\\n    def forward(\\n        self,\\n        log_probs: Tensor,\\n        targets: Tensor,\\n        input_lengths: Tensor,\\n        target_lengths: Tensor,\\n    ) -> Tensor:\\n        return F.ctc_loss(\\n            log_probs,\\n            targets,\\n            input_lengths,\\n            target_lengths,\\n            self.blank,\\n            self.reduction,\\n            self.zero_infinity,\\n        )\\n\\n\\n# TODO: L1HingeEmbeddingCriterion\\n# TODO: MSECriterion weight\\n# TODO: ClassSimplexCriterion\\n\\n\\n# mypy: allow-untyped-defs\\nimport itertools\\nfrom typing import Any, Optional, Protocol, Type\\n\\nimport torch\\nfrom torch.nn.parameter import is_lazy\\n\\n\\n__all__ = [\\\"LazyModuleMixin\\\"]\\n\\n\\nclass _LazyProtocol(Protocol):\\n    \\\"\\\"\\\"This class is used to avoid errors with mypy checks for the attributes in a mixin.\\n\\n    https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes\\n    \\\"\\\"\\\"\\n\\n    def _register_load_state_dict_pre_hook(self, hook):\\n        ...\\n\\n    def register_forward_pre_hook(self, hook, *, prepend=False, with_kwargs=False):\\n        ...\\n\\n    def _lazy_load_hook(\\n        self,\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ):\\n        ...\\n\\n    def _get_name(self):\\n        ...\\n\\n    def _infer_parameters(self, module, input):\\n        ...\\n\\n    @property\\n    def _parameters(self):\\n        ...\\n\\n    @property\\n    def _buffers(self):\\n        ...\\n\\n    @property\\n    def _non_persistent_buffers_set(self):\\n        ...\\n\\n    @property\\n    def _load_hook(self):\\n        ...\\n\\n    @property\\n    def _initialize_hook(self):\\n        ...\\n\\n\\nclass LazyModuleMixin:\\n    r\\\"\\\"\\\"A mixin for modules that lazily initialize parameters, also known as \\\"lazy modules\\\".\\n\\n    .. warning:\\n        Lazy modules are an experimental new feature under active development,\\n        and their API is likely to change.\\n\\n    Modules that lazily initialize parameters, or \\\"lazy modules\\\",\\n    derive the shapes of their parameters from the first input(s)\\n    to their forward method. Until that first forward they contain\\n    :class:`torch.nn.UninitializedParameter` s that should not be accessed\\n    or used, and afterward they contain regular :class:`torch.nn.Parameter` s.\\n    Lazy modules are convenient since they don't require computing some\\n    module arguments, like the :attr:`in_features` argument of a\\n    typical :class:`torch.nn.Linear`.\\n\\n    After construction, networks with lazy modules should first\\n    be converted to the desired dtype and placed on the expected device.\\n    This is because lazy modules only perform shape inference so the usual dtype\\n    and device placement behavior applies.\\n    The lazy modules should then perform \\\"dry runs\\\" to initialize all the components in the module.\\n    These \\\"dry runs\\\" send inputs of the correct size, dtype, and device through\\n    the network and to each one of its lazy modules. After this the network can be used as usual.\\n\\n    >>> # xdoctest: +SKIP\\n    >>> class LazyMLP(torch.nn.Module):\\n    ...    def __init__(self) -> None:\\n    ...        super().__init__()\\n    ...        self.fc1 = torch.nn.LazyLinear(10)\\n    ...        self.relu1 = torch.nn.ReLU()\\n    ...        self.fc2 = torch.nn.LazyLinear(1)\\n    ...        self.relu2 = torch.nn.ReLU()\\n    ...\\n    ...    def forward(self, input):\\n    ...        x = self.relu1(self.fc1(input))\\n    ...        y = self.relu2(self.fc2(x))\\n    ...        return y\\n    >>> # constructs a network with lazy modules\\n    >>> lazy_mlp = LazyMLP()\\n    >>> # transforms the network's device and dtype\\n    >>> # NOTE: these transforms can and should be applied after construction and before any 'dry runs'\\n    >>> lazy_mlp = lazy_mlp.cuda().double()\\n    >>> lazy_mlp\\n    LazyMLP( (fc1): LazyLinear(in_features=0, out_features=10, bias=True)\\n      (relu1): ReLU()\\n      (fc2): LazyLinear(in_features=0, out_features=1, bias=True)\\n      (relu2): ReLU()\\n    )\\n    >>> # performs a dry run to initialize the network's lazy modules\\n    >>> lazy_mlp(torch.ones(10,10).cuda())\\n    >>> # after initialization, LazyLinear modules become regular Linear modules\\n    >>> lazy_mlp\\n    LazyMLP(\\n      (fc1): Linear(in_features=10, out_features=10, bias=True)\\n      (relu1): ReLU()\\n      (fc2): Linear(in_features=10, out_features=1, bias=True)\\n      (relu2): ReLU()\\n    )\\n    >>> # attaches an optimizer, since parameters can now be used as usual\\n    >>> optim = torch.optim.SGD(mlp.parameters(), lr=0.01)\\n\\n    A final caveat when using lazy modules is that the order of initialization of a network's\\n    parameters may change, since the lazy modules are always initialized after other modules.\\n    For example, if the LazyMLP class defined above had a :class:`torch.nn.LazyLinear` module\\n    first and then a regular :class:`torch.nn.Linear` second, the second module would be\\n    initialized on construction and the first module would be initialized during the first dry run.\\n    This can cause the parameters of a network using lazy modules to be initialized differently\\n    than the parameters of a network without lazy modules as the order of parameter initializations,\\n    which often depends on a stateful random number generator, is different.\\n    Check :doc:`/notes/randomness` for more details.\\n\\n    Lazy modules can be serialized with a state dict like other modules. For example:\\n\\n    >>> lazy_mlp = LazyMLP()\\n    >>> # The state dict shows the uninitialized parameters\\n    >>> lazy_mlp.state_dict()\\n    OrderedDict([('fc1.weight', Uninitialized parameter),\\n                 ('fc1.bias',\\n                  tensor([-1.8832e+25,  4.5636e-41, -1.8832e+25,  4.5636e-41, -6.1598e-30,\\n                           4.5637e-41, -1.8788e+22,  4.5636e-41, -2.0042e-31,  4.5637e-41])),\\n                 ('fc2.weight', Uninitialized parameter),\\n                 ('fc2.bias', tensor([0.0019]))])\\n\\n\\n    Lazy modules can load regular :class:`torch.nn.Parameter` s (i.e. you can serialize/deserialize\\n    initialized LazyModules and they will remain initialized)\\n\\n\\n    >>> full_mlp = LazyMLP()\\n    >>> # Dry run to initialize another module\\n    >>> full_mlp.forward(torch.ones(10, 1))\\n    >>> # Load an initialized state into a lazy module\\n    >>> lazy_mlp.load_state_dict(full_mlp.state_dict())\\n    >>> # The state dict now holds valid values\\n    >>> lazy_mlp.state_dict()\\n    OrderedDict([('fc1.weight',\\n                  tensor([[-0.3837],\\n                          [ 0.0907],\\n                          [ 0.6708],\\n                          [-0.5223],\\n                          [-0.9028],\\n                          [ 0.2851],\\n                          [-0.4537],\\n                          [ 0.6813],\\n                          [ 0.5766],\\n                          [-0.8678]])),\\n                 ('fc1.bias',\\n                  tensor([-1.8832e+25,  4.5636e-41, -1.8832e+25,  4.5636e-41, -6.1598e-30,\\n                           4.5637e-41, -1.8788e+22,  4.5636e-41, -2.0042e-31,  4.5637e-41])),\\n                 ('fc2.weight',\\n                  tensor([[ 0.1320,  0.2938,  0.0679,  0.2793,  0.1088, -0.1795, -0.2301,  0.2807,\\n                            0.2479,  0.1091]])),\\n                 ('fc2.bias', tensor([0.0019]))])\\n\\n    Note, however, that the loaded parameters will not be replaced when doing a \\\"dry run\\\" if they are initialized\\n    when the state is loaded. This prevents using initialized modules in different contexts.\\n    \\\"\\\"\\\"\\n\\n    # modules inheriting from this will change their __class__ to the specified\\n    # one after they are fully initialized\\n    cls_to_become: Optional[Type[Any]] = None\\n\\n    def __init__(self: _LazyProtocol, *args, **kwargs):\\n        # Mypy doesnt like this super call in a mixin\\n        super().__init__(*args, **kwargs)  # type: ignore[misc]\\n        self._load_hook = self._register_load_state_dict_pre_hook(self._lazy_load_hook)\\n        self._initialize_hook = self.register_forward_pre_hook(\\n            self._infer_parameters, with_kwargs=True\\n        )\\n\\n    def _save_to_state_dict(self: _LazyProtocol, destination, prefix, keep_vars):\\n        # This should be ideally implemented as a hook,\\n        # but we should override `detach` in the UninitializedParameter to return itself\\n        # which is not clean\\n        for name, param in self._parameters.items():\\n            if param is not None:\\n                if not (is_lazy(param) or keep_vars):\\n                    param = param.detach()\\n                destination[prefix + name] = param\\n        for name, buf in self._buffers.items():\\n            if buf is not None and name not in self._non_persistent_buffers_set:\\n                if not (is_lazy(buf) or keep_vars):\\n                    buf = buf.detach()\\n                destination[prefix + name] = buf\\n\\n    def _lazy_load_hook(\\n        self: _LazyProtocol,\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ):\\n        \\\"\\\"\\\"load_state_dict pre-hook function for lazy buffers and parameters.\\n\\n        The purpose of this hook is to adjust the current state and/or\\n        ``state_dict`` being loaded so that a module instance serialized in\\n        both un/initialized state can be deserialized onto both un/initialized\\n        module instance.\\n        See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``\\n        for the details of the hook specification.\\n        \\\"\\\"\\\"\\n        for name, param in itertools.chain(\\n            self._parameters.items(), self._buffers.items()\\n        ):\\n            key = prefix + name\\n            if key in state_dict and param is not None:\\n                input_param = state_dict[key]\\n                if is_lazy(param):\\n                    # The current parameter is not initialized but the one being loaded one is\\n                    # create a new parameter based on the uninitialized one\\n                    if not is_lazy(input_param):\\n                        with torch.no_grad():\\n                            param.materialize(input_param.shape)\\n\\n    def initialize_parameters(self: _LazyProtocol, *args, **kwargs):\\n        r\\\"\\\"\\\"Initialize parameters according to the input batch properties.\\n\\n        This adds an interface to isolate parameter initialization from the\\n        forward pass when doing parameter shape inference.\\n        \\\"\\\"\\\"\\n        raise NotImplementedError(\\n            f\\\"initialize_parameters is not implemented for {self.__class__.__name__}\\\"\\n        )\\n\\n    def has_uninitialized_params(self: _LazyProtocol):\\n        r\\\"\\\"\\\"Check if a module has parameters that are not initialized.\\\"\\\"\\\"\\n        # This is to avoid the JIT to track this parameter and force\\n        # custom modules __setstate__ to add it\\n        params = self._parameters.values()\\n        buffers = self._buffers.values()\\n        for param in itertools.chain(params, buffers):\\n            if is_lazy(param):\\n                return True\\n        return False\\n\\n    # torchrec tests the code consistency with the following code\\n    # fmt: off\\n    def _infer_parameters(self: _LazyProtocol, module, args, kwargs=None):\\n        r\\\"\\\"\\\"Infers the size and initializes the parameters according to the provided input batch.\\n\\n        Given a module that contains parameters that were declared inferrable\\n        using :class:`torch.nn.parameter.ParameterMode.Infer`, runs a forward pass\\n        in the complete module using the provided input to initialize all the parameters\\n        as needed.\\n        The module is set into evaluation mode before running the forward pass in order\\n        to avoid saving statistics or calculating gradients\\n        \\\"\\\"\\\"\\n        kwargs = kwargs if kwargs else {}\\n        module.initialize_parameters(*args, **kwargs)\\n        if module.has_uninitialized_params():\\n            raise RuntimeError(f'module {self._get_name()} has not been fully initialized')\\n        module._initialize_hook.remove()\\n        module._load_hook.remove()\\n        delattr(module, '_initialize_hook')\\n        delattr(module, '_load_hook')\\n        if module.cls_to_become is not None:\\n            module.__class__ = module.cls_to_become\\n    # fmt: on\\n\\n    def _replicate_for_data_parallel(self: _LazyProtocol):\\n        raise RuntimeError(\\n            \\\"Modules with uninitialized parameters can't be used with `DataParallel`. \\\"\\n            \\\"Run a dummy forward pass to correctly initialize the modules\\\"\\n        )\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Sequence, Tuple\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.common_types import _size_2_t, _size_4_t, _size_6_t\\n\\nfrom .module import Module\\nfrom .utils import _ntuple, _pair, _quadruple\\n\\n\\n# TODO: grad_output size asserts in THNN\\n\\n__all__ = [\\n    \\\"CircularPad1d\\\",\\n    \\\"CircularPad2d\\\",\\n    \\\"CircularPad3d\\\",\\n    \\\"ConstantPad1d\\\",\\n    \\\"ConstantPad2d\\\",\\n    \\\"ConstantPad3d\\\",\\n    \\\"ReflectionPad1d\\\",\\n    \\\"ReflectionPad2d\\\",\\n    \\\"ReflectionPad3d\\\",\\n    \\\"ReplicationPad1d\\\",\\n    \\\"ReplicationPad2d\\\",\\n    \\\"ReplicationPad3d\\\",\\n    \\\"ZeroPad1d\\\",\\n    \\\"ZeroPad2d\\\",\\n    \\\"ZeroPad3d\\\",\\n]\\n\\n\\nclass _CircularPadNd(Module):\\n    __constants__ = [\\\"padding\\\"]\\n    padding: Sequence[int]\\n\\n    def _check_input_dim(self, input):\\n        raise NotImplementedError\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        self._check_input_dim(input)\\n        return F.pad(input, self.padding, \\\"circular\\\")\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.padding}\\\"\\n\\n\\nclass CircularPad1d(_CircularPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using circular padding of the input boundary.\\n\\n    Tensor values at the beginning of the dimension are used to pad the end,\\n    and values at the end are used to pad the beginning. If negative padding is\\n    applied then the ends of the tensor get removed.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 2-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`)\\n\\n    Shape:\\n        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.\\n        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"not sure why xdoctest is choking on this\\\")\\n        >>> m = nn.CircularPad1d(2)\\n        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)\\n        >>> input\\n        tensor([[[0., 1., 2., 3.],\\n                 [4., 5., 6., 7.]]])\\n        >>> m(input)\\n        tensor([[[2., 3., 0., 1., 2., 3., 0., 1.],\\n                 [6., 7., 4., 5., 6., 7., 4., 5.]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.CircularPad1d((3, 1))\\n        >>> m(input)\\n        tensor([[[1., 2., 3., 0., 1., 2., 3., 0.],\\n                 [5., 6., 7., 4., 5., 6., 7., 4.]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int]\\n\\n    def __init__(self, padding: _size_2_t) -> None:\\n        super().__init__()\\n        self.padding = _pair(padding)\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 2 and input.dim() != 3:\\n            raise ValueError(f\\\"expected 2D or 3D input (got {input.dim()}D input)\\\")\\n\\n\\nclass CircularPad2d(_CircularPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using circular padding of the input boundary.\\n\\n    Tensor values at the beginning of the dimension are used to pad the end,\\n    and values at the end are used to pad the beginning. If negative padding is\\n    applied then the ends of the tensor get removed.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 4-`tuple`, uses (:math:`\\\\text{padding\\\\_left}`,\\n            :math:`\\\\text{padding\\\\_right}`, :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> m = nn.CircularPad2d(2)\\n        >>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)\\n        >>> input\\n        tensor([[[[0., 1., 2.],\\n                  [3., 4., 5.],\\n                  [6., 7., 8.]]]])\\n        >>> m(input)\\n        tensor([[[[4., 5., 3., 4., 5., 3., 4.],\\n                  [7., 8., 6., 7., 8., 6., 7.],\\n                  [1., 2., 0., 1., 2., 0., 1.],\\n                  [4., 5., 3., 4., 5., 3., 4.],\\n                  [7., 8., 6., 7., 8., 6., 7.],\\n                  [1., 2., 0., 1., 2., 0., 1.],\\n                  [4., 5., 3., 4., 5., 3., 4.]]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.CircularPad2d((1, 1, 2, 0))\\n        >>> m(input)\\n        tensor([[[[5., 3., 4., 5., 3.],\\n                  [8., 6., 7., 8., 6.],\\n                  [2., 0., 1., 2., 0.],\\n                  [5., 3., 4., 5., 3.],\\n                  [8., 6., 7., 8., 6.]]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int]\\n\\n    def __init__(self, padding: _size_4_t) -> None:\\n        super().__init__()\\n        self.padding = _quadruple(padding)\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 3 and input.dim() != 4:\\n            raise ValueError(f\\\"expected 3D or 4D input (got {input.dim()}D input)\\\")\\n\\n\\nclass CircularPad3d(_CircularPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using circular padding of the input boundary.\\n\\n    Tensor values at the beginning of the dimension are used to pad the end,\\n    and values at the end are used to pad the beginning. If negative padding is\\n    applied then the ends of the tensor get removed.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 6-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`,\\n            :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`,\\n            :math:`\\\\text{padding\\\\_front}`, :math:`\\\\text{padding\\\\_back}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,\\n          where\\n\\n          :math:`D_{out} = D_{in} + \\\\text{padding\\\\_front} + \\\\text{padding\\\\_back}`\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> m = nn.CircularPad3d(3)\\n        >>> input = torch.randn(16, 3, 8, 320, 480)\\n        >>> output = m(input)\\n        >>> # using different paddings for different sides\\n        >>> m = nn.CircularPad3d((3, 3, 6, 6, 1, 1))\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int, int, int]\\n\\n    def __init__(self, padding: _size_6_t) -> None:\\n        super().__init__()\\n        self.padding = _ntuple(6)(padding)\\n\\n    def _check_input_dim(self, input):\\n        if input.dim() != 4 and input.dim() != 5:\\n            raise ValueError(f\\\"expected 4D or 5D input (got {input.dim()}D input)\\\")\\n\\n\\nclass _ConstantPadNd(Module):\\n    __constants__ = [\\\"padding\\\", \\\"value\\\"]\\n    value: float\\n    padding: Sequence[int]\\n\\n    def __init__(self, value: float) -> None:\\n        super().__init__()\\n        self.value = value\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.pad(input, self.padding, \\\"constant\\\", self.value)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"padding={self.padding}, value={self.value}\\\"\\n\\n\\nclass ConstantPad1d(_ConstantPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor boundaries with a constant value.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in both boundaries. If a 2-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`)\\n\\n    Shape:\\n        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.\\n        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> m = nn.ConstantPad1d(2, 3.5)\\n        >>> input = torch.randn(1, 2, 4)\\n        >>> input\\n        tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],\\n                 [-1.3287,  1.8966,  0.1466, -0.2771]]])\\n        >>> m(input)\\n        tensor([[[ 3.5000,  3.5000, -1.0491, -0.7152, -0.0749,  0.8530,  3.5000,\\n                   3.5000],\\n                 [ 3.5000,  3.5000, -1.3287,  1.8966,  0.1466, -0.2771,  3.5000,\\n                   3.5000]]])\\n        >>> m = nn.ConstantPad1d(2, 3.5)\\n        >>> input = torch.randn(1, 2, 3)\\n        >>> input\\n        tensor([[[ 1.6616,  1.4523, -1.1255],\\n                 [-3.6372,  0.1182, -1.8652]]])\\n        >>> m(input)\\n        tensor([[[ 3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000,  3.5000]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ConstantPad1d((3, 1), 3.5)\\n        >>> m(input)\\n        tensor([[[ 3.5000,  3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000],\\n                 [ 3.5000,  3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int]\\n\\n    def __init__(self, padding: _size_2_t, value: float):\\n        super().__init__(value)\\n        self.padding = _pair(padding)\\n\\n\\nclass ConstantPad2d(_ConstantPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor boundaries with a constant value.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 4-`tuple`, uses (:math:`\\\\text{padding\\\\_left}`,\\n            :math:`\\\\text{padding\\\\_right}`, :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> m = nn.ConstantPad2d(2, 3.5)\\n        >>> input = torch.randn(1, 2, 2)\\n        >>> input\\n        tensor([[[ 1.6585,  0.4320],\\n                 [-0.8701, -0.4649]]])\\n        >>> m(input)\\n        tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000,  1.6585,  0.4320,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000, -0.8701, -0.4649,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ConstantPad2d((3, 0, 2, 1), 3.5)\\n        >>> m(input)\\n        tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],\\n                 [ 3.5000,  3.5000,  3.5000,  1.6585,  0.4320],\\n                 [ 3.5000,  3.5000,  3.5000, -0.8701, -0.4649],\\n                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\\"padding\\\", \\\"value\\\"]\\n    padding: Tuple[int, int, int, int]\\n\\n    def __init__(self, padding: _size_4_t, value: float) -> None:\\n        super().__init__(value)\\n        self.padding = _quadruple(padding)\\n\\n\\nclass ConstantPad3d(_ConstantPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor boundaries with a constant value.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 6-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`,\\n            :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`,\\n            :math:`\\\\text{padding\\\\_front}`, :math:`\\\\text{padding\\\\_back}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or\\n          :math:`(C, D_{out}, H_{out}, W_{out})`, where\\n\\n          :math:`D_{out} = D_{in} + \\\\text{padding\\\\_front} + \\\\text{padding\\\\_back}`\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> m = nn.ConstantPad3d(3, 3.5)\\n        >>> input = torch.randn(16, 3, 10, 20, 30)\\n        >>> output = m(input)\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ConstantPad3d((3, 3, 6, 6, 0, 1), 3.5)\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int, int, int]\\n\\n    def __init__(self, padding: _size_6_t, value: float) -> None:\\n        super().__init__(value)\\n        self.padding = _ntuple(6)(padding)\\n\\n\\nclass _ReflectionPadNd(Module):\\n    __constants__ = [\\\"padding\\\"]\\n    padding: Sequence[int]\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.pad(input, self.padding, \\\"reflect\\\")\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.padding}\\\"\\n\\n\\nclass ReflectionPad1d(_ReflectionPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using the reflection of the input boundary.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 2-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`)\\n\\n    Shape:\\n        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.\\n        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> m = nn.ReflectionPad1d(2)\\n        >>> # xdoctest: +IGNORE_WANT(\\\"other tests seem to modify printing styles\\\")\\n        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)\\n        >>> input\\n        tensor([[[0., 1., 2., 3.],\\n                 [4., 5., 6., 7.]]])\\n        >>> m(input)\\n        tensor([[[2., 1., 0., 1., 2., 3., 2., 1.],\\n                 [6., 5., 4., 5., 6., 7., 6., 5.]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ReflectionPad1d((3, 1))\\n        >>> m(input)\\n        tensor([[[3., 2., 1., 0., 1., 2., 3., 2.],\\n                 [7., 6., 5., 4., 5., 6., 7., 6.]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int]\\n\\n    def __init__(self, padding: _size_2_t) -> None:\\n        super().__init__()\\n        self.padding = _pair(padding)\\n\\n\\nclass ReflectionPad2d(_ReflectionPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using the reflection of the input boundary.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 4-`tuple`, uses (:math:`\\\\text{padding\\\\_left}`,\\n            :math:`\\\\text{padding\\\\_right}`, :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`)\\n            Note that padding size should be less than the corresponding input dimension.\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})` where\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"not sure why xdoctest is choking on this\\\")\\n        >>> m = nn.ReflectionPad2d(2)\\n        >>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)\\n        >>> input\\n        tensor([[[[0., 1., 2.],\\n                  [3., 4., 5.],\\n                  [6., 7., 8.]]]])\\n        >>> m(input)\\n        tensor([[[[8., 7., 6., 7., 8., 7., 6.],\\n                  [5., 4., 3., 4., 5., 4., 3.],\\n                  [2., 1., 0., 1., 2., 1., 0.],\\n                  [5., 4., 3., 4., 5., 4., 3.],\\n                  [8., 7., 6., 7., 8., 7., 6.],\\n                  [5., 4., 3., 4., 5., 4., 3.],\\n                  [2., 1., 0., 1., 2., 1., 0.]]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ReflectionPad2d((1, 1, 2, 0))\\n        >>> m(input)\\n        tensor([[[[7., 6., 7., 8., 7.],\\n                  [4., 3., 4., 5., 4.],\\n                  [1., 0., 1., 2., 1.],\\n                  [4., 3., 4., 5., 4.],\\n                  [7., 6., 7., 8., 7.]]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int]\\n\\n    def __init__(self, padding: _size_4_t) -> None:\\n        super().__init__()\\n        self.padding = _quadruple(padding)\\n\\n\\nclass ReflectionPad3d(_ReflectionPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using the reflection of the input boundary.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 6-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`,\\n            :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`,\\n            :math:`\\\\text{padding\\\\_front}`, :math:`\\\\text{padding\\\\_back}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,\\n          where\\n\\n          :math:`D_{out} = D_{in} + \\\\text{padding\\\\_front} + \\\\text{padding\\\\_back}`\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"not sure why xdoctest is choking on this\\\")\\n        >>> m = nn.ReflectionPad3d(1)\\n        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 1, 2, 2, 2)\\n        >>> m(input)\\n        tensor([[[[[7., 6., 7., 6.],\\n                   [5., 4., 5., 4.],\\n                   [7., 6., 7., 6.],\\n                   [5., 4., 5., 4.]],\\n                  [[3., 2., 3., 2.],\\n                   [1., 0., 1., 0.],\\n                   [3., 2., 3., 2.],\\n                   [1., 0., 1., 0.]],\\n                  [[7., 6., 7., 6.],\\n                   [5., 4., 5., 4.],\\n                   [7., 6., 7., 6.],\\n                   [5., 4., 5., 4.]],\\n                  [[3., 2., 3., 2.],\\n                   [1., 0., 1., 0.],\\n                   [3., 2., 3., 2.],\\n                   [1., 0., 1., 0.]]]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int, int, int]\\n\\n    def __init__(self, padding: _size_6_t) -> None:\\n        super().__init__()\\n        self.padding = _ntuple(6)(padding)\\n\\n\\nclass _ReplicationPadNd(Module):\\n    __constants__ = [\\\"padding\\\"]\\n    padding: Sequence[int]\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.pad(input, self.padding, \\\"replicate\\\")\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.padding}\\\"\\n\\n\\nclass ReplicationPad1d(_ReplicationPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using replication of the input boundary.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 2-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`)\\n\\n    Shape:\\n        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.\\n        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"not sure why xdoctest is choking on this\\\")\\n        >>> m = nn.ReplicationPad1d(2)\\n        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)\\n        >>> input\\n        tensor([[[0., 1., 2., 3.],\\n                 [4., 5., 6., 7.]]])\\n        >>> m(input)\\n        tensor([[[0., 0., 0., 1., 2., 3., 3., 3.],\\n                 [4., 4., 4., 5., 6., 7., 7., 7.]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ReplicationPad1d((3, 1))\\n        >>> m(input)\\n        tensor([[[0., 0., 0., 0., 1., 2., 3., 3.],\\n                 [4., 4., 4., 4., 5., 6., 7., 7.]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int]\\n\\n    def __init__(self, padding: _size_2_t) -> None:\\n        super().__init__()\\n        self.padding = _pair(padding)\\n\\n\\nclass ReplicationPad2d(_ReplicationPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using replication of the input boundary.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 4-`tuple`, uses (:math:`\\\\text{padding\\\\_left}`,\\n            :math:`\\\\text{padding\\\\_right}`, :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> m = nn.ReplicationPad2d(2)\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)\\n        >>> input\\n        tensor([[[[0., 1., 2.],\\n                  [3., 4., 5.],\\n                  [6., 7., 8.]]]])\\n        >>> m(input)\\n        tensor([[[[0., 0., 0., 1., 2., 2., 2.],\\n                  [0., 0., 0., 1., 2., 2., 2.],\\n                  [0., 0., 0., 1., 2., 2., 2.],\\n                  [3., 3., 3., 4., 5., 5., 5.],\\n                  [6., 6., 6., 7., 8., 8., 8.],\\n                  [6., 6., 6., 7., 8., 8., 8.],\\n                  [6., 6., 6., 7., 8., 8., 8.]]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ReplicationPad2d((1, 1, 2, 0))\\n        >>> m(input)\\n        tensor([[[[0., 0., 1., 2., 2.],\\n                  [0., 0., 1., 2., 2.],\\n                  [0., 0., 1., 2., 2.],\\n                  [3., 3., 4., 5., 5.],\\n                  [6., 6., 7., 8., 8.]]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int]\\n\\n    def __init__(self, padding: _size_4_t) -> None:\\n        super().__init__()\\n        self.padding = _quadruple(padding)\\n\\n\\nclass ReplicationPad3d(_ReplicationPadNd):\\n    r\\\"\\\"\\\"Pads the input tensor using replication of the input boundary.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 6-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`,\\n            :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`,\\n            :math:`\\\\text{padding\\\\_front}`, :math:`\\\\text{padding\\\\_back}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,\\n          where\\n\\n          :math:`D_{out} = D_{in} + \\\\text{padding\\\\_front} + \\\\text{padding\\\\_back}`\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> m = nn.ReplicationPad3d(3)\\n        >>> input = torch.randn(16, 3, 8, 320, 480)\\n        >>> output = m(input)\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ReplicationPad3d((3, 3, 6, 6, 1, 1))\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int, int, int]\\n\\n    def __init__(self, padding: _size_6_t) -> None:\\n        super().__init__()\\n        self.padding = _ntuple(6)(padding)\\n\\n\\nclass ZeroPad1d(ConstantPad1d):\\n    r\\\"\\\"\\\"Pads the input tensor boundaries with zero.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in both boundaries. If a 2-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`)\\n\\n    Shape:\\n        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.\\n        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> m = nn.ZeroPad1d(2)\\n        >>> input = torch.randn(1, 2, 4)\\n        >>> input\\n        tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],\\n                 [-1.3287,  1.8966,  0.1466, -0.2771]]])\\n        >>> m(input)\\n        tensor([[[ 0.0000,  0.0000, -1.0491, -0.7152, -0.0749,  0.8530,  0.0000,\\n                   0.0000],\\n                 [ 0.0000,  0.0000, -1.3287,  1.8966,  0.1466, -0.2771,  0.0000,\\n                   0.0000]]])\\n        >>> m = nn.ZeroPad1d(2)\\n        >>> input = torch.randn(1, 2, 3)\\n        >>> input\\n        tensor([[[ 1.6616,  1.4523, -1.1255],\\n                 [-3.6372,  0.1182, -1.8652]]])\\n        >>> m(input)\\n        tensor([[[ 0.0000,  0.0000,  1.6616,  1.4523, -1.1255,  0.0000,  0.0000],\\n                 [ 0.0000,  0.0000, -3.6372,  0.1182, -1.8652,  0.0000,  0.0000]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ZeroPad1d((3, 1))\\n        >>> m(input)\\n        tensor([[[ 0.0000,  0.0000,  0.0000,  1.6616,  1.4523, -1.1255,  0.0000],\\n                 [ 0.0000,  0.0000,  0.0000, -3.6372,  0.1182, -1.8652,  0.0000]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int]\\n\\n    def __init__(self, padding: _size_2_t) -> None:\\n        super().__init__(padding, 0.0)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.padding}\\\"\\n\\n\\nclass ZeroPad2d(ConstantPad2d):\\n    r\\\"\\\"\\\"Pads the input tensor boundaries with zero.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 4-`tuple`, uses (:math:`\\\\text{padding\\\\_left}`,\\n            :math:`\\\\text{padding\\\\_right}`, :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> m = nn.ZeroPad2d(2)\\n        >>> input = torch.randn(1, 1, 3, 3)\\n        >>> input\\n        tensor([[[[-0.1678, -0.4418,  1.9466],\\n                  [ 0.9604, -0.4219, -0.5241],\\n                  [-0.9162, -0.5436, -0.6446]]]])\\n        >>> m(input)\\n        tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000, -0.1678, -0.4418,  1.9466,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000,  0.9604, -0.4219, -0.5241,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000, -0.9162, -0.5436, -0.6446,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]])\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ZeroPad2d((1, 1, 2, 0))\\n        >>> m(input)\\n        tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\\n                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\\n                  [ 0.0000, -0.1678, -0.4418,  1.9466,  0.0000],\\n                  [ 0.0000,  0.9604, -0.4219, -0.5241,  0.0000],\\n                  [ 0.0000, -0.9162, -0.5436, -0.6446,  0.0000]]]])\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int]\\n\\n    def __init__(self, padding: _size_4_t) -> None:\\n        super().__init__(padding, 0.0)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.padding}\\\"\\n\\n\\nclass ZeroPad3d(ConstantPad3d):\\n    r\\\"\\\"\\\"Pads the input tensor boundaries with zero.\\n\\n    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.\\n\\n    Args:\\n        padding (int, tuple): the size of the padding. If is `int`, uses the same\\n            padding in all boundaries. If a 6-`tuple`, uses\\n            (:math:`\\\\text{padding\\\\_left}`, :math:`\\\\text{padding\\\\_right}`,\\n            :math:`\\\\text{padding\\\\_top}`, :math:`\\\\text{padding\\\\_bottom}`,\\n            :math:`\\\\text{padding\\\\_front}`, :math:`\\\\text{padding\\\\_back}`)\\n\\n    Shape:\\n        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.\\n        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or\\n          :math:`(C, D_{out}, H_{out}, W_{out})`, where\\n\\n          :math:`D_{out} = D_{in} + \\\\text{padding\\\\_front} + \\\\text{padding\\\\_back}`\\n\\n          :math:`H_{out} = H_{in} + \\\\text{padding\\\\_top} + \\\\text{padding\\\\_bottom}`\\n\\n          :math:`W_{out} = W_{in} + \\\\text{padding\\\\_left} + \\\\text{padding\\\\_right}`\\n\\n    Examples::\\n\\n        >>> m = nn.ZeroPad3d(3)\\n        >>> input = torch.randn(16, 3, 10, 20, 30)\\n        >>> output = m(input)\\n        >>> # using different paddings for different sides\\n        >>> m = nn.ZeroPad3d((3, 3, 6, 6, 0, 1))\\n        >>> output = m(input)\\n    \\\"\\\"\\\"\\n\\n    padding: Tuple[int, int, int, int, int, int]\\n\\n    def __init__(self, padding: _size_6_t) -> None:\\n        super().__init__(padding, 0.0)\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"{self.padding}\\\"\\n\\n\\n# mypy: allow-untyped-decorators\\n# mypy: allow-untyped-defs\\nimport operator\\nfrom collections import abc as container_abcs, OrderedDict\\nfrom itertools import chain, islice\\nfrom typing import (\\n    Any,\\n    Dict,\\n    Iterable,\\n    Iterator,\\n    Mapping,\\n    Optional,\\n    overload,\\n    Tuple,\\n    TypeVar,\\n    Union,\\n)\\nfrom typing_extensions import deprecated, Self\\n\\nimport torch\\nfrom torch._jit_internal import _copy_to_script_wrapper\\nfrom torch.nn.parameter import Parameter\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"Container\\\",\\n    \\\"Sequential\\\",\\n    \\\"ModuleList\\\",\\n    \\\"ModuleDict\\\",\\n    \\\"ParameterList\\\",\\n    \\\"ParameterDict\\\",\\n]\\n\\nT = TypeVar(\\\"T\\\", bound=Module)\\n\\n\\n# Copied from torch.nn.modules.module, required for a custom __repr__ for ModuleList\\ndef _addindent(s_, numSpaces):\\n    s = s_.split(\\\"\\\\n\\\")\\n    # don't do anything for single-line stuff\\n    if len(s) == 1:\\n        return s_\\n    first = s.pop(0)\\n    s = [(numSpaces * \\\" \\\") + line for line in s]\\n    s = \\\"\\\\n\\\".join(s)\\n    s = first + \\\"\\\\n\\\" + s\\n    return s\\n\\n\\n@deprecated(\\n    \\\"`nn.Container` is deprecated. \\\"\\n    \\\"All of it's functionality is now implemented in `nn.Module`. Subclass that instead.\\\",\\n    category=FutureWarning,\\n)\\nclass Container(Module):\\n    def __init__(self, **kwargs: Any) -> None:\\n        super().__init__()\\n        for key, value in kwargs.items():\\n            self.add_module(key, value)\\n\\n\\nclass Sequential(Module):\\n    r\\\"\\\"\\\"A sequential container.\\n\\n    Modules will be added to it in the order they are passed in the\\n    constructor. Alternatively, an ``OrderedDict`` of modules can be\\n    passed in. The ``forward()`` method of ``Sequential`` accepts any\\n    input and forwards it to the first module it contains. It then\\n    \\\"chains\\\" outputs to inputs sequentially for each subsequent module,\\n    finally returning the output of the last module.\\n\\n    The value a ``Sequential`` provides over manually calling a sequence\\n    of modules is that it allows treating the whole container as a\\n    single module, such that performing a transformation on the\\n    ``Sequential`` applies to each of the modules it stores (which are\\n    each a registered submodule of the ``Sequential``).\\n\\n    What's the difference between a ``Sequential`` and a\\n    :class:`torch.nn.ModuleList`? A ``ModuleList`` is exactly what it\\n    sounds like--a list for storing ``Module`` s! On the other hand,\\n    the layers in a ``Sequential`` are connected in a cascading way.\\n\\n    Example::\\n\\n        # Using Sequential to create a small model. When `model` is run,\\n        # input will first be passed to `Conv2d(1,20,5)`. The output of\\n        # `Conv2d(1,20,5)` will be used as the input to the first\\n        # `ReLU`; the output of the first `ReLU` will become the input\\n        # for `Conv2d(20,64,5)`. Finally, the output of\\n        # `Conv2d(20,64,5)` will be used as input to the second `ReLU`\\n        model = nn.Sequential(\\n                  nn.Conv2d(1,20,5),\\n                  nn.ReLU(),\\n                  nn.Conv2d(20,64,5),\\n                  nn.ReLU()\\n                )\\n\\n        # Using Sequential with OrderedDict. This is functionally the\\n        # same as the above code\\n        model = nn.Sequential(OrderedDict([\\n                  ('conv1', nn.Conv2d(1,20,5)),\\n                  ('relu1', nn.ReLU()),\\n                  ('conv2', nn.Conv2d(20,64,5)),\\n                  ('relu2', nn.ReLU())\\n                ]))\\n    \\\"\\\"\\\"\\n\\n    _modules: Dict[str, Module]  # type: ignore[assignment]\\n\\n    @overload\\n    def __init__(self, *args: Module) -> None:\\n        ...\\n\\n    @overload\\n    def __init__(self, arg: \\\"OrderedDict[str, Module]\\\") -> None:\\n        ...\\n\\n    def __init__(self, *args):\\n        super().__init__()\\n        if len(args) == 1 and isinstance(args[0], OrderedDict):\\n            for key, module in args[0].items():\\n                self.add_module(key, module)\\n        else:\\n            for idx, module in enumerate(args):\\n                self.add_module(str(idx), module)\\n\\n    def _get_item_by_idx(self, iterator, idx) -> T:  # type: ignore[misc, type-var]\\n        \\\"\\\"\\\"Get the idx-th item of the iterator.\\\"\\\"\\\"\\n        size = len(self)\\n        idx = operator.index(idx)\\n        if not -size <= idx < size:\\n            raise IndexError(f\\\"index {idx} is out of range\\\")\\n        idx %= size\\n        return next(islice(iterator, idx, None))\\n\\n    @_copy_to_script_wrapper\\n    def __getitem__(self, idx: Union[slice, int]) -> Union[\\\"Sequential\\\", T]:\\n        if isinstance(idx, slice):\\n            return self.__class__(OrderedDict(list(self._modules.items())[idx]))\\n        else:\\n            return self._get_item_by_idx(self._modules.values(), idx)\\n\\n    def __setitem__(self, idx: int, module: Module) -> None:\\n        key: str = self._get_item_by_idx(self._modules.keys(), idx)\\n        return setattr(self, key, module)\\n\\n    def __delitem__(self, idx: Union[slice, int]) -> None:\\n        if isinstance(idx, slice):\\n            for key in list(self._modules.keys())[idx]:\\n                delattr(self, key)\\n        else:\\n            key = self._get_item_by_idx(self._modules.keys(), idx)\\n            delattr(self, key)\\n        # To preserve numbering\\n        str_indices = [str(i) for i in range(len(self._modules))]\\n        self._modules = OrderedDict(list(zip(str_indices, self._modules.values())))\\n\\n    @_copy_to_script_wrapper\\n    def __len__(self) -> int:\\n        return len(self._modules)\\n\\n    def __add__(self, other) -> \\\"Sequential\\\":\\n        if isinstance(other, Sequential):\\n            ret = Sequential()\\n            for layer in self:\\n                ret.append(layer)\\n            for layer in other:\\n                ret.append(layer)\\n            return ret\\n        else:\\n            raise ValueError(\\n                \\\"add operator supports only objects \\\"\\n                f\\\"of Sequential class, but {str(type(other))} is given.\\\"\\n            )\\n\\n    def pop(self, key: Union[int, slice]) -> Module:\\n        v = self[key]\\n        del self[key]\\n        return v\\n\\n    def __iadd__(self, other) -> Self:\\n        if isinstance(other, Sequential):\\n            offset = len(self)\\n            for i, module in enumerate(other):\\n                self.add_module(str(i + offset), module)\\n            return self\\n        else:\\n            raise ValueError(\\n                \\\"add operator supports only objects \\\"\\n                f\\\"of Sequential class, but {str(type(other))} is given.\\\"\\n            )\\n\\n    def __mul__(self, other: int) -> \\\"Sequential\\\":\\n        if not isinstance(other, int):\\n            raise TypeError(\\n                f\\\"unsupported operand type(s) for *: {type(self)} and {type(other)}\\\"\\n            )\\n        elif other <= 0:\\n            raise ValueError(\\n                f\\\"Non-positive multiplication factor {other} for {type(self)}\\\"\\n            )\\n        else:\\n            combined = Sequential()\\n            offset = 0\\n            for _ in range(other):\\n                for module in self:\\n                    combined.add_module(str(offset), module)\\n                    offset += 1\\n            return combined\\n\\n    def __rmul__(self, other: int) -> \\\"Sequential\\\":\\n        return self.__mul__(other)\\n\\n    def __imul__(self, other: int) -> Self:\\n        if not isinstance(other, int):\\n            raise TypeError(\\n                f\\\"unsupported operand type(s) for *: {type(self)} and {type(other)}\\\"\\n            )\\n        elif other <= 0:\\n            raise ValueError(\\n                f\\\"Non-positive multiplication factor {other} for {type(self)}\\\"\\n            )\\n        else:\\n            len_original = len(self)\\n            offset = len(self)\\n            for _ in range(other - 1):\\n                for i in range(len_original):\\n                    self.add_module(str(i + offset), self._modules[str(i)])\\n                offset += len_original\\n            return self\\n\\n    @_copy_to_script_wrapper\\n    def __dir__(self):\\n        keys = super().__dir__()\\n        keys = [key for key in keys if not key.isdigit()]\\n        return keys\\n\\n    @_copy_to_script_wrapper\\n    def __iter__(self) -> Iterator[Module]:\\n        return iter(self._modules.values())\\n\\n    # NB: We can't really type check this function as the type of input\\n    # may change dynamically (as is tested in\\n    # TestScript.test_sequential_intermediary_types).  Cannot annotate\\n    # with Any as TorchScript expects a more precise type\\n    def forward(self, input):\\n        for module in self:\\n            input = module(input)\\n        return input\\n\\n    def append(self, module: Module) -> \\\"Sequential\\\":\\n        r\\\"\\\"\\\"Append a given module to the end.\\n\\n        Args:\\n            module (nn.Module): module to append\\n        \\\"\\\"\\\"\\n        self.add_module(str(len(self)), module)\\n        return self\\n\\n    def insert(self, index: int, module: Module) -> \\\"Sequential\\\":\\n        if not isinstance(module, Module):\\n            raise AssertionError(f\\\"module should be of type: {Module}\\\")\\n        n = len(self._modules)\\n        if not (-n <= index <= n):\\n            raise IndexError(f\\\"Index out of range: {index}\\\")\\n        if index < 0:\\n            index += n\\n        for i in range(n, index, -1):\\n            self._modules[str(i)] = self._modules[str(i - 1)]\\n        self._modules[str(index)] = module\\n        return self\\n\\n    def extend(self, sequential) -> \\\"Sequential\\\":\\n        for layer in sequential:\\n            self.append(layer)\\n        return self\\n\\n\\nclass ModuleList(Module):\\n    r\\\"\\\"\\\"Holds submodules in a list.\\n\\n    :class:`~torch.nn.ModuleList` can be indexed like a regular Python list, but\\n    modules it contains are properly registered, and will be visible by all\\n    :class:`~torch.nn.Module` methods.\\n\\n    Args:\\n        modules (iterable, optional): an iterable of modules to add\\n\\n    Example::\\n\\n        class MyModule(nn.Module):\\n            def __init__(self) -> None:\\n                super().__init__()\\n                self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])\\n\\n            def forward(self, x):\\n                # ModuleList can act as an iterable, or be indexed using ints\\n                for i, l in enumerate(self.linears):\\n                    x = self.linears[i // 2](x) + l(x)\\n                return x\\n    \\\"\\\"\\\"\\n\\n    _modules: Dict[str, Module]  # type: ignore[assignment]\\n\\n    def __init__(self, modules: Optional[Iterable[Module]] = None) -> None:\\n        super().__init__()\\n        if modules is not None:\\n            self += modules\\n\\n    def _get_abs_string_index(self, idx):\\n        \\\"\\\"\\\"Get the absolute index for the list of modules.\\\"\\\"\\\"\\n        idx = operator.index(idx)\\n        if not (-len(self) <= idx < len(self)):\\n            raise IndexError(f\\\"index {idx} is out of range\\\")\\n        if idx < 0:\\n            idx += len(self)\\n        return str(idx)\\n\\n    @overload\\n    def __getitem__(self, idx: slice) -> \\\"ModuleList\\\":\\n        ...\\n\\n    @overload\\n    def __getitem__(self, idx: int) -> Module:\\n        ...\\n\\n    @_copy_to_script_wrapper\\n    def __getitem__(self, idx: Union[int, slice]) -> Union[Module, \\\"ModuleList\\\"]:\\n        if isinstance(idx, slice):\\n            return self.__class__(list(self._modules.values())[idx])\\n        else:\\n            return self._modules[self._get_abs_string_index(idx)]\\n\\n    def __setitem__(self, idx: int, module: Module) -> None:\\n        idx = self._get_abs_string_index(idx)\\n        return setattr(self, str(idx), module)\\n\\n    def __delitem__(self, idx: Union[int, slice]) -> None:\\n        if isinstance(idx, slice):\\n            for k in range(len(self._modules))[idx]:\\n                delattr(self, str(k))\\n        else:\\n            delattr(self, self._get_abs_string_index(idx))\\n        # To preserve numbering, self._modules is being reconstructed with modules after deletion\\n        str_indices = [str(i) for i in range(len(self._modules))]\\n        self._modules = OrderedDict(list(zip(str_indices, self._modules.values())))\\n\\n    @_copy_to_script_wrapper\\n    def __len__(self) -> int:\\n        return len(self._modules)\\n\\n    @_copy_to_script_wrapper\\n    def __iter__(self) -> Iterator[Module]:\\n        return iter(self._modules.values())\\n\\n    def __iadd__(self, modules: Iterable[Module]) -> Self:\\n        return self.extend(modules)\\n\\n    def __add__(self, other: Iterable[Module]) -> \\\"ModuleList\\\":\\n        combined = ModuleList()\\n        for i, module in enumerate(chain(self, other)):\\n            combined.add_module(str(i), module)\\n        return combined\\n\\n    def __repr__(self):\\n        \\\"\\\"\\\"Return a custom repr for ModuleList that compresses repeated module representations.\\\"\\\"\\\"\\n        list_of_reprs = [repr(item) for item in self]\\n        if len(list_of_reprs) == 0:\\n            return self._get_name() + \\\"()\\\"\\n\\n        start_end_indices = [[0, 0]]\\n        repeated_blocks = [list_of_reprs[0]]\\n        for i, r in enumerate(list_of_reprs[1:], 1):\\n            if r == repeated_blocks[-1]:\\n                start_end_indices[-1][1] += 1\\n                continue\\n\\n            start_end_indices.append([i, i])\\n            repeated_blocks.append(r)\\n\\n        lines = []\\n        main_str = self._get_name() + \\\"(\\\"\\n        for (start_id, end_id), b in zip(start_end_indices, repeated_blocks):\\n            local_repr = f\\\"({start_id}): {b}\\\"  # default repr\\n\\n            if start_id != end_id:\\n                n = end_id - start_id + 1\\n                local_repr = f\\\"({start_id}-{end_id}): {n} x {b}\\\"\\n\\n            local_repr = _addindent(local_repr, 2)\\n            lines.append(local_repr)\\n\\n        main_str += \\\"\\\\n  \\\" + \\\"\\\\n  \\\".join(lines) + \\\"\\\\n\\\"\\n        main_str += \\\")\\\"\\n        return main_str\\n\\n    @_copy_to_script_wrapper\\n    def __dir__(self):\\n        keys = super().__dir__()\\n        keys = [key for key in keys if not key.isdigit()]\\n        return keys\\n\\n    def insert(self, index: int, module: Module) -> None:\\n        r\\\"\\\"\\\"Insert a given module before a given index in the list.\\n\\n        Args:\\n            index (int): index to insert.\\n            module (nn.Module): module to insert\\n        \\\"\\\"\\\"\\n        for i in range(len(self._modules), index, -1):\\n            self._modules[str(i)] = self._modules[str(i - 1)]\\n        self._modules[str(index)] = module\\n\\n    def append(self, module: Module) -> \\\"ModuleList\\\":\\n        r\\\"\\\"\\\"Append a given module to the end of the list.\\n\\n        Args:\\n            module (nn.Module): module to append\\n        \\\"\\\"\\\"\\n        self.add_module(str(len(self)), module)\\n        return self\\n\\n    def pop(self, key: Union[int, slice]) -> Module:\\n        v = self[key]\\n        del self[key]\\n        return v\\n\\n    def extend(self, modules: Iterable[Module]) -> Self:\\n        r\\\"\\\"\\\"Append modules from a Python iterable to the end of the list.\\n\\n        Args:\\n            modules (iterable): iterable of modules to append\\n        \\\"\\\"\\\"\\n        if not isinstance(modules, container_abcs.Iterable):\\n            raise TypeError(\\n                \\\"ModuleList.extend should be called with an \\\"\\n                \\\"iterable, but got \\\" + type(modules).__name__\\n            )\\n        offset = len(self)\\n        for i, module in enumerate(modules):\\n            self.add_module(str(offset + i), module)\\n        return self\\n\\n    # remove forward alltogether to fallback on Module's _forward_unimplemented\\n\\n\\nclass ModuleDict(Module):\\n    r\\\"\\\"\\\"Holds submodules in a dictionary.\\n\\n    :class:`~torch.nn.ModuleDict` can be indexed like a regular Python dictionary,\\n    but modules it contains are properly registered, and will be visible by all\\n    :class:`~torch.nn.Module` methods.\\n\\n    :class:`~torch.nn.ModuleDict` is an **ordered** dictionary that respects\\n\\n    * the order of insertion, and\\n\\n    * in :meth:`~torch.nn.ModuleDict.update`, the order of the merged\\n      ``OrderedDict``, ``dict`` (started from Python 3.6) or another\\n      :class:`~torch.nn.ModuleDict` (the argument to\\n      :meth:`~torch.nn.ModuleDict.update`).\\n\\n    Note that :meth:`~torch.nn.ModuleDict.update` with other unordered mapping\\n    types (e.g., Python's plain ``dict`` before Python version 3.6) does not\\n    preserve the order of the merged mapping.\\n\\n    Args:\\n        modules (iterable, optional): a mapping (dictionary) of (string: module)\\n            or an iterable of key-value pairs of type (string, module)\\n\\n    Example::\\n\\n        class MyModule(nn.Module):\\n            def __init__(self) -> None:\\n                super().__init__()\\n                self.choices = nn.ModuleDict({\\n                        'conv': nn.Conv2d(10, 10, 3),\\n                        'pool': nn.MaxPool2d(3)\\n                })\\n                self.activations = nn.ModuleDict([\\n                        ['lrelu', nn.LeakyReLU()],\\n                        ['prelu', nn.PReLU()]\\n                ])\\n\\n            def forward(self, x, choice, act):\\n                x = self.choices[choice](x)\\n                x = self.activations[act](x)\\n                return x\\n    \\\"\\\"\\\"\\n\\n    _modules: Dict[str, Module]  # type: ignore[assignment]\\n\\n    def __init__(self, modules: Optional[Mapping[str, Module]] = None) -> None:\\n        super().__init__()\\n        if modules is not None:\\n            self.update(modules)\\n\\n    @_copy_to_script_wrapper\\n    def __getitem__(self, key: str) -> Module:\\n        return self._modules[key]\\n\\n    def __setitem__(self, key: str, module: Module) -> None:\\n        self.add_module(key, module)\\n\\n    def __delitem__(self, key: str) -> None:\\n        del self._modules[key]\\n\\n    @_copy_to_script_wrapper\\n    def __len__(self) -> int:\\n        return len(self._modules)\\n\\n    @_copy_to_script_wrapper\\n    def __iter__(self) -> Iterator[str]:\\n        return iter(self._modules)\\n\\n    @_copy_to_script_wrapper\\n    def __contains__(self, key: str) -> bool:\\n        return key in self._modules\\n\\n    def clear(self) -> None:\\n        \\\"\\\"\\\"Remove all items from the ModuleDict.\\\"\\\"\\\"\\n        self._modules.clear()\\n\\n    def pop(self, key: str) -> Module:\\n        r\\\"\\\"\\\"Remove key from the ModuleDict and return its module.\\n\\n        Args:\\n            key (str): key to pop from the ModuleDict\\n        \\\"\\\"\\\"\\n        v = self[key]\\n        del self[key]\\n        return v\\n\\n    @_copy_to_script_wrapper\\n    def keys(self) -> Iterable[str]:\\n        r\\\"\\\"\\\"Return an iterable of the ModuleDict keys.\\\"\\\"\\\"\\n        return self._modules.keys()\\n\\n    @_copy_to_script_wrapper\\n    def items(self) -> Iterable[Tuple[str, Module]]:\\n        r\\\"\\\"\\\"Return an iterable of the ModuleDict key/value pairs.\\\"\\\"\\\"\\n        return self._modules.items()\\n\\n    @_copy_to_script_wrapper\\n    def values(self) -> Iterable[Module]:\\n        r\\\"\\\"\\\"Return an iterable of the ModuleDict values.\\\"\\\"\\\"\\n        return self._modules.values()\\n\\n    def update(self, modules: Mapping[str, Module]) -> None:\\n        r\\\"\\\"\\\"Update the :class:`~torch.nn.ModuleDict` with key-value pairs from a mapping, overwriting existing keys.\\n\\n        .. note::\\n            If :attr:`modules` is an ``OrderedDict``, a :class:`~torch.nn.ModuleDict`, or\\n            an iterable of key-value pairs, the order of new elements in it is preserved.\\n\\n        Args:\\n            modules (iterable): a mapping (dictionary) from string to :class:`~torch.nn.Module`,\\n                or an iterable of key-value pairs of type (string, :class:`~torch.nn.Module`)\\n        \\\"\\\"\\\"\\n        if not isinstance(modules, container_abcs.Iterable):\\n            raise TypeError(\\n                \\\"ModuleDict.update should be called with an \\\"\\n                \\\"iterable of key/value pairs, but got \\\" + type(modules).__name__\\n            )\\n\\n        if isinstance(modules, (OrderedDict, ModuleDict, container_abcs.Mapping)):\\n            for key, module in modules.items():\\n                self[key] = module\\n        else:\\n            # modules here can be a list with two items\\n            for j, m in enumerate(modules):\\n                if not isinstance(m, container_abcs.Iterable):\\n                    raise TypeError(\\n                        \\\"ModuleDict update sequence element \\\"\\n                        \\\"#\\\" + str(j) + \\\" should be Iterable; is\\\" + type(m).__name__\\n                    )\\n                if not len(m) == 2:\\n                    raise ValueError(\\n                        \\\"ModuleDict update sequence element \\\"\\n                        \\\"#\\\" + str(j) + \\\" has length \\\" + str(len(m)) + \\\"; 2 is required\\\"\\n                    )\\n                # modules can be Mapping (what it's typed at), or a list: [(name1, module1), (name2, module2)]\\n                # that's too cumbersome to type correctly with overloads, so we add an ignore here\\n                self[m[0]] = m[1]  # type: ignore[assignment]\\n\\n    # remove forward alltogether to fallback on Module's _forward_unimplemented\\n\\n\\nclass ParameterList(Module):\\n    r\\\"\\\"\\\"Holds parameters in a list.\\n\\n    :class:`~torch.nn.ParameterList` can be used like a regular Python\\n    list, but Tensors that are :class:`~torch.nn.Parameter` are properly registered,\\n    and will be visible by all :class:`~torch.nn.Module` methods.\\n\\n    Note that the constructor, assigning an element of the list, the\\n    :meth:`~torch.nn.ParameterList.append` method and the :meth:`~torch.nn.ParameterList.extend`\\n    method will convert any :class:`~torch.Tensor` into :class:`~torch.nn.Parameter`.\\n\\n    Args:\\n        parameters (iterable, optional): an iterable of elements to add to the list.\\n\\n    Example::\\n\\n        class MyModule(nn.Module):\\n            def __init__(self) -> None:\\n                super().__init__()\\n                self.params = nn.ParameterList([nn.Parameter(torch.randn(10, 10)) for i in range(10)])\\n\\n            def forward(self, x):\\n                # ParameterList can act as an iterable, or be indexed using ints\\n                for i, p in enumerate(self.params):\\n                    x = self.params[i // 2].mm(x) + p.mm(x)\\n                return x\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, values: Optional[Iterable[Any]] = None) -> None:\\n        super().__init__()\\n        self._size = 0\\n        if values is not None:\\n            self += values\\n\\n    def _get_abs_string_index(self, idx):\\n        \\\"\\\"\\\"Get the absolute index for the list of modules.\\\"\\\"\\\"\\n        idx = operator.index(idx)\\n        if not (-len(self) <= idx < len(self)):\\n            raise IndexError(f\\\"index {idx} is out of range\\\")\\n        if idx < 0:\\n            idx += len(self)\\n        return str(idx)\\n\\n    @overload\\n    def __getitem__(self, idx: int) -> Any:\\n        ...\\n\\n    @overload\\n    def __getitem__(self: T, idx: slice) -> T:\\n        ...\\n\\n    def __getitem__(self, idx):\\n        if isinstance(idx, slice):\\n            start, stop, step = idx.indices(len(self))\\n            out = self.__class__()\\n            for i in range(start, stop, step):\\n                out.append(self[i])\\n            return out\\n        else:\\n            idx = self._get_abs_string_index(idx)\\n            return getattr(self, str(idx))\\n\\n    def __setitem__(self, idx: int, param: Any) -> None:\\n        # Note that all other function that add an entry to the list part of\\n        # the ParameterList end up here. So this is the only place where we need\\n        # to wrap things into Parameter if needed.\\n        # Objects added via setattr() are not in the list part and thus won't\\n        # call into this function.\\n        idx = self._get_abs_string_index(idx)\\n        if isinstance(param, torch.Tensor) and not isinstance(param, Parameter):\\n            param = Parameter(param)\\n        return setattr(self, str(idx), param)\\n\\n    def __len__(self) -> int:\\n        return self._size\\n\\n    def __iter__(self) -> Iterator[Any]:\\n        return iter(self[i] for i in range(len(self)))\\n\\n    def __iadd__(self, parameters: Iterable[Any]) -> Self:\\n        return self.extend(parameters)\\n\\n    def __dir__(self):\\n        keys = super().__dir__()\\n        keys = [key for key in keys if not key.isdigit()]\\n        return keys\\n\\n    def append(self, value: Any) -> \\\"ParameterList\\\":\\n        \\\"\\\"\\\"Append a given value at the end of the list.\\n\\n        Args:\\n            value (Any): value to append\\n        \\\"\\\"\\\"\\n        new_idx = len(self)\\n        self._size += 1\\n        self[new_idx] = value\\n        return self\\n\\n    def extend(self, values: Iterable[Any]) -> Self:\\n        \\\"\\\"\\\"Append values from a Python iterable to the end of the list.\\n\\n        Args:\\n            values (iterable): iterable of values to append\\n        \\\"\\\"\\\"\\n        # Tensor is an iterable but we never want to unpack it here\\n        if not isinstance(values, container_abcs.Iterable) or isinstance(\\n            values, torch.Tensor\\n        ):\\n            raise TypeError(\\n                \\\"ParameterList.extend should be called with an \\\"\\n                \\\"iterable, but got \\\" + type(values).__name__\\n            )\\n        for value in values:\\n            self.append(value)\\n        return self\\n\\n    def extra_repr(self) -> str:\\n        child_lines = []\\n        for k, p in enumerate(self):\\n            if isinstance(p, torch.Tensor):\\n                size_str = \\\"x\\\".join(str(size) for size in p.size())\\n                if p.device.type in [\\\"cuda\\\", torch._C._get_privateuse1_backend_name()]:\\n                    device_str = f\\\" ({p.device})\\\"\\n                else:\\n                    device_str = \\\"\\\"\\n                parastr = \\\"{} containing: [{} of size {}{}]\\\".format(\\n                    \\\"Parameter\\\" if isinstance(p, Parameter) else \\\"Tensor\\\",\\n                    p.dtype,\\n                    size_str,\\n                    device_str,\\n                )\\n                child_lines.append(\\\"  (\\\" + str(k) + \\\"): \\\" + parastr)\\n            else:\\n                child_lines.append(\\n                    \\\"  (\\\" + str(k) + \\\"): Object of type: \\\" + type(p).__name__\\n                )\\n\\n        tmpstr = \\\"\\\\n\\\".join(child_lines)\\n        return tmpstr\\n\\n    def __call__(self, *args, **kwargs):\\n        raise RuntimeError(\\\"ParameterList should not be called.\\\")\\n\\n\\nclass ParameterDict(Module):\\n    r\\\"\\\"\\\"Holds parameters in a dictionary.\\n\\n    ParameterDict can be indexed like a regular Python dictionary, but Parameters it\\n    contains are properly registered, and will be visible by all Module methods.\\n    Other objects are treated as would be done by a regular Python dictionary\\n\\n    :class:`~torch.nn.ParameterDict` is an **ordered** dictionary.\\n    :meth:`~torch.nn.ParameterDict.update` with other unordered mapping\\n    types (e.g., Python's plain ``dict``) does not preserve the order of the\\n    merged mapping. On the other hand, ``OrderedDict`` or another :class:`~torch.nn.ParameterDict`\\n    will preserve their ordering.\\n\\n    Note that the constructor, assigning an element of the dictionary and the\\n    :meth:`~torch.nn.ParameterDict.update` method will convert any :class:`~torch.Tensor` into\\n    :class:`~torch.nn.Parameter`.\\n\\n    Args:\\n        values (iterable, optional): a mapping (dictionary) of\\n            (string : Any) or an iterable of key-value pairs\\n            of type (string, Any)\\n\\n    Example::\\n\\n        class MyModule(nn.Module):\\n            def __init__(self) -> None:\\n                super().__init__()\\n                self.params = nn.ParameterDict({\\n                        'left': nn.Parameter(torch.randn(5, 10)),\\n                        'right': nn.Parameter(torch.randn(5, 10))\\n                })\\n\\n            def forward(self, x, choice):\\n                x = self.params[choice].mm(x)\\n                return x\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, parameters: Any = None) -> None:\\n        super().__init__()\\n        self._keys: Dict[str, None] = {}\\n        if parameters is not None:\\n            self.update(parameters)\\n\\n    def _key_to_attr(self, key: str) -> str:\\n        if not isinstance(key, str):\\n            raise TypeError(\\n                \\\"Index given to ParameterDict cannot be used as a key as it is \\\"\\n                f\\\"not a string (type is '{type(key).__name__}'). Open an issue on \\\"\\n                \\\"github if you need non-string keys.\\\"\\n            )\\n        else:\\n            # Use the key as-is so that `.named_parameters()` returns the right thing\\n            return key\\n\\n    def __getitem__(self, key: str) -> Any:\\n        attr = self._key_to_attr(key)\\n        return getattr(self, attr)\\n\\n    def __setitem__(self, key: str, value: Any) -> None:\\n        # Note that all other function that add an entry to the dictionary part of\\n        # the ParameterDict end up here. So this is the only place where we need\\n        # to wrap things into Parameter if needed.\\n        # Objects added via setattr() are not in the dictionary part and thus won't\\n        # call into this function.\\n        self._keys[key] = None\\n        attr = self._key_to_attr(key)\\n        if isinstance(value, torch.Tensor) and not isinstance(value, Parameter):\\n            value = Parameter(value)\\n        setattr(self, attr, value)\\n\\n    def __delitem__(self, key: str) -> None:\\n        del self._keys[key]\\n        attr = self._key_to_attr(key)\\n        delattr(self, attr)\\n\\n    def __len__(self) -> int:\\n        return len(self._keys)\\n\\n    def __iter__(self) -> Iterator[str]:\\n        return iter(self._keys)\\n\\n    def __reversed__(self) -> Iterator[str]:\\n        return reversed(list(self._keys))\\n\\n    def copy(self) -> \\\"ParameterDict\\\":\\n        \\\"\\\"\\\"Return a copy of this :class:`~torch.nn.ParameterDict` instance.\\\"\\\"\\\"\\n        # We have to use an OrderedDict because the ParameterDict constructor\\n        # behaves differently on plain dict vs OrderedDict\\n        return ParameterDict(OrderedDict((k, self[k]) for k in self._keys))\\n\\n    def __contains__(self, key: str) -> bool:\\n        return key in self._keys\\n\\n    def setdefault(self, key: str, default: Optional[Any] = None) -> Any:\\n        \\\"\\\"\\\"Set the default for a key in the Parameterdict.\\n\\n        If key is in the ParameterDict, return its value.\\n        If not, insert `key` with a parameter `default` and return `default`.\\n        `default` defaults to `None`.\\n\\n        Args:\\n            key (str): key to set default for\\n            default (Any): the parameter set to the key\\n        \\\"\\\"\\\"\\n        if key not in self:\\n            self[key] = default\\n        return self[key]\\n\\n    def clear(self) -> None:\\n        \\\"\\\"\\\"Remove all items from the ParameterDict.\\\"\\\"\\\"\\n        for k in self._keys.copy():\\n            del self[k]\\n\\n    def pop(self, key: str) -> Any:\\n        r\\\"\\\"\\\"Remove key from the ParameterDict and return its parameter.\\n\\n        Args:\\n            key (str): key to pop from the ParameterDict\\n        \\\"\\\"\\\"\\n        v = self[key]\\n        del self[key]\\n        return v\\n\\n    def popitem(self) -> Tuple[str, Any]:\\n        \\\"\\\"\\\"Remove and return the last inserted `(key, parameter)` pair from the ParameterDict.\\\"\\\"\\\"\\n        k, _ = self._keys.popitem()\\n        # We need the key in the _keys to be able to access/del\\n        self._keys[k] = None\\n        val = self[k]\\n        del self[k]\\n        return k, val\\n\\n    def get(self, key: str, default: Optional[Any] = None) -> Any:\\n        r\\\"\\\"\\\"Return the parameter associated with key if present. Otherwise return default if provided, None if not.\\n\\n        Args:\\n            key (str): key to get from the ParameterDict\\n            default (Parameter, optional): value to return if key not present\\n        \\\"\\\"\\\"\\n        return self[key] if key in self else default\\n\\n    def fromkeys(\\n        self, keys: Iterable[str], default: Optional[Any] = None\\n    ) -> \\\"ParameterDict\\\":\\n        r\\\"\\\"\\\"Return a new ParameterDict with the keys provided.\\n\\n        Args:\\n            keys (iterable, string): keys to make the new ParameterDict from\\n            default (Parameter, optional): value to set for all keys\\n        \\\"\\\"\\\"\\n        return ParameterDict((k, default) for k in keys)\\n\\n    def keys(self) -> Iterable[str]:\\n        r\\\"\\\"\\\"Return an iterable of the ParameterDict keys.\\\"\\\"\\\"\\n        return self._keys.keys()\\n\\n    def items(self) -> Iterable[Tuple[str, Any]]:\\n        r\\\"\\\"\\\"Return an iterable of the ParameterDict key/value pairs.\\\"\\\"\\\"\\n        return ((k, self[k]) for k in self._keys)\\n\\n    def values(self) -> Iterable[Any]:\\n        r\\\"\\\"\\\"Return an iterable of the ParameterDict values.\\\"\\\"\\\"\\n        return (self[k] for k in self._keys)\\n\\n    def update(self, parameters: Union[Mapping[str, Any], \\\"ParameterDict\\\"]) -> None:\\n        r\\\"\\\"\\\"Update the :class:`~torch.nn.ParameterDict` with key-value pairs from ``parameters``, overwriting existing keys.\\n\\n        .. note::\\n            If :attr:`parameters` is an ``OrderedDict``, a :class:`~torch.nn.ParameterDict`, or\\n            an iterable of key-value pairs, the order of new elements in it is preserved.\\n\\n        Args:\\n            parameters (iterable): a mapping (dictionary) from string to\\n                :class:`~torch.nn.Parameter`, or an iterable of\\n                key-value pairs of type (string, :class:`~torch.nn.Parameter`)\\n        \\\"\\\"\\\"\\n        if not isinstance(parameters, container_abcs.Iterable):\\n            raise TypeError(\\n                \\\"ParametersDict.update should be called with an \\\"\\n                \\\"iterable of key/value pairs, but got \\\" + type(parameters).__name__\\n            )\\n\\n        if isinstance(parameters, (OrderedDict, ParameterDict)):\\n            for key, parameter in parameters.items():\\n                self[key] = parameter\\n        elif isinstance(parameters, container_abcs.Mapping):\\n            for key, parameter in sorted(parameters.items()):\\n                self[key] = parameter\\n        else:\\n            for j, p in enumerate(parameters):\\n                if not isinstance(p, container_abcs.Iterable):\\n                    raise TypeError(\\n                        \\\"ParameterDict update sequence element \\\"\\n                        \\\"#\\\" + str(j) + \\\" should be Iterable; is\\\" + type(p).__name__\\n                    )\\n                if not len(p) == 2:\\n                    raise ValueError(\\n                        \\\"ParameterDict update sequence element \\\"\\n                        \\\"#\\\" + str(j) + \\\" has length \\\" + str(len(p)) + \\\"; 2 is required\\\"\\n                    )\\n                # parameters as length-2 list too cumbersome to type, see ModuleDict.update comment\\n                self[p[0]] = p[1]  # type: ignore[assignment]\\n\\n    def extra_repr(self) -> str:\\n        child_lines = []\\n        for k, p in self.items():\\n            if isinstance(p, torch.Tensor):\\n                size_str = \\\"x\\\".join(str(size) for size in p.size())\\n                if p.device.type in [\\\"cuda\\\", torch._C._get_privateuse1_backend_name()]:\\n                    device_str = f\\\" ({p.device})\\\"\\n                else:\\n                    device_str = \\\"\\\"\\n                parastr = \\\"{} containing: [{} of size {}{}]\\\".format(\\n                    \\\"Parameter\\\" if isinstance(p, Parameter) else \\\"Tensor\\\",\\n                    torch.typename(p),\\n                    size_str,\\n                    device_str,\\n                )\\n                child_lines.append(\\\"  (\\\" + str(k) + \\\"): \\\" + parastr)\\n            else:\\n                child_lines.append(\\n                    \\\"  (\\\" + str(k) + \\\"): Object of type: \\\" + type(p).__name__\\n                )\\n        tmpstr = \\\"\\\\n\\\".join(child_lines)\\n        return tmpstr\\n\\n    def __call__(self, input):\\n        raise RuntimeError(\\\"ParameterDict should not be called.\\\")\\n\\n    def __or__(self, other: \\\"ParameterDict\\\") -> \\\"ParameterDict\\\":\\n        copy = self.copy()\\n        copy.update(other)\\n        return copy\\n\\n    def __ror__(self, other: \\\"ParameterDict\\\") -> \\\"ParameterDict\\\":\\n        copy = other.copy()\\n        copy.update(self)\\n        return copy\\n\\n    def __ior__(self, other: \\\"ParameterDict\\\") -> Self:\\n        self.update(other)\\n        return self\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Optional\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.common_types import _ratio_2_t, _ratio_any_t, _size_2_t, _size_any_t\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"Upsample\\\", \\\"UpsamplingNearest2d\\\", \\\"UpsamplingBilinear2d\\\"]\\n\\n\\nclass Upsample(Module):\\n    r\\\"\\\"\\\"Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data.\\n\\n    The input data is assumed to be of the form\\n    `minibatch x channels x [optional depth] x [optional height] x width`.\\n    Hence, for spatial inputs, we expect a 4D Tensor and for volumetric inputs, we expect a 5D Tensor.\\n\\n    The algorithms available for upsampling are nearest neighbor and linear,\\n    bilinear, bicubic and trilinear for 3D, 4D and 5D input Tensor,\\n    respectively.\\n\\n    One can either give a :attr:`scale_factor` or the target output :attr:`size` to\\n    calculate the output size. (You cannot give both, as it is ambiguous)\\n\\n    Args:\\n        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], optional):\\n            output spatial sizes\\n        scale_factor (float or Tuple[float] or Tuple[float, float] or Tuple[float, float, float], optional):\\n            multiplier for spatial size. Has to match input size if it is a tuple.\\n        mode (str, optional): the upsampling algorithm: one of ``'nearest'``,\\n            ``'linear'``, ``'bilinear'``, ``'bicubic'`` and ``'trilinear'``.\\n            Default: ``'nearest'``\\n        align_corners (bool, optional): if ``True``, the corner pixels of the input\\n            and output tensors are aligned, and thus preserving the values at\\n            those pixels. This only has effect when :attr:`mode` is\\n            ``'linear'``, ``'bilinear'``, ``'bicubic'``, or ``'trilinear'``.\\n            Default: ``False``\\n        recompute_scale_factor (bool, optional): recompute the scale_factor for use in the\\n            interpolation calculation. If `recompute_scale_factor` is ``True``, then\\n            `scale_factor` must be passed in and `scale_factor` is used to compute the\\n            output `size`. The computed output `size` will be used to infer new scales for\\n            the interpolation. Note that when `scale_factor` is floating-point, it may differ\\n            from the recomputed `scale_factor` due to rounding and precision issues.\\n            If `recompute_scale_factor` is ``False``, then `size` or `scale_factor` will\\n            be used directly for interpolation.\\n\\n    Shape:\\n        - Input: :math:`(N, C, W_{in})`, :math:`(N, C, H_{in}, W_{in})` or :math:`(N, C, D_{in}, H_{in}, W_{in})`\\n        - Output: :math:`(N, C, W_{out})`, :math:`(N, C, H_{out}, W_{out})`\\n          or :math:`(N, C, D_{out}, H_{out}, W_{out})`, where\\n\\n    .. math::\\n        D_{out} = \\\\left\\\\lfloor D_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    .. math::\\n        H_{out} = \\\\left\\\\lfloor H_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    .. math::\\n        W_{out} = \\\\left\\\\lfloor W_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    .. warning::\\n        With ``align_corners = True``, the linearly interpolating modes\\n        (`linear`, `bilinear`, `bicubic`, and `trilinear`) don't proportionally\\n        align the output and input pixels, and thus the output values can depend\\n        on the input size. This was the default behavior for these modes up to\\n        version 0.3.1. Since then, the default behavior is\\n        ``align_corners = False``. See below for concrete examples on how this\\n        affects the outputs.\\n\\n    .. note::\\n        If you want downsampling/general resizing, you should use :func:`~nn.functional.interpolate`.\\n\\n    Examples::\\n\\n        >>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)\\n        >>> input\\n        tensor([[[[1., 2.],\\n                  [3., 4.]]]])\\n\\n        >>> m = nn.Upsample(scale_factor=2, mode='nearest')\\n        >>> m(input)\\n        tensor([[[[1., 1., 2., 2.],\\n                  [1., 1., 2., 2.],\\n                  [3., 3., 4., 4.],\\n                  [3., 3., 4., 4.]]]])\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"other tests seem to modify printing styles\\\")\\n        >>> m = nn.Upsample(scale_factor=2, mode='bilinear')  # align_corners=False\\n        >>> m(input)\\n        tensor([[[[1.0000, 1.2500, 1.7500, 2.0000],\\n                  [1.5000, 1.7500, 2.2500, 2.5000],\\n                  [2.5000, 2.7500, 3.2500, 3.5000],\\n                  [3.0000, 3.2500, 3.7500, 4.0000]]]])\\n\\n        >>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\\n        >>> m(input)\\n        tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],\\n                  [1.6667, 2.0000, 2.3333, 2.6667],\\n                  [2.3333, 2.6667, 3.0000, 3.3333],\\n                  [3.0000, 3.3333, 3.6667, 4.0000]]]])\\n\\n        >>> # Try scaling the same data in a larger tensor\\n        >>> input_3x3 = torch.zeros(3, 3).view(1, 1, 3, 3)\\n        >>> input_3x3[:, :, :2, :2].copy_(input)\\n        tensor([[[[1., 2.],\\n                  [3., 4.]]]])\\n        >>> input_3x3\\n        tensor([[[[1., 2., 0.],\\n                  [3., 4., 0.],\\n                  [0., 0., 0.]]]])\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"seems to fail when other tests are run in the same session\\\")\\n        >>> m = nn.Upsample(scale_factor=2, mode='bilinear')  # align_corners=False\\n        >>> # Notice that values in top left corner are the same with the small input (except at boundary)\\n        >>> m(input_3x3)\\n        tensor([[[[1.0000, 1.2500, 1.7500, 1.5000, 0.5000, 0.0000],\\n                  [1.5000, 1.7500, 2.2500, 1.8750, 0.6250, 0.0000],\\n                  [2.5000, 2.7500, 3.2500, 2.6250, 0.8750, 0.0000],\\n                  [2.2500, 2.4375, 2.8125, 2.2500, 0.7500, 0.0000],\\n                  [0.7500, 0.8125, 0.9375, 0.7500, 0.2500, 0.0000],\\n                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])\\n\\n        >>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\\n        >>> # Notice that values in top left corner are now changed\\n        >>> m(input_3x3)\\n        tensor([[[[1.0000, 1.4000, 1.8000, 1.6000, 0.8000, 0.0000],\\n                  [1.8000, 2.2000, 2.6000, 2.2400, 1.1200, 0.0000],\\n                  [2.6000, 3.0000, 3.4000, 2.8800, 1.4400, 0.0000],\\n                  [2.4000, 2.7200, 3.0400, 2.5600, 1.2800, 0.0000],\\n                  [1.2000, 1.3600, 1.5200, 1.2800, 0.6400, 0.0000],\\n                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\n        \\\"size\\\",\\n        \\\"scale_factor\\\",\\n        \\\"mode\\\",\\n        \\\"align_corners\\\",\\n        \\\"name\\\",\\n        \\\"recompute_scale_factor\\\",\\n    ]\\n    name: str\\n    size: Optional[_size_any_t]\\n    scale_factor: Optional[_ratio_any_t]\\n    mode: str\\n    align_corners: Optional[bool]\\n    recompute_scale_factor: Optional[bool]\\n\\n    def __init__(\\n        self,\\n        size: Optional[_size_any_t] = None,\\n        scale_factor: Optional[_ratio_any_t] = None,\\n        mode: str = \\\"nearest\\\",\\n        align_corners: Optional[bool] = None,\\n        recompute_scale_factor: Optional[bool] = None,\\n    ) -> None:\\n        super().__init__()\\n        self.name = type(self).__name__\\n        self.size = size\\n        if isinstance(scale_factor, tuple):\\n            self.scale_factor = tuple(float(factor) for factor in scale_factor)\\n        else:\\n            self.scale_factor = float(scale_factor) if scale_factor else None\\n        self.mode = mode\\n        self.align_corners = align_corners\\n        self.recompute_scale_factor = recompute_scale_factor\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.interpolate(\\n            input,\\n            self.size,\\n            self.scale_factor,\\n            self.mode,\\n            self.align_corners,\\n            recompute_scale_factor=self.recompute_scale_factor,\\n        )\\n\\n    def __setstate__(self, state):\\n        if \\\"recompute_scale_factor\\\" not in state:\\n            state[\\\"recompute_scale_factor\\\"] = True\\n\\n        super().__setstate__(state)\\n\\n    def extra_repr(self) -> str:\\n        if self.scale_factor is not None:\\n            info = \\\"scale_factor=\\\" + repr(self.scale_factor)\\n        else:\\n            info = \\\"size=\\\" + repr(self.size)\\n        info += \\\", mode=\\\" + repr(self.mode)\\n        return info\\n\\n\\nclass UpsamplingNearest2d(Upsample):\\n    r\\\"\\\"\\\"Applies a 2D nearest neighbor upsampling to an input signal composed of several input channels.\\n\\n    To specify the scale, it takes either the :attr:`size` or the :attr:`scale_factor`\\n    as it's constructor argument.\\n\\n    When :attr:`size` is given, it is the output size of the image `(h, w)`.\\n\\n    Args:\\n        size (int or Tuple[int, int], optional): output spatial sizes\\n        scale_factor (float or Tuple[float, float], optional): multiplier for\\n            spatial size.\\n\\n    .. warning::\\n        This class is deprecated in favor of :func:`~nn.functional.interpolate`.\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})`\\n        - Output: :math:`(N, C, H_{out}, W_{out})` where\\n\\n    .. math::\\n          H_{out} = \\\\left\\\\lfloor H_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    .. math::\\n          W_{out} = \\\\left\\\\lfloor W_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)\\n        >>> input\\n        tensor([[[[1., 2.],\\n                  [3., 4.]]]])\\n\\n        >>> m = nn.UpsamplingNearest2d(scale_factor=2)\\n        >>> m(input)\\n        tensor([[[[1., 1., 2., 2.],\\n                  [1., 1., 2., 2.],\\n                  [3., 3., 4., 4.],\\n                  [3., 3., 4., 4.]]]])\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        size: Optional[_size_2_t] = None,\\n        scale_factor: Optional[_ratio_2_t] = None,\\n    ) -> None:\\n        super().__init__(size, scale_factor, mode=\\\"nearest\\\")\\n\\n\\nclass UpsamplingBilinear2d(Upsample):\\n    r\\\"\\\"\\\"Applies a 2D bilinear upsampling to an input signal composed of several input channels.\\n\\n    To specify the scale, it takes either the :attr:`size` or the :attr:`scale_factor`\\n    as it's constructor argument.\\n\\n    When :attr:`size` is given, it is the output size of the image `(h, w)`.\\n\\n    Args:\\n        size (int or Tuple[int, int], optional): output spatial sizes\\n        scale_factor (float or Tuple[float, float], optional): multiplier for\\n            spatial size.\\n\\n    .. warning::\\n        This class is deprecated in favor of :func:`~nn.functional.interpolate`. It is\\n        equivalent to ``nn.functional.interpolate(..., mode='bilinear', align_corners=True)``.\\n\\n    Shape:\\n        - Input: :math:`(N, C, H_{in}, W_{in})`\\n        - Output: :math:`(N, C, H_{out}, W_{out})` where\\n\\n    .. math::\\n        H_{out} = \\\\left\\\\lfloor H_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    .. math::\\n        W_{out} = \\\\left\\\\lfloor W_{in} \\\\times \\\\text{scale\\\\_factor} \\\\right\\\\rfloor\\n\\n    Examples::\\n\\n        >>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)\\n        >>> input\\n        tensor([[[[1., 2.],\\n                  [3., 4.]]]])\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"do other tests modify the global state?\\\")\\n        >>> m = nn.UpsamplingBilinear2d(scale_factor=2)\\n        >>> m(input)\\n        tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],\\n                  [1.6667, 2.0000, 2.3333, 2.6667],\\n                  [2.3333, 2.6667, 3.0000, 3.3333],\\n                  [3.0000, 3.3333, 3.6667, 4.0000]]]])\\n    \\\"\\\"\\\"\\n\\n    def __init__(\\n        self,\\n        size: Optional[_size_2_t] = None,\\n        scale_factor: Optional[_ratio_2_t] = None,\\n    ) -> None:\\n        super().__init__(size, scale_factor, mode=\\\"bilinear\\\", align_corners=True)\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Optional\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch.nn import functional as F, init\\nfrom torch.nn.parameter import Parameter\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"Embedding\\\", \\\"EmbeddingBag\\\"]\\n\\n\\nclass Embedding(Module):\\n    r\\\"\\\"\\\"A simple lookup table that stores embeddings of a fixed dictionary and size.\\n\\n    This module is often used to store word embeddings and retrieve them using indices.\\n    The input to the module is a list of indices, and the output is the corresponding\\n    word embeddings.\\n\\n    Args:\\n        num_embeddings (int): size of the dictionary of embeddings\\n        embedding_dim (int): the size of each embedding vector\\n        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;\\n                                     therefore, the embedding vector at :attr:`padding_idx` is not updated during training,\\n                                     i.e. it remains as a fixed \\\"pad\\\". For a newly constructed Embedding,\\n                                     the embedding vector at :attr:`padding_idx` will default to all zeros,\\n                                     but can be updated to another value to be used as the padding vector.\\n        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`\\n                                    is renormalized to have norm :attr:`max_norm`.\\n        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.\\n        scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse of frequency of\\n                                                the words in the mini-batch. Default ``False``.\\n        sparse (bool, optional): If ``True``, gradient w.r.t. :attr:`weight` matrix will be a sparse tensor.\\n                                 See Notes for more details regarding sparse gradients.\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape (num_embeddings, embedding_dim)\\n                         initialized from :math:`\\\\mathcal{N}(0, 1)`\\n\\n    Shape:\\n        - Input: :math:`(*)`, IntTensor or LongTensor of arbitrary shape containing the indices to extract\\n        - Output: :math:`(*, H)`, where `*` is the input shape and :math:`H=\\\\text{embedding\\\\_dim}`\\n\\n    .. note::\\n        Keep in mind that only a limited number of optimizers support\\n        sparse gradients: currently it's :class:`optim.SGD` (`CUDA` and `CPU`),\\n        :class:`optim.SparseAdam` (`CUDA` and `CPU`) and :class:`optim.Adagrad` (`CPU`)\\n\\n    .. note::\\n        When :attr:`max_norm` is not ``None``, :class:`Embedding`'s forward method will modify the\\n        :attr:`weight` tensor in-place. Since tensors needed for gradient computations cannot be\\n        modified in-place, performing a differentiable operation on ``Embedding.weight`` before\\n        calling :class:`Embedding`'s forward method requires cloning ``Embedding.weight`` when\\n        :attr:`max_norm` is not ``None``. For example::\\n\\n            n, d, m = 3, 5, 7\\n            embedding = nn.Embedding(n, d, max_norm=1.0)\\n            W = torch.randn((m, d), requires_grad=True)\\n            idx = torch.tensor([1, 2])\\n            a = embedding.weight.clone() @ W.t()  # weight must be cloned for this to be differentiable\\n            b = embedding(idx) @ W.t()  # modifies weight in-place\\n            out = (a.unsqueeze(0) + b.unsqueeze(1))\\n            loss = out.sigmoid().prod()\\n            loss.backward()\\n\\n    Examples::\\n\\n        >>> # an Embedding module containing 10 tensors of size 3\\n        >>> embedding = nn.Embedding(10, 3)\\n        >>> # a batch of 2 samples of 4 indices each\\n        >>> input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> embedding(input)\\n        tensor([[[-0.0251, -1.6902,  0.7172],\\n                 [-0.6431,  0.0748,  0.6969],\\n                 [ 1.4970,  1.3448, -0.9685],\\n                 [-0.3677, -2.7265, -0.1685]],\\n\\n                [[ 1.4970,  1.3448, -0.9685],\\n                 [ 0.4362, -0.4004,  0.9400],\\n                 [-0.6431,  0.0748,  0.6969],\\n                 [ 0.9124, -2.3616,  1.1151]]])\\n\\n\\n        >>> # example with padding_idx\\n        >>> embedding = nn.Embedding(10, 3, padding_idx=0)\\n        >>> input = torch.LongTensor([[0, 2, 0, 5]])\\n        >>> embedding(input)\\n        tensor([[[ 0.0000,  0.0000,  0.0000],\\n                 [ 0.1535, -2.0309,  0.9315],\\n                 [ 0.0000,  0.0000,  0.0000],\\n                 [-0.1655,  0.9897,  0.0635]]])\\n\\n        >>> # example of changing `pad` vector\\n        >>> padding_idx = 0\\n        >>> embedding = nn.Embedding(3, 3, padding_idx=padding_idx)\\n        >>> embedding.weight\\n        Parameter containing:\\n        tensor([[ 0.0000,  0.0000,  0.0000],\\n                [-0.7895, -0.7089, -0.0364],\\n                [ 0.6778,  0.5803,  0.2678]], requires_grad=True)\\n        >>> with torch.no_grad():\\n        ...     embedding.weight[padding_idx] = torch.ones(3)\\n        >>> embedding.weight\\n        Parameter containing:\\n        tensor([[ 1.0000,  1.0000,  1.0000],\\n                [-0.7895, -0.7089, -0.0364],\\n                [ 0.6778,  0.5803,  0.2678]], requires_grad=True)\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\n        \\\"num_embeddings\\\",\\n        \\\"embedding_dim\\\",\\n        \\\"padding_idx\\\",\\n        \\\"max_norm\\\",\\n        \\\"norm_type\\\",\\n        \\\"scale_grad_by_freq\\\",\\n        \\\"sparse\\\",\\n    ]\\n\\n    num_embeddings: int\\n    embedding_dim: int\\n    padding_idx: Optional[int]\\n    max_norm: Optional[float]\\n    norm_type: float\\n    scale_grad_by_freq: bool\\n    weight: Tensor\\n    freeze: bool\\n    sparse: bool\\n\\n    def __init__(\\n        self,\\n        num_embeddings: int,\\n        embedding_dim: int,\\n        padding_idx: Optional[int] = None,\\n        max_norm: Optional[float] = None,\\n        norm_type: float = 2.0,\\n        scale_grad_by_freq: bool = False,\\n        sparse: bool = False,\\n        _weight: Optional[Tensor] = None,\\n        _freeze: bool = False,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.num_embeddings = num_embeddings\\n        self.embedding_dim = embedding_dim\\n        if padding_idx is not None:\\n            if padding_idx > 0:\\n                assert (\\n                    padding_idx < self.num_embeddings\\n                ), \\\"Padding_idx must be within num_embeddings\\\"\\n            elif padding_idx < 0:\\n                assert (\\n                    padding_idx >= -self.num_embeddings\\n                ), \\\"Padding_idx must be within num_embeddings\\\"\\n                padding_idx = self.num_embeddings + padding_idx\\n        self.padding_idx = padding_idx\\n        self.max_norm = max_norm\\n        self.norm_type = norm_type\\n        self.scale_grad_by_freq = scale_grad_by_freq\\n        if _weight is None:\\n            self.weight = Parameter(\\n                torch.empty((num_embeddings, embedding_dim), **factory_kwargs),\\n                requires_grad=not _freeze,\\n            )\\n            self.reset_parameters()\\n        else:\\n            assert list(_weight.shape) == [\\n                num_embeddings,\\n                embedding_dim,\\n            ], \\\"Shape of weight does not match num_embeddings and embedding_dim\\\"\\n            self.weight = Parameter(_weight, requires_grad=not _freeze)\\n\\n        self.sparse = sparse\\n\\n    def reset_parameters(self) -> None:\\n        init.normal_(self.weight)\\n        self._fill_padding_idx_with_zero()\\n\\n    def _fill_padding_idx_with_zero(self) -> None:\\n        if self.padding_idx is not None:\\n            with torch.no_grad():\\n                self.weight[self.padding_idx].fill_(0)\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.embedding(\\n            input,\\n            self.weight,\\n            self.padding_idx,\\n            self.max_norm,\\n            self.norm_type,\\n            self.scale_grad_by_freq,\\n            self.sparse,\\n        )\\n\\n    def extra_repr(self) -> str:\\n        s = \\\"{num_embeddings}, {embedding_dim}\\\"\\n        if self.padding_idx is not None:\\n            s += \\\", padding_idx={padding_idx}\\\"\\n        if self.max_norm is not None:\\n            s += \\\", max_norm={max_norm}\\\"\\n        if self.norm_type != 2:\\n            s += \\\", norm_type={norm_type}\\\"\\n        if self.scale_grad_by_freq is not False:\\n            s += \\\", scale_grad_by_freq={scale_grad_by_freq}\\\"\\n        if self.sparse is not False:\\n            s += \\\", sparse=True\\\"\\n        return s.format(**self.__dict__)\\n\\n    @classmethod\\n    def from_pretrained(\\n        cls,\\n        embeddings,\\n        freeze=True,\\n        padding_idx=None,\\n        max_norm=None,\\n        norm_type=2.0,\\n        scale_grad_by_freq=False,\\n        sparse=False,\\n    ):\\n        r\\\"\\\"\\\"Create Embedding instance from given 2-dimensional FloatTensor.\\n\\n        Args:\\n            embeddings (Tensor): FloatTensor containing weights for the Embedding.\\n                First dimension is being passed to Embedding as ``num_embeddings``, second as ``embedding_dim``.\\n            freeze (bool, optional): If ``True``, the tensor does not get updated in the learning process.\\n                Equivalent to ``embedding.weight.requires_grad = False``. Default: ``True``\\n            padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;\\n                                         therefore, the embedding vector at :attr:`padding_idx` is not updated during training,\\n                                         i.e. it remains as a fixed \\\"pad\\\".\\n            max_norm (float, optional): See module initialization documentation.\\n            norm_type (float, optional): See module initialization documentation. Default ``2``.\\n            scale_grad_by_freq (bool, optional): See module initialization documentation. Default ``False``.\\n            sparse (bool, optional): See module initialization documentation.\\n\\n        Examples::\\n\\n            >>> # FloatTensor containing pretrained weights\\n            >>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])\\n            >>> embedding = nn.Embedding.from_pretrained(weight)\\n            >>> # Get embeddings for index 1\\n            >>> input = torch.LongTensor([1])\\n            >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n            >>> embedding(input)\\n            tensor([[ 4.0000,  5.1000,  6.3000]])\\n        \\\"\\\"\\\"\\n        assert (\\n            embeddings.dim() == 2\\n        ), \\\"Embeddings parameter is expected to be 2-dimensional\\\"\\n        rows, cols = embeddings.shape\\n        embedding = cls(\\n            num_embeddings=rows,\\n            embedding_dim=cols,\\n            _weight=embeddings,\\n            _freeze=freeze,\\n            padding_idx=padding_idx,\\n            max_norm=max_norm,\\n            norm_type=norm_type,\\n            scale_grad_by_freq=scale_grad_by_freq,\\n            sparse=sparse,\\n        )\\n        return embedding\\n\\n\\nclass EmbeddingBag(Module):\\n    r\\\"\\\"\\\"Compute sums or means of 'bags' of embeddings, without instantiating the intermediate embeddings.\\n\\n    For bags of constant length, no :attr:`per_sample_weights`, no indices equal to :attr:`padding_idx`,\\n    and with 2D inputs, this class\\n\\n        * with ``mode=\\\"sum\\\"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.sum(dim=1)``,\\n        * with ``mode=\\\"mean\\\"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.mean(dim=1)``,\\n        * with ``mode=\\\"max\\\"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.max(dim=1)``.\\n\\n    However, :class:`~torch.nn.EmbeddingBag` is much more time and memory efficient than using a chain of these\\n    operations.\\n\\n    EmbeddingBag also supports per-sample weights as an argument to the forward\\n    pass. This scales the output of the Embedding before performing a weighted\\n    reduction as specified by ``mode``. If :attr:`per_sample_weights` is passed, the\\n    only supported ``mode`` is ``\\\"sum\\\"``, which computes a weighted sum according to\\n    :attr:`per_sample_weights`.\\n\\n    Args:\\n        num_embeddings (int): size of the dictionary of embeddings\\n        embedding_dim (int): the size of each embedding vector\\n        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`\\n                                    is renormalized to have norm :attr:`max_norm`.\\n        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.\\n        scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of\\n                                                the words in the mini-batch. Default ``False``.\\n                                                Note: this option is not supported when ``mode=\\\"max\\\"``.\\n        mode (str, optional): ``\\\"sum\\\"``, ``\\\"mean\\\"`` or ``\\\"max\\\"``. Specifies the way to reduce the bag.\\n                                 ``\\\"sum\\\"`` computes the weighted sum, taking :attr:`per_sample_weights`\\n                                 into consideration. ``\\\"mean\\\"`` computes the average of the values\\n                                 in the bag, ``\\\"max\\\"`` computes the max value over each bag.\\n                                 Default: ``\\\"mean\\\"``\\n        sparse (bool, optional): if ``True``, gradient w.r.t. :attr:`weight` matrix will be a sparse tensor. See\\n                                 Notes for more details regarding sparse gradients. Note: this option is not\\n                                 supported when ``mode=\\\"max\\\"``.\\n        include_last_offset (bool, optional): if ``True``, :attr:`offsets` has one additional element, where the last element\\n                                      is equivalent to the size of `indices`. This matches the CSR format.\\n        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the\\n                                     gradient; therefore, the embedding vector at :attr:`padding_idx` is not updated\\n                                     during training, i.e. it remains as a fixed \\\"pad\\\". For a newly constructed\\n                                     EmbeddingBag, the embedding vector at :attr:`padding_idx` will default to all\\n                                     zeros, but can be updated to another value to be used as the padding vector.\\n                                     Note that the embedding vector at :attr:`padding_idx` is excluded from the\\n                                     reduction.\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape `(num_embeddings, embedding_dim)`\\n                         initialized from :math:`\\\\mathcal{N}(0, 1)`.\\n\\n    Examples::\\n\\n        >>> # an EmbeddingBag module containing 10 tensors of size 3\\n        >>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum')\\n        >>> # a batch of 2 samples of 4 indices each\\n        >>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)\\n        >>> offsets = torch.tensor([0, 4], dtype=torch.long)\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> embedding_sum(input, offsets)\\n        tensor([[-0.8861, -5.4350, -0.0523],\\n                [ 1.1306, -2.5798, -1.0044]])\\n\\n        >>> # Example with padding_idx\\n        >>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum', padding_idx=2)\\n        >>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9], dtype=torch.long)\\n        >>> offsets = torch.tensor([0, 4], dtype=torch.long)\\n        >>> embedding_sum(input, offsets)\\n        tensor([[ 0.0000,  0.0000,  0.0000],\\n                [-0.7082,  3.2145, -2.6251]])\\n\\n        >>> # An EmbeddingBag can be loaded from an Embedding like so\\n        >>> embedding = nn.Embedding(10, 3, padding_idx=2)\\n        >>> embedding_sum = nn.EmbeddingBag.from_pretrained(\\n                embedding.weight,\\n                padding_idx=embedding.padding_idx,\\n                mode='sum')\\n    \\\"\\\"\\\"\\n\\n    __constants__ = [\\n        \\\"num_embeddings\\\",\\n        \\\"embedding_dim\\\",\\n        \\\"max_norm\\\",\\n        \\\"norm_type\\\",\\n        \\\"scale_grad_by_freq\\\",\\n        \\\"mode\\\",\\n        \\\"sparse\\\",\\n        \\\"include_last_offset\\\",\\n        \\\"padding_idx\\\",\\n    ]\\n\\n    num_embeddings: int\\n    embedding_dim: int\\n    max_norm: Optional[float]\\n    norm_type: float\\n    scale_grad_by_freq: bool\\n    weight: Tensor\\n    mode: str\\n    sparse: bool\\n    include_last_offset: bool\\n    padding_idx: Optional[int]\\n\\n    def __init__(\\n        self,\\n        num_embeddings: int,\\n        embedding_dim: int,\\n        max_norm: Optional[float] = None,\\n        norm_type: float = 2.0,\\n        scale_grad_by_freq: bool = False,\\n        mode: str = \\\"mean\\\",\\n        sparse: bool = False,\\n        _weight: Optional[Tensor] = None,\\n        include_last_offset: bool = False,\\n        padding_idx: Optional[int] = None,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        self.num_embeddings = num_embeddings\\n        self.embedding_dim = embedding_dim\\n        self.max_norm = max_norm\\n        self.norm_type = norm_type\\n        self.scale_grad_by_freq = scale_grad_by_freq\\n        if padding_idx is not None:\\n            if padding_idx > 0:\\n                assert (\\n                    padding_idx < self.num_embeddings\\n                ), \\\"padding_idx must be within num_embeddings\\\"\\n            elif padding_idx < 0:\\n                assert (\\n                    padding_idx >= -self.num_embeddings\\n                ), \\\"padding_idx must be within num_embeddings\\\"\\n                padding_idx = self.num_embeddings + padding_idx\\n        self.padding_idx = padding_idx\\n        if _weight is None:\\n            self.weight = Parameter(\\n                torch.empty((num_embeddings, embedding_dim), **factory_kwargs)\\n            )\\n            self.reset_parameters()\\n        else:\\n            assert list(_weight.shape) == [\\n                num_embeddings,\\n                embedding_dim,\\n            ], \\\"Shape of weight does not match num_embeddings and embedding_dim\\\"\\n            self.weight = Parameter(_weight)\\n        self.mode = mode\\n        self.sparse = sparse\\n        self.include_last_offset = include_last_offset\\n\\n    def reset_parameters(self) -> None:\\n        init.normal_(self.weight)\\n        self._fill_padding_idx_with_zero()\\n\\n    def _fill_padding_idx_with_zero(self) -> None:\\n        if self.padding_idx is not None:\\n            with torch.no_grad():\\n                self.weight[self.padding_idx].fill_(0)\\n\\n    def forward(\\n        self,\\n        input: Tensor,\\n        offsets: Optional[Tensor] = None,\\n        per_sample_weights: Optional[Tensor] = None,\\n    ) -> Tensor:\\n        \\\"\\\"\\\"Forward pass of EmbeddingBag.\\n\\n        Args:\\n            input (Tensor): Tensor containing bags of indices into the embedding matrix.\\n            offsets (Tensor, optional): Only used when :attr:`input` is 1D. :attr:`offsets` determines\\n                the starting index position of each bag (sequence) in :attr:`input`.\\n            per_sample_weights (Tensor, optional): a tensor of float / double weights, or None\\n                to indicate all weights should be taken to be ``1``. If specified, :attr:`per_sample_weights`\\n                must have exactly the same shape as input and is treated as having the same\\n                :attr:`offsets`, if those are not ``None``. Only supported for ``mode='sum'``.\\n\\n        Returns:\\n            Tensor output shape of `(B, embedding_dim)`.\\n\\n        .. note::\\n\\n            A few notes about ``input`` and ``offsets``:\\n\\n            - :attr:`input` and :attr:`offsets` have to be of the same type, either int or long\\n\\n            - If :attr:`input` is 2D of shape `(B, N)`, it will be treated as ``B`` bags (sequences)\\n              each of fixed length ``N``, and this will return ``B`` values aggregated in a way\\n              depending on the :attr:`mode`. :attr:`offsets` is ignored and required to be ``None`` in this case.\\n\\n            - If :attr:`input` is 1D of shape `(N)`, it will be treated as a concatenation of\\n              multiple bags (sequences).  :attr:`offsets` is required to be a 1D tensor containing the\\n              starting index positions of each bag in :attr:`input`. Therefore, for :attr:`offsets` of shape `(B)`,\\n              :attr:`input` will be viewed as having ``B`` bags. Empty bags (i.e., having 0-length) will have\\n              returned vectors filled by zeros.\\n        \\\"\\\"\\\"\\n        return F.embedding_bag(\\n            input,\\n            self.weight,\\n            offsets,\\n            self.max_norm,\\n            self.norm_type,\\n            self.scale_grad_by_freq,\\n            self.mode,\\n            self.sparse,\\n            per_sample_weights,\\n            self.include_last_offset,\\n            self.padding_idx,\\n        )\\n\\n    def extra_repr(self) -> str:\\n        s = \\\"{num_embeddings}, {embedding_dim}\\\"\\n        if self.max_norm is not None:\\n            s += \\\", max_norm={max_norm}\\\"\\n        if self.norm_type != 2:\\n            s += \\\", norm_type={norm_type}\\\"\\n        if self.scale_grad_by_freq is not False:\\n            s += \\\", scale_grad_by_freq={scale_grad_by_freq}\\\"\\n        s += \\\", mode={mode}\\\"\\n        if self.padding_idx is not None:\\n            s += \\\", padding_idx={padding_idx}\\\"\\n        return s.format(**{k: repr(v) for k, v in self.__dict__.items()})\\n\\n    @classmethod\\n    def from_pretrained(\\n        cls,\\n        embeddings: Tensor,\\n        freeze: bool = True,\\n        max_norm: Optional[float] = None,\\n        norm_type: float = 2.0,\\n        scale_grad_by_freq: bool = False,\\n        mode: str = \\\"mean\\\",\\n        sparse: bool = False,\\n        include_last_offset: bool = False,\\n        padding_idx: Optional[int] = None,\\n    ) -> \\\"EmbeddingBag\\\":\\n        r\\\"\\\"\\\"Create EmbeddingBag instance from given 2-dimensional FloatTensor.\\n\\n        Args:\\n            embeddings (Tensor): FloatTensor containing weights for the EmbeddingBag.\\n                First dimension is being passed to EmbeddingBag as 'num_embeddings', second as 'embedding_dim'.\\n            freeze (bool, optional): If ``True``, the tensor does not get updated in the learning process.\\n                Equivalent to ``embeddingbag.weight.requires_grad = False``. Default: ``True``\\n            max_norm (float, optional): See module initialization documentation. Default: ``None``\\n            norm_type (float, optional): See module initialization documentation. Default ``2``.\\n            scale_grad_by_freq (bool, optional): See module initialization documentation. Default ``False``.\\n            mode (str, optional): See module initialization documentation. Default: ``\\\"mean\\\"``\\n            sparse (bool, optional): See module initialization documentation. Default: ``False``.\\n            include_last_offset (bool, optional): See module initialization documentation. Default: ``False``.\\n            padding_idx (int, optional): See module initialization documentation. Default: ``None``.\\n\\n        Examples::\\n\\n            >>> # FloatTensor containing pretrained weights\\n            >>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])\\n            >>> embeddingbag = nn.EmbeddingBag.from_pretrained(weight)\\n            >>> # Get embeddings for index 1\\n            >>> input = torch.LongTensor([[1, 0]])\\n            >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n            >>> embeddingbag(input)\\n            tensor([[ 2.5000,  3.7000,  4.6500]])\\n        \\\"\\\"\\\"\\n        assert (\\n            embeddings.dim() == 2\\n        ), \\\"Embeddings parameter is expected to be 2-dimensional\\\"\\n        rows, cols = embeddings.shape\\n        embeddingbag = cls(\\n            num_embeddings=rows,\\n            embedding_dim=cols,\\n            _weight=embeddings,\\n            max_norm=max_norm,\\n            norm_type=norm_type,\\n            scale_grad_by_freq=scale_grad_by_freq,\\n            mode=mode,\\n            sparse=sparse,\\n            include_last_offset=include_last_offset,\\n            padding_idx=padding_idx,\\n        )\\n        embeddingbag.weight.requires_grad = not freeze\\n        return embeddingbag\\n\\n\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\n\\nfrom .module import Module\\n\\n\\n__all__ = [\\n    \\\"Dropout\\\",\\n    \\\"Dropout1d\\\",\\n    \\\"Dropout2d\\\",\\n    \\\"Dropout3d\\\",\\n    \\\"AlphaDropout\\\",\\n    \\\"FeatureAlphaDropout\\\",\\n]\\n\\n\\nclass _DropoutNd(Module):\\n    __constants__ = [\\\"p\\\", \\\"inplace\\\"]\\n    p: float\\n    inplace: bool\\n\\n    def __init__(self, p: float = 0.5, inplace: bool = False) -> None:\\n        super().__init__()\\n        if p < 0 or p > 1:\\n            raise ValueError(\\n                f\\\"dropout probability has to be between 0 and 1, but got {p}\\\"\\n            )\\n        self.p = p\\n        self.inplace = inplace\\n\\n    def extra_repr(self) -> str:\\n        return f\\\"p={self.p}, inplace={self.inplace}\\\"\\n\\n\\nclass Dropout(_DropoutNd):\\n    r\\\"\\\"\\\"During training, randomly zeroes some of the elements of the input tensor with probability :attr:`p`.\\n\\n    The zeroed elements are chosen independently for each forward call and are sampled from a Bernoulli distribution.\\n\\n    Each channel will be zeroed out independently on every forward call.\\n\\n    This has proven to be an effective technique for regularization and\\n    preventing the co-adaptation of neurons as described in the paper\\n    `Improving neural networks by preventing co-adaptation of feature\\n    detectors`_ .\\n\\n    Furthermore, the outputs are scaled by a factor of :math:`\\\\frac{1}{1-p}` during\\n    training. This means that during evaluation the module simply computes an\\n    identity function.\\n\\n    Args:\\n        p: probability of an element to be zeroed. Default: 0.5\\n        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\\n\\n    Shape:\\n        - Input: :math:`(*)`. Input can be of any shape\\n        - Output: :math:`(*)`. Output is of the same shape as input\\n\\n    Examples::\\n\\n        >>> m = nn.Dropout(p=0.2)\\n        >>> input = torch.randn(20, 16)\\n        >>> output = m(input)\\n\\n    .. _Improving neural networks by preventing co-adaptation of feature\\n        detectors: https://arxiv.org/abs/1207.0580\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.dropout(input, self.p, self.training, self.inplace)\\n\\n\\nclass Dropout1d(_DropoutNd):\\n    r\\\"\\\"\\\"Randomly zero out entire channels.\\n\\n    A channel is a 1D feature map,\\n    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the\\n    batched input is a 1D tensor :math:`\\\\text{input}[i, j]`.\\n\\n    Each channel will be zeroed out independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n\\n    Usually the input comes from :class:`nn.Conv1d` modules.\\n\\n    As described in the paper\\n    `Efficient Object Localization Using Convolutional Networks`_ ,\\n    if adjacent pixels within feature maps are strongly correlated\\n    (as is normally the case in early convolution layers) then i.i.d. dropout\\n    will not regularize the activations and will otherwise just result\\n    in an effective learning rate decrease.\\n\\n    In this case, :func:`nn.Dropout1d` will help promote independence between\\n    feature maps and should be used instead.\\n\\n    Args:\\n        p (float, optional): probability of an element to be zero-ed.\\n        inplace (bool, optional): If set to ``True``, will do this operation\\n            in-place\\n\\n    Shape:\\n        - Input: :math:`(N, C, L)` or :math:`(C, L)`.\\n        - Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input).\\n\\n    Examples::\\n\\n        >>> m = nn.Dropout1d(p=0.2)\\n        >>> input = torch.randn(20, 16, 32)\\n        >>> output = m(input)\\n\\n    .. _Efficient Object Localization Using Convolutional Networks:\\n       https://arxiv.org/abs/1411.4280\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.dropout1d(input, self.p, self.training, self.inplace)\\n\\n\\nclass Dropout2d(_DropoutNd):\\n    r\\\"\\\"\\\"Randomly zero out entire channels.\\n\\n    A channel is a 2D feature map,\\n    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the\\n    batched input is a 2D tensor :math:`\\\\text{input}[i, j]`.\\n\\n    Each channel will be zeroed out independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n\\n    Usually the input comes from :class:`nn.Conv2d` modules.\\n\\n    As described in the paper\\n    `Efficient Object Localization Using Convolutional Networks`_ ,\\n    if adjacent pixels within feature maps are strongly correlated\\n    (as is normally the case in early convolution layers) then i.i.d. dropout\\n    will not regularize the activations and will otherwise just result\\n    in an effective learning rate decrease.\\n\\n    In this case, :func:`nn.Dropout2d` will help promote independence between\\n    feature maps and should be used instead.\\n\\n    Args:\\n        p (float, optional): probability of an element to be zero-ed.\\n        inplace (bool, optional): If set to ``True``, will do this operation\\n            in-place\\n\\n    .. warning ::\\n        Due to historical reasons, this class will perform 1D channel-wise dropout\\n        for 3D inputs (as done by :class:`nn.Dropout1d`). Thus, it currently does NOT\\n        support inputs without a batch dimension of shape :math:`(C, H, W)`. This\\n        behavior will change in a future release to interpret 3D inputs as no-batch-dim\\n        inputs. To maintain the old behavior, switch to :class:`nn.Dropout1d`.\\n\\n    Shape:\\n        - Input: :math:`(N, C, H, W)` or :math:`(N, C, L)`.\\n        - Output: :math:`(N, C, H, W)` or :math:`(N, C, L)` (same shape as input).\\n\\n    Examples::\\n\\n        >>> m = nn.Dropout2d(p=0.2)\\n        >>> input = torch.randn(20, 16, 32, 32)\\n        >>> output = m(input)\\n\\n    .. _Efficient Object Localization Using Convolutional Networks:\\n       https://arxiv.org/abs/1411.4280\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.dropout2d(input, self.p, self.training, self.inplace)\\n\\n\\nclass Dropout3d(_DropoutNd):\\n    r\\\"\\\"\\\"Randomly zero out entire channels.\\n\\n    A channel is a 3D feature map,\\n    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the\\n    batched input is a 3D tensor :math:`\\\\text{input}[i, j]`.\\n\\n    Each channel will be zeroed out independently on every forward call with\\n    probability :attr:`p` using samples from a Bernoulli distribution.\\n\\n    Usually the input comes from :class:`nn.Conv3d` modules.\\n\\n    As described in the paper\\n    `Efficient Object Localization Using Convolutional Networks`_ ,\\n    if adjacent pixels within feature maps are strongly correlated\\n    (as is normally the case in early convolution layers) then i.i.d. dropout\\n    will not regularize the activations and will otherwise just result\\n    in an effective learning rate decrease.\\n\\n    In this case, :func:`nn.Dropout3d` will help promote independence between\\n    feature maps and should be used instead.\\n\\n    Args:\\n        p (float, optional): probability of an element to be zeroed.\\n        inplace (bool, optional): If set to ``True``, will do this operation\\n            in-place\\n\\n    Shape:\\n        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`.\\n        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input).\\n\\n    Examples::\\n\\n        >>> m = nn.Dropout3d(p=0.2)\\n        >>> input = torch.randn(20, 16, 4, 32, 32)\\n        >>> output = m(input)\\n\\n    .. _Efficient Object Localization Using Convolutional Networks:\\n       https://arxiv.org/abs/1411.4280\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.dropout3d(input, self.p, self.training, self.inplace)\\n\\n\\nclass AlphaDropout(_DropoutNd):\\n    r\\\"\\\"\\\"Applies Alpha Dropout over the input.\\n\\n    Alpha Dropout is a type of Dropout that maintains the self-normalizing\\n    property.\\n    For an input with zero mean and unit standard deviation, the output of\\n    Alpha Dropout maintains the original mean and standard deviation of the\\n    input.\\n    Alpha Dropout goes hand-in-hand with SELU activation function, which ensures\\n    that the outputs have zero mean and unit standard deviation.\\n\\n    During training, it randomly masks some of the elements of the input\\n    tensor with probability *p* using samples from a bernoulli distribution.\\n    The elements to masked are randomized on every forward call, and scaled\\n    and shifted to maintain zero mean and unit standard deviation.\\n\\n    During evaluation the module simply computes an identity function.\\n\\n    More details can be found in the paper `Self-Normalizing Neural Networks`_ .\\n\\n    Args:\\n        p (float): probability of an element to be dropped. Default: 0.5\\n        inplace (bool, optional): If set to ``True``, will do this operation\\n            in-place\\n\\n    Shape:\\n        - Input: :math:`(*)`. Input can be of any shape\\n        - Output: :math:`(*)`. Output is of the same shape as input\\n\\n    Examples::\\n\\n        >>> m = nn.AlphaDropout(p=0.2)\\n        >>> input = torch.randn(20, 16)\\n        >>> output = m(input)\\n\\n    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.alpha_dropout(input, self.p, self.training)\\n\\n\\nclass FeatureAlphaDropout(_DropoutNd):\\n    r\\\"\\\"\\\"Randomly masks out entire channels.\\n\\n    A channel is a feature map,\\n    e.g. the :math:`j`-th channel of the :math:`i`-th sample in the batch input\\n    is a tensor :math:`\\\\text{input}[i, j]` of the input tensor). Instead of\\n    setting activations to zero, as in regular Dropout, the activations are set\\n    to the negative saturation value of the SELU activation function. More details\\n    can be found in the paper `Self-Normalizing Neural Networks`_ .\\n\\n    Each element will be masked independently for each sample on every forward\\n    call with probability :attr:`p` using samples from a Bernoulli distribution.\\n    The elements to be masked are randomized on every forward call, and scaled\\n    and shifted to maintain zero mean and unit variance.\\n\\n    Usually the input comes from :class:`nn.AlphaDropout` modules.\\n\\n    As described in the paper\\n    `Efficient Object Localization Using Convolutional Networks`_ ,\\n    if adjacent pixels within feature maps are strongly correlated\\n    (as is normally the case in early convolution layers) then i.i.d. dropout\\n    will not regularize the activations and will otherwise just result\\n    in an effective learning rate decrease.\\n\\n    In this case, :func:`nn.AlphaDropout` will help promote independence between\\n    feature maps and should be used instead.\\n\\n    Args:\\n        p (float, optional): probability of an element to be zeroed. Default: 0.5\\n        inplace (bool, optional): If set to ``True``, will do this operation\\n            in-place\\n\\n    Shape:\\n        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`.\\n        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input).\\n\\n    Examples::\\n\\n        >>> m = nn.FeatureAlphaDropout(p=0.2)\\n        >>> input = torch.randn(20, 16, 4, 32, 32)\\n        >>> output = m(input)\\n\\n    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515\\n    .. _Efficient Object Localization Using Convolutional Networks:\\n       https://arxiv.org/abs/1411.4280\\n    \\\"\\\"\\\"\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return F.feature_alpha_dropout(input, self.p, self.training)\\n\\n\\n# mypy: allow-untyped-defs\\n\\nfrom collections import namedtuple\\nfrom typing import List, Sequence\\n\\nimport torch\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\n\\nfrom .container import ModuleList, Sequential\\nfrom .linear import Linear\\nfrom .module import Module\\n\\n\\n__all__ = [\\\"AdaptiveLogSoftmaxWithLoss\\\"]\\n\\n_ASMoutput = namedtuple(\\\"_ASMoutput\\\", [\\\"output\\\", \\\"loss\\\"])\\n\\n\\nclass AdaptiveLogSoftmaxWithLoss(Module):\\n    \\\"\\\"\\\"Efficient softmax approximation.\\n\\n    As described in\\n    `Efficient softmax approximation for GPUs by Edouard Grave, Armand Joulin,\\n    Moustapha Ciss\\\\u00e9, David Grangier, and Herv\\\\u00e9 J\\\\u00e9gou\\n    <https://arxiv.org/abs/1609.04309>`__.\\n\\\"\\\"\\\" r\\\"\\\"\\\"\\n    Adaptive softmax is an approximate strategy for training models with large\\n    output spaces. It is most effective when the label distribution is highly\\n    imbalanced, for example in natural language modelling, where the word\\n    frequency distribution approximately follows the `Zipf's law`_.\\n\\n    Adaptive softmax partitions the labels into several clusters, according to\\n    their frequency. These clusters may contain different number of targets\\n    each.\\n    Additionally, clusters containing less frequent labels assign lower\\n    dimensional embeddings to those labels, which speeds up the computation.\\n    For each minibatch, only clusters for which at least one target is\\n    present are evaluated.\\n\\n    The idea is that the clusters which are accessed frequently\\n    (like the first one, containing most frequent labels), should also be cheap\\n    to compute -- that is, contain a small number of assigned labels.\\n\\n    We highly recommend taking a look at the original paper for more details.\\n\\n    * :attr:`cutoffs` should be an ordered Sequence of integers sorted\\n      in the increasing order.\\n      It controls number of clusters and the partitioning of targets into\\n      clusters. For example setting ``cutoffs = [10, 100, 1000]``\\n      means that first `10` targets will be assigned\\n      to the 'head' of the adaptive softmax, targets `11, 12, ..., 100` will be\\n      assigned to the first cluster, and targets `101, 102, ..., 1000` will be\\n      assigned to the second cluster, while targets\\n      `1001, 1002, ..., n_classes - 1` will be assigned\\n      to the last, third cluster.\\n\\n    * :attr:`div_value` is used to compute the size of each additional cluster,\\n      which is given as\\n      :math:`\\\\left\\\\lfloor\\\\frac{\\\\texttt{in\\\\_features}}{\\\\texttt{div\\\\_value}^{idx}}\\\\right\\\\rfloor`,\\n      where :math:`idx` is the cluster index (with clusters\\n      for less frequent words having larger indices,\\n      and indices starting from :math:`1`).\\n\\n    * :attr:`head_bias` if set to True, adds a bias term to the 'head' of the\\n      adaptive softmax. See paper for details. Set to False in the official\\n      implementation.\\n\\n    .. warning::\\n        Labels passed as inputs to this module should be sorted according to\\n        their frequency. This means that the most frequent label should be\\n        represented by the index `0`, and the least frequent\\n        label should be represented by the index `n_classes - 1`.\\n\\n    .. note::\\n        This module returns a ``NamedTuple`` with ``output``\\n        and ``loss`` fields. See further documentation for details.\\n\\n    .. note::\\n        To compute log-probabilities for all classes, the ``log_prob``\\n        method can be used.\\n\\n    Args:\\n        in_features (int): Number of features in the input tensor\\n        n_classes (int): Number of classes in the dataset\\n        cutoffs (Sequence): Cutoffs used to assign targets to their buckets\\n        div_value (float, optional): value used as an exponent to compute sizes\\n            of the clusters. Default: 4.0\\n        head_bias (bool, optional): If ``True``, adds a bias term to the 'head' of the\\n            adaptive softmax. Default: ``False``\\n\\n    Returns:\\n        ``NamedTuple`` with ``output`` and ``loss`` fields:\\n            * **output** is a Tensor of size ``N`` containing computed target\\n              log probabilities for each example\\n            * **loss** is a Scalar representing the computed negative\\n              log likelihood loss\\n\\n    Shape:\\n        - input: :math:`(N, \\\\texttt{in\\\\_features})` or :math:`(\\\\texttt{in\\\\_features})`\\n        - target: :math:`(N)` or :math:`()` where each value satisfies :math:`0 <= \\\\texttt{target[i]} <= \\\\texttt{n\\\\_classes}`\\n        - output1: :math:`(N)` or :math:`()`\\n        - output2: ``Scalar``\\n\\n    .. _Zipf's law: https://en.wikipedia.org/wiki/Zipf%27s_law\\n    \\\"\\\"\\\"\\n\\n    in_features: int\\n    n_classes: int\\n    cutoffs: List[int]\\n    div_value: float\\n    head_bias: bool\\n    head: Linear\\n    tail: ModuleList\\n\\n    def __init__(\\n        self,\\n        in_features: int,\\n        n_classes: int,\\n        cutoffs: Sequence[int],\\n        div_value: float = 4.0,\\n        head_bias: bool = False,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n\\n        cutoffs = list(cutoffs)\\n\\n        if len(cutoffs) == 0:\\n            raise ValueError(\\\"cutoffs should be a sequence of length larger than 0\\\")\\n\\n        if (\\n            (cutoffs != sorted(cutoffs))\\n            or (min(cutoffs) <= 0)\\n            or (max(cutoffs) > (n_classes - 1))\\n            or (len(set(cutoffs)) != len(cutoffs))\\n            or any(int(c) != c for c in cutoffs)\\n        ):\\n            raise ValueError(\\n                \\\"cutoffs should be a sequence of unique, positive \\\"\\n                \\\"integers sorted in an increasing order, where \\\"\\n                \\\"each value is between 1 and n_classes-1\\\"\\n            )\\n\\n        self.in_features = in_features\\n        self.n_classes = n_classes\\n        self.cutoffs = cutoffs + [n_classes]\\n        self.div_value = div_value\\n        self.head_bias = head_bias\\n\\n        self.shortlist_size = self.cutoffs[0]\\n        self.n_clusters = len(self.cutoffs) - 1\\n        self.head_size = self.shortlist_size + self.n_clusters\\n\\n        self.head = Linear(\\n            self.in_features, self.head_size, bias=self.head_bias, **factory_kwargs\\n        )\\n        self.tail = ModuleList()\\n\\n        for i in range(self.n_clusters):\\n            hsz = int(self.in_features // (self.div_value ** (i + 1)))\\n            osz = self.cutoffs[i + 1] - self.cutoffs[i]\\n\\n            projection = Sequential(\\n                Linear(self.in_features, hsz, bias=False, **factory_kwargs),\\n                Linear(hsz, osz, bias=False, **factory_kwargs),\\n            )\\n\\n            self.tail.append(projection)\\n\\n    def reset_parameters(self) -> None:\\n        self.head.reset_parameters()\\n        for i2h, h2o in self.tail:\\n            i2h.reset_parameters()\\n            h2o.reset_parameters()\\n\\n    def forward(self, input_: Tensor, target_: Tensor) -> _ASMoutput:\\n        targ_dim = target_.dim()\\n\\n        if targ_dim == 1:\\n            if input_.size(0) != target_.size(0):\\n                raise RuntimeError(\\n                    \\\"Input and target should have the same size \\\"\\n                    \\\"in the batch dimension.\\\"\\n                )\\n            if input_.dim() != 2:\\n                raise RuntimeError(\\n                    \\\"1D target tensor expects 2D input tensors, \\\"\\n                    \\\"but found inputs with size\\\",\\n                    input_.size(),\\n                )\\n        elif targ_dim == 0:\\n            if input_.dim() != 1:\\n                raise RuntimeError(\\n                    \\\"0D target tensor expects 1D input tensors, \\\"\\n                    \\\"but found inputs with size\\\",\\n                    input_.size(),\\n                )\\n        else:\\n            raise RuntimeError(\\n                \\\"0D or 1D target tensor expected, \\\" \\\"multi-target not supported\\\"\\n            )\\n\\n        is_batched = targ_dim > 0\\n        input = input_ if is_batched else input_.unsqueeze(0)\\n        target = target_ if is_batched else target_.unsqueeze(0)\\n\\n        used_rows = 0\\n        batch_size = target.size(0)\\n\\n        output = input.new_zeros(batch_size)\\n        gather_inds = target.new_empty(batch_size)\\n\\n        cutoff_values = [0] + self.cutoffs\\n        for i in range(len(cutoff_values) - 1):\\n            low_idx = cutoff_values[i]\\n            high_idx = cutoff_values[i + 1]\\n\\n            target_mask = (target >= low_idx) & (target < high_idx)\\n            row_indices = target_mask.nonzero().squeeze()\\n\\n            if row_indices.numel() == 0:\\n                continue\\n\\n            if i == 0:\\n                gather_inds.index_copy_(0, row_indices, target[target_mask])\\n\\n            else:\\n                relative_target = target[target_mask] - low_idx\\n                input_subset = input.index_select(0, row_indices)\\n\\n                cluster_output = self.tail[i - 1](input_subset)\\n                cluster_index = self.shortlist_size + i - 1\\n\\n                gather_inds.index_fill_(0, row_indices, cluster_index)\\n                cluster_logprob = F.log_softmax(cluster_output, dim=1)\\n                local_logprob = cluster_logprob.gather(1, relative_target.unsqueeze(1))\\n                output.index_copy_(0, row_indices, local_logprob.squeeze(1))\\n\\n            used_rows += row_indices.numel()\\n\\n        if used_rows != batch_size:\\n            raise RuntimeError(\\n                f\\\"Target values should be in [0, {self.n_classes - 1}], \\\"\\n                f\\\"but values in range [{target.min().item()}, {target.max().item()}] \\\"\\n                \\\"were found. \\\"\\n            )\\n\\n        head_output = self.head(input)\\n        head_logprob = F.log_softmax(head_output, dim=1)\\n        output += head_logprob.gather(1, gather_inds.unsqueeze(1)).squeeze()\\n        loss = (-output).mean()\\n\\n        if not is_batched:\\n            output = output.squeeze(0)\\n\\n        return _ASMoutput(output, loss)\\n\\n    def _get_full_log_prob(self, input, head_output):\\n        \\\"\\\"\\\"Given input tensor, and output of ``self.head``, compute the log of the full distribution.\\\"\\\"\\\"\\n        out = input.new_empty((head_output.size(0), self.n_classes))\\n        head_logprob = F.log_softmax(head_output, dim=1)\\n\\n        out[:, : self.shortlist_size] = head_logprob[:, : self.shortlist_size]\\n\\n        for i, (start_idx, stop_idx) in enumerate(zip(self.cutoffs, self.cutoffs[1:])):\\n            cluster_output = self.tail[i](input)\\n            cluster_logprob = F.log_softmax(cluster_output, dim=1)\\n            output_logprob = cluster_logprob + head_logprob[\\n                :, self.shortlist_size + i\\n            ].unsqueeze(1)\\n\\n            out[:, start_idx:stop_idx] = output_logprob\\n\\n        return out\\n\\n    def log_prob(self, input: Tensor) -> Tensor:\\n        r\\\"\\\"\\\"Compute log probabilities for all :math:`\\\\texttt{n\\\\_classes}`.\\n\\n        Args:\\n            input (Tensor): a minibatch of examples\\n\\n        Returns:\\n            log-probabilities of for each class :math:`c`\\n            in range :math:`0 <= c <= \\\\texttt{n\\\\_classes}`, where :math:`\\\\texttt{n\\\\_classes}` is a\\n            parameter passed to ``AdaptiveLogSoftmaxWithLoss`` constructor.\\n\\n        Shape:\\n            - Input: :math:`(N, \\\\texttt{in\\\\_features})`\\n            - Output: :math:`(N, \\\\texttt{n\\\\_classes})`\\n\\n        \\\"\\\"\\\"\\n        head_output = self.head(input)\\n        return self._get_full_log_prob(input, head_output)\\n\\n    def predict(self, input: Tensor) -> Tensor:\\n        r\\\"\\\"\\\"Return the class with the highest probability for each example in the input minibatch.\\n\\n        This is equivalent to ``self.log_prob(input).argmax(dim=1)``, but is more efficient in some cases.\\n\\n        Args:\\n            input (Tensor): a minibatch of examples\\n\\n        Returns:\\n            output (Tensor): a class with the highest probability for each example\\n\\n        Shape:\\n            - Input: :math:`(N, \\\\texttt{in\\\\_features})`\\n            - Output: :math:`(N)`\\n        \\\"\\\"\\\"\\n        head_output = self.head(input)\\n        output = torch.argmax(head_output, dim=1)\\n        not_in_shortlist = output >= self.shortlist_size\\n        all_in_shortlist = not (not_in_shortlist.any())\\n\\n        if all_in_shortlist:\\n            return output\\n\\n        elif not_in_shortlist.all():\\n            log_prob = self._get_full_log_prob(input, head_output)\\n            return torch.argmax(log_prob, dim=1)\\n\\n        else:\\n            log_prob = self._get_full_log_prob(\\n                input[not_in_shortlist], head_output[not_in_shortlist]\\n            )\\n            output[not_in_shortlist] = torch.argmax(log_prob, dim=1)\\n            return output\\n\\n\\n# mypy: allow-untyped-defs\\nimport math\\nfrom typing import List, Optional, Tuple, Union\\nfrom typing_extensions import deprecated\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch._torch_docs import reproducibility_notes\\nfrom torch.nn import functional as F, init\\nfrom torch.nn.common_types import _size_1_t, _size_2_t, _size_3_t\\nfrom torch.nn.parameter import Parameter, UninitializedParameter\\n\\nfrom .lazy import LazyModuleMixin\\nfrom .module import Module\\nfrom .utils import _pair, _reverse_repeat_tuple, _single, _triple\\n\\n\\n__all__ = [\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n    \\\"LazyConv1d\\\",\\n    \\\"LazyConv2d\\\",\\n    \\\"LazyConv3d\\\",\\n    \\\"LazyConvTranspose1d\\\",\\n    \\\"LazyConvTranspose2d\\\",\\n    \\\"LazyConvTranspose3d\\\",\\n]\\n\\nconvolution_notes = {\\n    \\\"groups_note\\\": r\\\"\\\"\\\"* :attr:`groups` controls the connections between inputs and outputs.\\n      :attr:`in_channels` and :attr:`out_channels` must both be divisible by\\n      :attr:`groups`. For example,\\n\\n        * At groups=1, all inputs are convolved to all outputs.\\n        * At groups=2, the operation becomes equivalent to having two conv\\n          layers side by side, each seeing half the input channels\\n          and producing half the output channels, and both subsequently\\n          concatenated.\\n        * At groups= :attr:`in_channels`, each input channel is convolved with\\n          its own set of filters (of size\\n          :math:`\\\\frac{\\\\text{out\\\\_channels}}{\\\\text{in\\\\_channels}}`).\\\"\\\"\\\",\\n    \\\"depthwise_separable_note\\\": r\\\"\\\"\\\"When `groups == in_channels` and `out_channels == K * in_channels`,\\n        where `K` is a positive integer, this operation is also known as a \\\"depthwise convolution\\\".\\n\\n        In other words, for an input of size :math:`(N, C_{in}, L_{in})`,\\n        a depthwise convolution with a depthwise multiplier `K` can be performed with the arguments\\n        :math:`(C_\\\\text{in}=C_\\\\text{in}, C_\\\\text{out}=C_\\\\text{in} \\\\times \\\\text{K}, ..., \\\\text{groups}=C_\\\\text{in})`.\\\"\\\"\\\",\\n}  # noqa: B950\\n\\n\\nclass _ConvNd(Module):\\n    __constants__ = [\\n        \\\"stride\\\",\\n        \\\"padding\\\",\\n        \\\"dilation\\\",\\n        \\\"groups\\\",\\n        \\\"padding_mode\\\",\\n        \\\"output_padding\\\",\\n        \\\"in_channels\\\",\\n        \\\"out_channels\\\",\\n        \\\"kernel_size\\\",\\n    ]\\n    __annotations__ = {\\\"bias\\\": Optional[torch.Tensor]}\\n\\n    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor:  # type: ignore[empty-body]\\n        ...\\n\\n    in_channels: int\\n    _reversed_padding_repeated_twice: List[int]\\n    out_channels: int\\n    kernel_size: Tuple[int, ...]\\n    stride: Tuple[int, ...]\\n    padding: Union[str, Tuple[int, ...]]\\n    dilation: Tuple[int, ...]\\n    transposed: bool\\n    output_padding: Tuple[int, ...]\\n    groups: int\\n    padding_mode: str\\n    weight: Tensor\\n    bias: Optional[Tensor]\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: Tuple[int, ...],\\n        stride: Tuple[int, ...],\\n        padding: Tuple[int, ...],\\n        dilation: Tuple[int, ...],\\n        transposed: bool,\\n        output_padding: Tuple[int, ...],\\n        groups: int,\\n        bias: bool,\\n        padding_mode: str,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__()\\n        if groups <= 0:\\n            raise ValueError(\\\"groups must be a positive integer\\\")\\n        if in_channels % groups != 0:\\n            raise ValueError(\\\"in_channels must be divisible by groups\\\")\\n        if out_channels % groups != 0:\\n            raise ValueError(\\\"out_channels must be divisible by groups\\\")\\n        valid_padding_strings = {\\\"same\\\", \\\"valid\\\"}\\n        if isinstance(padding, str):\\n            if padding not in valid_padding_strings:\\n                raise ValueError(\\n                    f\\\"Invalid padding string {padding!r}, should be one of {valid_padding_strings}\\\"\\n                )\\n            if padding == \\\"same\\\" and any(s != 1 for s in stride):\\n                raise ValueError(\\n                    \\\"padding='same' is not supported for strided convolutions\\\"\\n                )\\n\\n        valid_padding_modes = {\\\"zeros\\\", \\\"reflect\\\", \\\"replicate\\\", \\\"circular\\\"}\\n        if padding_mode not in valid_padding_modes:\\n            raise ValueError(\\n                f\\\"padding_mode must be one of {valid_padding_modes}, but got padding_mode='{padding_mode}'\\\"\\n            )\\n        self.in_channels = in_channels\\n        self.out_channels = out_channels\\n        self.kernel_size = kernel_size\\n        self.stride = stride\\n        self.padding = padding\\n        self.dilation = dilation\\n        self.transposed = transposed\\n        self.output_padding = output_padding\\n        self.groups = groups\\n        self.padding_mode = padding_mode\\n        # `_reversed_padding_repeated_twice` is the padding to be passed to\\n        # `F.pad` if needed (e.g., for non-zero padding types that are\\n        # implemented as two ops: padding + conv). `F.pad` accepts paddings in\\n        # reverse order than the dimension.\\n        if isinstance(self.padding, str):\\n            self._reversed_padding_repeated_twice = [0, 0] * len(kernel_size)\\n            if padding == \\\"same\\\":\\n                for d, k, i in zip(\\n                    dilation, kernel_size, range(len(kernel_size) - 1, -1, -1)\\n                ):\\n                    total_padding = d * (k - 1)\\n                    left_pad = total_padding // 2\\n                    self._reversed_padding_repeated_twice[2 * i] = left_pad\\n                    self._reversed_padding_repeated_twice[2 * i + 1] = (\\n                        total_padding - left_pad\\n                    )\\n        else:\\n            self._reversed_padding_repeated_twice = _reverse_repeat_tuple(\\n                self.padding, 2\\n            )\\n\\n        if transposed:\\n            self.weight = Parameter(\\n                torch.empty(\\n                    (in_channels, out_channels // groups, *kernel_size),\\n                    **factory_kwargs,\\n                )\\n            )\\n        else:\\n            self.weight = Parameter(\\n                torch.empty(\\n                    (out_channels, in_channels // groups, *kernel_size),\\n                    **factory_kwargs,\\n                )\\n            )\\n        if bias:\\n            self.bias = Parameter(torch.empty(out_channels, **factory_kwargs))\\n        else:\\n            self.register_parameter(\\\"bias\\\", None)\\n\\n        self.reset_parameters()\\n\\n    def reset_parameters(self) -> None:\\n        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with\\n        # uniform(-1/sqrt(k), 1/sqrt(k)), where k = weight.size(1) * prod(*kernel_size)\\n        # For more details see: https://github.com/pytorch/pytorch/issues/15314#issuecomment-477448573\\n        init.kaiming_uniform_(self.weight, a=math.sqrt(5))\\n        if self.bias is not None:\\n            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)\\n            if fan_in != 0:\\n                bound = 1 / math.sqrt(fan_in)\\n                init.uniform_(self.bias, -bound, bound)\\n\\n    def extra_repr(self):\\n        s = (\\n            \\\"{in_channels}, {out_channels}, kernel_size={kernel_size}\\\"\\n            \\\", stride={stride}\\\"\\n        )\\n        if self.padding != (0,) * len(self.padding):\\n            s += \\\", padding={padding}\\\"\\n        if self.dilation != (1,) * len(self.dilation):\\n            s += \\\", dilation={dilation}\\\"\\n        if self.output_padding != (0,) * len(self.output_padding):\\n            s += \\\", output_padding={output_padding}\\\"\\n        if self.groups != 1:\\n            s += \\\", groups={groups}\\\"\\n        if self.bias is None:\\n            s += \\\", bias=False\\\"\\n        if self.padding_mode != \\\"zeros\\\":\\n            s += \\\", padding_mode={padding_mode}\\\"\\n        return s.format(**self.__dict__)\\n\\n    def __setstate__(self, state):\\n        super().__setstate__(state)\\n        if not hasattr(self, \\\"padding_mode\\\"):\\n            self.padding_mode = \\\"zeros\\\"\\n\\n\\nclass Conv1d(_ConvNd):\\n    __doc__ = (\\n        r\\\"\\\"\\\"Applies a 1D convolution over an input signal composed of several input\\n    planes.\\n\\n    In the simplest case, the output value of the layer with input size\\n    :math:`(N, C_{\\\\text{in}}, L)` and output :math:`(N, C_{\\\\text{out}}, L_{\\\\text{out}})` can be\\n    precisely described as:\\n\\n    .. math::\\n        \\\\text{out}(N_i, C_{\\\\text{out}_j}) = \\\\text{bias}(C_{\\\\text{out}_j}) +\\n        \\\\sum_{k = 0}^{C_{in} - 1} \\\\text{weight}(C_{\\\\text{out}_j}, k)\\n        \\\\star \\\\text{input}(N_i, k)\\n\\n    where :math:`\\\\star` is the valid `cross-correlation`_ operator,\\n    :math:`N` is a batch size, :math:`C` denotes a number of channels,\\n    :math:`L` is a length of signal sequence.\\n    \\\"\\\"\\\"\\n        + r\\\"\\\"\\\"\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    * :attr:`stride` controls the stride for the cross-correlation, a single\\n      number or a one-element tuple.\\n\\n    * :attr:`padding` controls the amount of padding applied to the input. It\\n      can be either a string {{'valid', 'same'}} or a tuple of ints giving the\\n      amount of implicit padding applied on both sides.\\n\\\"\\\"\\\"\\n        \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also\\n      known as the \\\\u00e0 trous algorithm. It is harder to describe, but this `link`_\\n      has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\"\\n        r\\\"\\\"\\\"\\n    {groups_note}\\n\\n    Note:\\n        {depthwise_separable_note}\\n    Note:\\n        {cudnn_reproducibility_note}\\n\\n    Note:\\n        ``padding='valid'`` is the same as no padding. ``padding='same'`` pads\\n        the input so the output has the shape as the input. However, this mode\\n        doesn't support any stride values other than 1.\\n\\n    Note:\\n        This module supports complex data types i.e. ``complex32, complex64, complex128``.\\n\\n    Args:\\n        in_channels (int): Number of channels in the input image\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int, tuple or str, optional): Padding added to both sides of\\n            the input. Default: 0\\n        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,\\n            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``\\n        dilation (int or tuple, optional): Spacing between kernel\\n            elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input\\n            channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the\\n            output. Default: ``True``\\n\\n    \\\"\\\"\\\".format(\\n            **reproducibility_notes, **convolution_notes\\n        )\\n        + r\\\"\\\"\\\"\\n\\n    Shape:\\n        - Input: :math:`(N, C_{in}, L_{in})` or :math:`(C_{in}, L_{in})`\\n        - Output: :math:`(N, C_{out}, L_{out})` or :math:`(C_{out}, L_{out})`, where\\n\\n          .. math::\\n              L_{out} = \\\\left\\\\lfloor\\\\frac{L_{in} + 2 \\\\times \\\\text{padding} - \\\\text{dilation}\\n                        \\\\times (\\\\text{kernel\\\\_size} - 1) - 1}{\\\\text{stride}} + 1\\\\right\\\\rfloor\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape\\n            :math:`(\\\\text{out\\\\_channels},\\n            \\\\frac{\\\\text{in\\\\_channels}}{\\\\text{groups}}, \\\\text{kernel\\\\_size})`.\\n            The values of these weights are sampled from\\n            :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n            :math:`k = \\\\frac{groups}{C_\\\\text{in} * \\\\text{kernel\\\\_size}}`\\n        bias (Tensor):   the learnable bias of the module of shape\\n            (out_channels). If :attr:`bias` is ``True``, then the values of these weights are\\n            sampled from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n            :math:`k = \\\\frac{groups}{C_\\\\text{in} * \\\\text{kernel\\\\_size}}`\\n\\n    Examples::\\n\\n        >>> m = nn.Conv1d(16, 33, 3, stride=2)\\n        >>> input = torch.randn(20, 16, 50)\\n        >>> output = m(input)\\n\\n    .. _cross-correlation:\\n        https://en.wikipedia.org/wiki/Cross-correlation\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n    \\\"\\\"\\\"\\n    )\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: _size_1_t,\\n        stride: _size_1_t = 1,\\n        padding: Union[str, _size_1_t] = 0,\\n        dilation: _size_1_t = 1,\\n        groups: int = 1,\\n        bias: bool = True,\\n        padding_mode: str = \\\"zeros\\\",  # TODO: refine this type\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        # we create new variables below to make mypy happy since kernel_size has\\n        # type Union[int, Tuple[int]] and kernel_size_ has type Tuple[int]\\n        kernel_size_ = _single(kernel_size)\\n        stride_ = _single(stride)\\n        padding_ = padding if isinstance(padding, str) else _single(padding)\\n        dilation_ = _single(dilation)\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size_,\\n            stride_,\\n            padding_,\\n            dilation_,\\n            False,\\n            _single(0),\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):\\n        if self.padding_mode != \\\"zeros\\\":\\n            return F.conv1d(\\n                F.pad(\\n                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode\\n                ),\\n                weight,\\n                bias,\\n                self.stride,\\n                _single(0),\\n                self.dilation,\\n                self.groups,\\n            )\\n        return F.conv1d(\\n            input, weight, bias, self.stride, self.padding, self.dilation, self.groups\\n        )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return self._conv_forward(input, self.weight, self.bias)\\n\\n\\nclass Conv2d(_ConvNd):\\n    __doc__ = (\\n        r\\\"\\\"\\\"Applies a 2D convolution over an input signal composed of several input\\n    planes.\\n\\n    In the simplest case, the output value of the layer with input size\\n    :math:`(N, C_{\\\\text{in}}, H, W)` and output :math:`(N, C_{\\\\text{out}}, H_{\\\\text{out}}, W_{\\\\text{out}})`\\n    can be precisely described as:\\n\\n    .. math::\\n        \\\\text{out}(N_i, C_{\\\\text{out}_j}) = \\\\text{bias}(C_{\\\\text{out}_j}) +\\n        \\\\sum_{k = 0}^{C_{\\\\text{in}} - 1} \\\\text{weight}(C_{\\\\text{out}_j}, k) \\\\star \\\\text{input}(N_i, k)\\n\\n\\n    where :math:`\\\\star` is the valid 2D `cross-correlation`_ operator,\\n    :math:`N` is a batch size, :math:`C` denotes a number of channels,\\n    :math:`H` is a height of input planes in pixels, and :math:`W` is\\n    width in pixels.\\n    \\\"\\\"\\\"\\n        + r\\\"\\\"\\\"\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    * :attr:`stride` controls the stride for the cross-correlation, a single\\n      number or a tuple.\\n\\n    * :attr:`padding` controls the amount of padding applied to the input. It\\n      can be either a string {{'valid', 'same'}} or an int / a tuple of ints giving the\\n      amount of implicit padding applied on both sides.\\n\\\"\\\"\\\"\\n        \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also\\n      known as the \\\\u00e0 trous algorithm. It is harder to describe, but this `link`_\\n      has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\"\\n        r\\\"\\\"\\\"\\n\\n    {groups_note}\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the height and width dimension\\n        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,\\n          and the second `int` for the width dimension\\n\\n    Note:\\n        {depthwise_separable_note}\\n\\n    Note:\\n        {cudnn_reproducibility_note}\\n\\n    Note:\\n        ``padding='valid'`` is the same as no padding. ``padding='same'`` pads\\n        the input so the output has the shape as the input. However, this mode\\n        doesn't support any stride values other than 1.\\n\\n    Note:\\n        This module supports complex data types i.e. ``complex32, complex64, complex128``.\\n\\n    Args:\\n        in_channels (int): Number of channels in the input image\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int, tuple or str, optional): Padding added to all four sides of\\n            the input. Default: 0\\n        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,\\n            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input\\n            channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the\\n            output. Default: ``True``\\n    \\\"\\\"\\\".format(\\n            **reproducibility_notes, **convolution_notes\\n        )\\n        + r\\\"\\\"\\\"\\n\\n    Shape:\\n        - Input: :math:`(N, C_{in}, H_{in}, W_{in})` or :math:`(C_{in}, H_{in}, W_{in})`\\n        - Output: :math:`(N, C_{out}, H_{out}, W_{out})` or :math:`(C_{out}, H_{out}, W_{out})`, where\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in}  + 2 \\\\times \\\\text{padding}[0] - \\\\text{dilation}[0]\\n                        \\\\times (\\\\text{kernel\\\\_size}[0] - 1) - 1}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in}  + 2 \\\\times \\\\text{padding}[1] - \\\\text{dilation}[1]\\n                        \\\\times (\\\\text{kernel\\\\_size}[1] - 1) - 1}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape\\n            :math:`(\\\\text{out\\\\_channels}, \\\\frac{\\\\text{in\\\\_channels}}{\\\\text{groups}},`\\n            :math:`\\\\text{kernel\\\\_size[0]}, \\\\text{kernel\\\\_size[1]})`.\\n            The values of these weights are sampled from\\n            :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n            :math:`k = \\\\frac{groups}{C_\\\\text{in} * \\\\prod_{i=0}^{1}\\\\text{kernel\\\\_size}[i]}`\\n        bias (Tensor):   the learnable bias of the module of shape\\n            (out_channels). If :attr:`bias` is ``True``,\\n            then the values of these weights are\\n            sampled from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n            :math:`k = \\\\frac{groups}{C_\\\\text{in} * \\\\prod_{i=0}^{1}\\\\text{kernel\\\\_size}[i]}`\\n\\n    Examples:\\n\\n        >>> # With square kernels and equal stride\\n        >>> m = nn.Conv2d(16, 33, 3, stride=2)\\n        >>> # non-square kernels and unequal stride and with padding\\n        >>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))\\n        >>> # non-square kernels and unequal stride and with padding and dilation\\n        >>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))\\n        >>> input = torch.randn(20, 16, 50, 100)\\n        >>> output = m(input)\\n\\n    .. _cross-correlation:\\n        https://en.wikipedia.org/wiki/Cross-correlation\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n    \\\"\\\"\\\"\\n    )\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: _size_2_t,\\n        stride: _size_2_t = 1,\\n        padding: Union[str, _size_2_t] = 0,\\n        dilation: _size_2_t = 1,\\n        groups: int = 1,\\n        bias: bool = True,\\n        padding_mode: str = \\\"zeros\\\",  # TODO: refine this type\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        kernel_size_ = _pair(kernel_size)\\n        stride_ = _pair(stride)\\n        padding_ = padding if isinstance(padding, str) else _pair(padding)\\n        dilation_ = _pair(dilation)\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size_,\\n            stride_,\\n            padding_,\\n            dilation_,\\n            False,\\n            _pair(0),\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):\\n        if self.padding_mode != \\\"zeros\\\":\\n            return F.conv2d(\\n                F.pad(\\n                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode\\n                ),\\n                weight,\\n                bias,\\n                self.stride,\\n                _pair(0),\\n                self.dilation,\\n                self.groups,\\n            )\\n        return F.conv2d(\\n            input, weight, bias, self.stride, self.padding, self.dilation, self.groups\\n        )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return self._conv_forward(input, self.weight, self.bias)\\n\\n\\nclass Conv3d(_ConvNd):\\n    __doc__ = (\\n        r\\\"\\\"\\\"Applies a 3D convolution over an input signal composed of several input\\n    planes.\\n\\n    In the simplest case, the output value of the layer with input size :math:`(N, C_{in}, D, H, W)`\\n    and output :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})` can be precisely described as:\\n\\n    .. math::\\n        out(N_i, C_{out_j}) = bias(C_{out_j}) +\\n                                \\\\sum_{k = 0}^{C_{in} - 1} weight(C_{out_j}, k) \\\\star input(N_i, k)\\n\\n    where :math:`\\\\star` is the valid 3D `cross-correlation`_ operator\\n    \\\"\\\"\\\"\\n        + r\\\"\\\"\\\"\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    * :attr:`stride` controls the stride for the cross-correlation.\\n\\n    * :attr:`padding` controls the amount of padding applied to the input. It\\n      can be either a string {{'valid', 'same'}} or a tuple of ints giving the\\n      amount of implicit padding applied on both sides.\\n\\\"\\\"\\\"\\n        \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also known as the \\\\u00e0 trous algorithm.\\n      It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\"\\n        r\\\"\\\"\\\"\\n\\n    {groups_note}\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the depth, height and width dimension\\n        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,\\n          the second `int` for the height dimension and the third `int` for the width dimension\\n\\n    Note:\\n        {depthwise_separable_note}\\n\\n    Note:\\n        {cudnn_reproducibility_note}\\n\\n    Note:\\n        ``padding='valid'`` is the same as no padding. ``padding='same'`` pads\\n        the input so the output has the shape as the input. However, this mode\\n        doesn't support any stride values other than 1.\\n\\n    Note:\\n        This module supports complex data types i.e. ``complex32, complex64, complex128``.\\n\\n    Args:\\n        in_channels (int): Number of channels in the input image\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int, tuple or str, optional): Padding added to all six sides of\\n            the input. Default: 0\\n        padding_mode (str, optional): ``'zeros'``, ``'reflect'``, ``'replicate'`` or ``'circular'``. Default: ``'zeros'``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n    \\\"\\\"\\\".format(\\n            **reproducibility_notes, **convolution_notes\\n        )\\n        + r\\\"\\\"\\\"\\n\\n    Shape:\\n        - Input: :math:`(N, C_{in}, D_{in}, H_{in}, W_{in})` or :math:`(C_{in}, D_{in}, H_{in}, W_{in})`\\n        - Output: :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})` or :math:`(C_{out}, D_{out}, H_{out}, W_{out})`,\\n          where\\n\\n          .. math::\\n              D_{out} = \\\\left\\\\lfloor\\\\frac{D_{in} + 2 \\\\times \\\\text{padding}[0] - \\\\text{dilation}[0]\\n                    \\\\times (\\\\text{kernel\\\\_size}[0] - 1) - 1}{\\\\text{stride}[0]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              H_{out} = \\\\left\\\\lfloor\\\\frac{H_{in} + 2 \\\\times \\\\text{padding}[1] - \\\\text{dilation}[1]\\n                    \\\\times (\\\\text{kernel\\\\_size}[1] - 1) - 1}{\\\\text{stride}[1]} + 1\\\\right\\\\rfloor\\n\\n          .. math::\\n              W_{out} = \\\\left\\\\lfloor\\\\frac{W_{in} + 2 \\\\times \\\\text{padding}[2] - \\\\text{dilation}[2]\\n                    \\\\times (\\\\text{kernel\\\\_size}[2] - 1) - 1}{\\\\text{stride}[2]} + 1\\\\right\\\\rfloor\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape\\n                         :math:`(\\\\text{out\\\\_channels}, \\\\frac{\\\\text{in\\\\_channels}}{\\\\text{groups}},`\\n                         :math:`\\\\text{kernel\\\\_size[0]}, \\\\text{kernel\\\\_size[1]}, \\\\text{kernel\\\\_size[2]})`.\\n                         The values of these weights are sampled from\\n                         :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{in} * \\\\prod_{i=0}^{2}\\\\text{kernel\\\\_size}[i]}`\\n        bias (Tensor):   the learnable bias of the module of shape (out_channels). If :attr:`bias` is ``True``,\\n                         then the values of these weights are\\n                         sampled from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{in} * \\\\prod_{i=0}^{2}\\\\text{kernel\\\\_size}[i]}`\\n\\n    Examples::\\n\\n        >>> # With square kernels and equal stride\\n        >>> m = nn.Conv3d(16, 33, 3, stride=2)\\n        >>> # non-square kernels and unequal stride and with padding\\n        >>> m = nn.Conv3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(4, 2, 0))\\n        >>> input = torch.randn(20, 16, 10, 50, 100)\\n        >>> output = m(input)\\n\\n    .. _cross-correlation:\\n        https://en.wikipedia.org/wiki/Cross-correlation\\n\\n    .. _link:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n    \\\"\\\"\\\"\\n    )\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: _size_3_t,\\n        stride: _size_3_t = 1,\\n        padding: Union[str, _size_3_t] = 0,\\n        dilation: _size_3_t = 1,\\n        groups: int = 1,\\n        bias: bool = True,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        kernel_size_ = _triple(kernel_size)\\n        stride_ = _triple(stride)\\n        padding_ = padding if isinstance(padding, str) else _triple(padding)\\n        dilation_ = _triple(dilation)\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size_,\\n            stride_,\\n            padding_,\\n            dilation_,\\n            False,\\n            _triple(0),\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):\\n        if self.padding_mode != \\\"zeros\\\":\\n            return F.conv3d(\\n                F.pad(\\n                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode\\n                ),\\n                weight,\\n                bias,\\n                self.stride,\\n                _triple(0),\\n                self.dilation,\\n                self.groups,\\n            )\\n        return F.conv3d(\\n            input, weight, bias, self.stride, self.padding, self.dilation, self.groups\\n        )\\n\\n    def forward(self, input: Tensor) -> Tensor:\\n        return self._conv_forward(input, self.weight, self.bias)\\n\\n\\nclass _ConvTransposeNd(_ConvNd):\\n    def __init__(\\n        self,\\n        in_channels,\\n        out_channels,\\n        kernel_size,\\n        stride,\\n        padding,\\n        dilation,\\n        transposed,\\n        output_padding,\\n        groups,\\n        bias,\\n        padding_mode,\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        if padding_mode != \\\"zeros\\\":\\n            raise ValueError(\\n                f'Only \\\"zeros\\\" padding mode is supported for {self.__class__.__name__}'\\n            )\\n\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            transposed,\\n            output_padding,\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    # dilation being an optional parameter is for backwards\\n    # compatibility\\n    def _output_padding(\\n        self,\\n        input: Tensor,\\n        output_size: Optional[List[int]],\\n        stride: List[int],\\n        padding: List[int],\\n        kernel_size: List[int],\\n        num_spatial_dims: int,\\n        dilation: Optional[List[int]] = None,\\n    ) -> List[int]:\\n        if output_size is None:\\n            ret = _single(self.output_padding)  # converting to list if was not already\\n        else:\\n            has_batch_dim = input.dim() == num_spatial_dims + 2\\n            num_non_spatial_dims = 2 if has_batch_dim else 1\\n            if len(output_size) == num_non_spatial_dims + num_spatial_dims:\\n                output_size = output_size[num_non_spatial_dims:]\\n            if len(output_size) != num_spatial_dims:\\n                raise ValueError(\\n                    f\\\"ConvTranspose{num_spatial_dims}D: for {input.dim()}D input, output_size must have {num_spatial_dims} \\\"\\n                    f\\\"or {num_non_spatial_dims + num_spatial_dims} elements (got {len(output_size)})\\\"\\n                )\\n\\n            min_sizes = torch.jit.annotate(List[int], [])\\n            max_sizes = torch.jit.annotate(List[int], [])\\n            for d in range(num_spatial_dims):\\n                dim_size = (\\n                    (input.size(d + num_non_spatial_dims) - 1) * stride[d]\\n                    - 2 * padding[d]\\n                    + (dilation[d] if dilation is not None else 1)\\n                    * (kernel_size[d] - 1)\\n                    + 1\\n                )\\n                min_sizes.append(dim_size)\\n                max_sizes.append(min_sizes[d] + stride[d] - 1)\\n\\n            for i in range(len(output_size)):\\n                size = output_size[i]\\n                min_size = min_sizes[i]\\n                max_size = max_sizes[i]\\n                if size < min_size or size > max_size:\\n                    raise ValueError(\\n                        f\\\"requested an output size of {output_size}, but valid sizes range \\\"\\n                        f\\\"from {min_sizes} to {max_sizes} (for an input of {input.size()[2:]})\\\"\\n                    )\\n\\n            res = torch.jit.annotate(List[int], [])\\n            for d in range(num_spatial_dims):\\n                res.append(output_size[d] - min_sizes[d])\\n\\n            ret = res\\n        return ret\\n\\n\\nclass ConvTranspose1d(_ConvTransposeNd):\\n    __doc__ = (\\n        r\\\"\\\"\\\"Applies a 1D transposed convolution operator over an input image\\n    composed of several input planes.\\n\\n    This module can be seen as the gradient of Conv1d with respect to its input.\\n    It is also known as a fractionally-strided convolution or\\n    a deconvolution (although it is not an actual deconvolution operation as it does\\n    not compute a true inverse of convolution). For more information, see the visualizations\\n    `here`_ and the `Deconvolutional Networks`_ paper.\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    * :attr:`stride` controls the stride for the cross-correlation.\\n\\n    * :attr:`padding` controls the amount of implicit zero padding on both\\n      sides for ``dilation * (kernel_size - 1) - padding`` number of points. See note\\n      below for details.\\n\\n    * :attr:`output_padding` controls the additional size added to one side\\n      of the output shape. See note below for details.\\n\\\"\\\"\\\"\\n        \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also known as the \\\\u00e0 trous algorithm.\\n      It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\"\\n        r\\\"\\\"\\\"\\n    {groups_note}\\n\\n    Note:\\n        The :attr:`padding` argument effectively adds ``dilation * (kernel_size - 1) - padding``\\n        amount of zero padding to both sizes of the input. This is set so that\\n        when a :class:`~torch.nn.Conv1d` and a :class:`~torch.nn.ConvTranspose1d`\\n        are initialized with same parameters, they are inverses of each other in\\n        regard to the input and output shapes. However, when ``stride > 1``,\\n        :class:`~torch.nn.Conv1d` maps multiple input shapes to the same output\\n        shape. :attr:`output_padding` is provided to resolve this ambiguity by\\n        effectively increasing the calculated output shape on one side. Note\\n        that :attr:`output_padding` is only used to find output shape, but does\\n        not actually add zero-padding to output.\\n\\n    Note:\\n        In some circumstances when using the CUDA backend with CuDNN, this operator\\n        may select a nondeterministic algorithm to increase performance. If this is\\n        undesirable, you can try to make the operation deterministic (potentially at\\n        a performance cost) by setting ``torch.backends.cudnn.deterministic =\\n        True``.\\n        Please see the notes on :doc:`/notes/randomness` for background.\\n\\n\\n    Args:\\n        in_channels (int): Number of channels in the input image\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding\\n            will be added to both sides of the input. Default: 0\\n        output_padding (int or tuple, optional): Additional size added to one side\\n            of the output shape. Default: 0\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n    \\\"\\\"\\\".format(\\n            **reproducibility_notes, **convolution_notes\\n        )\\n        + r\\\"\\\"\\\"\\n\\n    Shape:\\n        - Input: :math:`(N, C_{in}, L_{in})` or :math:`(C_{in}, L_{in})`\\n        - Output: :math:`(N, C_{out}, L_{out})` or :math:`(C_{out}, L_{out})`, where\\n\\n          .. math::\\n              L_{out} = (L_{in} - 1) \\\\times \\\\text{stride} - 2 \\\\times \\\\text{padding} + \\\\text{dilation}\\n                        \\\\times (\\\\text{kernel\\\\_size} - 1) + \\\\text{output\\\\_padding} + 1\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape\\n                         :math:`(\\\\text{in\\\\_channels}, \\\\frac{\\\\text{out\\\\_channels}}{\\\\text{groups}},`\\n                         :math:`\\\\text{kernel\\\\_size})`.\\n                         The values of these weights are sampled from\\n                         :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{out} * \\\\text{kernel\\\\_size}}`\\n        bias (Tensor):   the learnable bias of the module of shape (out_channels).\\n                         If :attr:`bias` is ``True``, then the values of these weights are\\n                         sampled from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{out} * \\\\text{kernel\\\\_size}}`\\n\\n    .. _`here`:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n\\n    .. _`Deconvolutional Networks`:\\n        https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf\\n    \\\"\\\"\\\"\\n    )\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: _size_1_t,\\n        stride: _size_1_t = 1,\\n        padding: _size_1_t = 0,\\n        output_padding: _size_1_t = 0,\\n        groups: int = 1,\\n        bias: bool = True,\\n        dilation: _size_1_t = 1,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        kernel_size = _single(kernel_size)\\n        stride = _single(stride)\\n        padding = _single(padding)\\n        dilation = _single(dilation)\\n        output_padding = _single(output_padding)\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            True,\\n            output_padding,\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    def forward(self, input: Tensor, output_size: Optional[List[int]] = None) -> Tensor:\\n        if self.padding_mode != \\\"zeros\\\":\\n            raise ValueError(\\n                \\\"Only `zeros` padding mode is supported for ConvTranspose1d\\\"\\n            )\\n\\n        assert isinstance(self.padding, tuple)\\n        # One cannot replace List by Tuple or Sequence in \\\"_output_padding\\\" because\\n        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.\\n        num_spatial_dims = 1\\n        output_padding = self._output_padding(\\n            input,\\n            output_size,\\n            self.stride,  # type: ignore[arg-type]\\n            self.padding,  # type: ignore[arg-type]\\n            self.kernel_size,  # type: ignore[arg-type]\\n            num_spatial_dims,\\n            self.dilation,  # type: ignore[arg-type]\\n        )\\n        return F.conv_transpose1d(\\n            input,\\n            self.weight,\\n            self.bias,\\n            self.stride,\\n            self.padding,\\n            output_padding,\\n            self.groups,\\n            self.dilation,\\n        )\\n\\n\\nclass ConvTranspose2d(_ConvTransposeNd):\\n    __doc__ = (\\n        r\\\"\\\"\\\"Applies a 2D transposed convolution operator over an input image\\n    composed of several input planes.\\n\\n    This module can be seen as the gradient of Conv2d with respect to its input.\\n    It is also known as a fractionally-strided convolution or\\n    a deconvolution (although it is not an actual deconvolution operation as it does\\n    not compute a true inverse of convolution). For more information, see the visualizations\\n    `here`_ and the `Deconvolutional Networks`_ paper.\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    * :attr:`stride` controls the stride for the cross-correlation.\\n\\n    * :attr:`padding` controls the amount of implicit zero padding on both\\n      sides for ``dilation * (kernel_size - 1) - padding`` number of points. See note\\n      below for details.\\n\\n    * :attr:`output_padding` controls the additional size added to one side\\n      of the output shape. See note below for details.\\n\\\"\\\"\\\"\\n        \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also known as the \\\\u00e0 trous algorithm.\\n      It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\"\\n        r\\\"\\\"\\\"\\n    {groups_note}\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`output_padding`\\n    can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the height and width dimensions\\n        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,\\n          and the second `int` for the width dimension\\n\\n    Note:\\n        The :attr:`padding` argument effectively adds ``dilation * (kernel_size - 1) - padding``\\n        amount of zero padding to both sizes of the input. This is set so that\\n        when a :class:`~torch.nn.Conv2d` and a :class:`~torch.nn.ConvTranspose2d`\\n        are initialized with same parameters, they are inverses of each other in\\n        regard to the input and output shapes. However, when ``stride > 1``,\\n        :class:`~torch.nn.Conv2d` maps multiple input shapes to the same output\\n        shape. :attr:`output_padding` is provided to resolve this ambiguity by\\n        effectively increasing the calculated output shape on one side. Note\\n        that :attr:`output_padding` is only used to find output shape, but does\\n        not actually add zero-padding to output.\\n\\n    Note:\\n        {cudnn_reproducibility_note}\\n\\n    Args:\\n        in_channels (int): Number of channels in the input image\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding\\n            will be added to both sides of each dimension in the input. Default: 0\\n        output_padding (int or tuple, optional): Additional size added to one side\\n            of each dimension in the output shape. Default: 0\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n    \\\"\\\"\\\".format(\\n            **reproducibility_notes, **convolution_notes\\n        )\\n        + r\\\"\\\"\\\"\\n\\n    Shape:\\n        - Input: :math:`(N, C_{in}, H_{in}, W_{in})` or :math:`(C_{in}, H_{in}, W_{in})`\\n        - Output: :math:`(N, C_{out}, H_{out}, W_{out})` or :math:`(C_{out}, H_{out}, W_{out})`, where\\n\\n        .. math::\\n              H_{out} = (H_{in} - 1) \\\\times \\\\text{stride}[0] - 2 \\\\times \\\\text{padding}[0] + \\\\text{dilation}[0]\\n                        \\\\times (\\\\text{kernel\\\\_size}[0] - 1) + \\\\text{output\\\\_padding}[0] + 1\\n        .. math::\\n              W_{out} = (W_{in} - 1) \\\\times \\\\text{stride}[1] - 2 \\\\times \\\\text{padding}[1] + \\\\text{dilation}[1]\\n                        \\\\times (\\\\text{kernel\\\\_size}[1] - 1) + \\\\text{output\\\\_padding}[1] + 1\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape\\n                         :math:`(\\\\text{in\\\\_channels}, \\\\frac{\\\\text{out\\\\_channels}}{\\\\text{groups}},`\\n                         :math:`\\\\text{kernel\\\\_size[0]}, \\\\text{kernel\\\\_size[1]})`.\\n                         The values of these weights are sampled from\\n                         :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{out} * \\\\prod_{i=0}^{1}\\\\text{kernel\\\\_size}[i]}`\\n        bias (Tensor):   the learnable bias of the module of shape (out_channels)\\n                         If :attr:`bias` is ``True``, then the values of these weights are\\n                         sampled from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{out} * \\\\prod_{i=0}^{1}\\\\text{kernel\\\\_size}[i]}`\\n\\n    Examples::\\n\\n        >>> # With square kernels and equal stride\\n        >>> m = nn.ConvTranspose2d(16, 33, 3, stride=2)\\n        >>> # non-square kernels and unequal stride and with padding\\n        >>> m = nn.ConvTranspose2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))\\n        >>> input = torch.randn(20, 16, 50, 100)\\n        >>> output = m(input)\\n        >>> # exact output size can be also specified as an argument\\n        >>> input = torch.randn(1, 16, 12, 12)\\n        >>> downsample = nn.Conv2d(16, 16, 3, stride=2, padding=1)\\n        >>> upsample = nn.ConvTranspose2d(16, 16, 3, stride=2, padding=1)\\n        >>> h = downsample(input)\\n        >>> h.size()\\n        torch.Size([1, 16, 6, 6])\\n        >>> output = upsample(h, output_size=input.size())\\n        >>> output.size()\\n        torch.Size([1, 16, 12, 12])\\n\\n    .. _`here`:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n\\n    .. _`Deconvolutional Networks`:\\n        https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf\\n    \\\"\\\"\\\"\\n    )\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: _size_2_t,\\n        stride: _size_2_t = 1,\\n        padding: _size_2_t = 0,\\n        output_padding: _size_2_t = 0,\\n        groups: int = 1,\\n        bias: bool = True,\\n        dilation: _size_2_t = 1,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        kernel_size = _pair(kernel_size)\\n        stride = _pair(stride)\\n        padding = _pair(padding)\\n        dilation = _pair(dilation)\\n        output_padding = _pair(output_padding)\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            True,\\n            output_padding,\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    def forward(self, input: Tensor, output_size: Optional[List[int]] = None) -> Tensor:\\n        if self.padding_mode != \\\"zeros\\\":\\n            raise ValueError(\\n                \\\"Only `zeros` padding mode is supported for ConvTranspose2d\\\"\\n            )\\n\\n        assert isinstance(self.padding, tuple)\\n        # One cannot replace List by Tuple or Sequence in \\\"_output_padding\\\" because\\n        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.\\n        num_spatial_dims = 2\\n        output_padding = self._output_padding(\\n            input,\\n            output_size,\\n            self.stride,  # type: ignore[arg-type]\\n            self.padding,  # type: ignore[arg-type]\\n            self.kernel_size,  # type: ignore[arg-type]\\n            num_spatial_dims,\\n            self.dilation,  # type: ignore[arg-type]\\n        )\\n\\n        return F.conv_transpose2d(\\n            input,\\n            self.weight,\\n            self.bias,\\n            self.stride,\\n            self.padding,\\n            output_padding,\\n            self.groups,\\n            self.dilation,\\n        )\\n\\n\\nclass ConvTranspose3d(_ConvTransposeNd):\\n    __doc__ = (\\n        r\\\"\\\"\\\"Applies a 3D transposed convolution operator over an input image composed of several input\\n    planes.\\n    The transposed convolution operator multiplies each input value element-wise by a learnable kernel,\\n    and sums over the outputs from all input feature planes.\\n\\n    This module can be seen as the gradient of Conv3d with respect to its input.\\n    It is also known as a fractionally-strided convolution or\\n    a deconvolution (although it is not an actual deconvolution operation as it does\\n    not compute a true inverse of convolution). For more information, see the visualizations\\n    `here`_ and the `Deconvolutional Networks`_ paper.\\n\\n    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.\\n\\n    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.\\n\\n    * :attr:`stride` controls the stride for the cross-correlation.\\n\\n    * :attr:`padding` controls the amount of implicit zero padding on both\\n      sides for ``dilation * (kernel_size - 1) - padding`` number of points. See note\\n      below for details.\\n\\n    * :attr:`output_padding` controls the additional size added to one side\\n      of the output shape. See note below for details.\\n\\\"\\\"\\\"\\n        \\\"\\\"\\\"\\n    * :attr:`dilation` controls the spacing between the kernel points; also known as the \\\\u00e0 trous algorithm.\\n      It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does.\\n\\\"\\\"\\\"\\n        r\\\"\\\"\\\"\\n    {groups_note}\\n\\n    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`output_padding`\\n    can either be:\\n\\n        - a single ``int`` -- in which case the same value is used for the depth, height and width dimensions\\n        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,\\n          the second `int` for the height dimension and the third `int` for the width dimension\\n\\n    Note:\\n        The :attr:`padding` argument effectively adds ``dilation * (kernel_size - 1) - padding``\\n        amount of zero padding to both sizes of the input. This is set so that\\n        when a :class:`~torch.nn.Conv3d` and a :class:`~torch.nn.ConvTranspose3d`\\n        are initialized with same parameters, they are inverses of each other in\\n        regard to the input and output shapes. However, when ``stride > 1``,\\n        :class:`~torch.nn.Conv3d` maps multiple input shapes to the same output\\n        shape. :attr:`output_padding` is provided to resolve this ambiguity by\\n        effectively increasing the calculated output shape on one side. Note\\n        that :attr:`output_padding` is only used to find output shape, but does\\n        not actually add zero-padding to output.\\n\\n    Note:\\n        {cudnn_reproducibility_note}\\n\\n    Args:\\n        in_channels (int): Number of channels in the input image\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding\\n            will be added to both sides of each dimension in the input. Default: 0\\n        output_padding (int or tuple, optional): Additional size added to one side\\n            of each dimension in the output shape. Default: 0\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n    \\\"\\\"\\\".format(\\n            **reproducibility_notes, **convolution_notes\\n        )\\n        + r\\\"\\\"\\\"\\n\\n    Shape:\\n        - Input: :math:`(N, C_{in}, D_{in}, H_{in}, W_{in})` or :math:`(C_{in}, D_{in}, H_{in}, W_{in})`\\n        - Output: :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})` or\\n          :math:`(C_{out}, D_{out}, H_{out}, W_{out})`, where\\n\\n        .. math::\\n              D_{out} = (D_{in} - 1) \\\\times \\\\text{stride}[0] - 2 \\\\times \\\\text{padding}[0] + \\\\text{dilation}[0]\\n                        \\\\times (\\\\text{kernel\\\\_size}[0] - 1) + \\\\text{output\\\\_padding}[0] + 1\\n        .. math::\\n              H_{out} = (H_{in} - 1) \\\\times \\\\text{stride}[1] - 2 \\\\times \\\\text{padding}[1] + \\\\text{dilation}[1]\\n                        \\\\times (\\\\text{kernel\\\\_size}[1] - 1) + \\\\text{output\\\\_padding}[1] + 1\\n        .. math::\\n              W_{out} = (W_{in} - 1) \\\\times \\\\text{stride}[2] - 2 \\\\times \\\\text{padding}[2] + \\\\text{dilation}[2]\\n                        \\\\times (\\\\text{kernel\\\\_size}[2] - 1) + \\\\text{output\\\\_padding}[2] + 1\\n\\n\\n    Attributes:\\n        weight (Tensor): the learnable weights of the module of shape\\n                         :math:`(\\\\text{in\\\\_channels}, \\\\frac{\\\\text{out\\\\_channels}}{\\\\text{groups}},`\\n                         :math:`\\\\text{kernel\\\\_size[0]}, \\\\text{kernel\\\\_size[1]}, \\\\text{kernel\\\\_size[2]})`.\\n                         The values of these weights are sampled from\\n                         :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{out} * \\\\prod_{i=0}^{2}\\\\text{kernel\\\\_size}[i]}`\\n        bias (Tensor):   the learnable bias of the module of shape (out_channels)\\n                         If :attr:`bias` is ``True``, then the values of these weights are\\n                         sampled from :math:`\\\\mathcal{U}(-\\\\sqrt{k}, \\\\sqrt{k})` where\\n                         :math:`k = \\\\frac{groups}{C_\\\\text{out} * \\\\prod_{i=0}^{2}\\\\text{kernel\\\\_size}[i]}`\\n\\n    Examples::\\n\\n        >>> # With square kernels and equal stride\\n        >>> m = nn.ConvTranspose3d(16, 33, 3, stride=2)\\n        >>> # non-square kernels and unequal stride and with padding\\n        >>> m = nn.ConvTranspose3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(0, 4, 2))\\n        >>> input = torch.randn(20, 16, 10, 50, 100)\\n        >>> output = m(input)\\n\\n    .. _`here`:\\n        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\\n\\n    .. _`Deconvolutional Networks`:\\n        https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf\\n    \\\"\\\"\\\"\\n    )\\n\\n    def __init__(\\n        self,\\n        in_channels: int,\\n        out_channels: int,\\n        kernel_size: _size_3_t,\\n        stride: _size_3_t = 1,\\n        padding: _size_3_t = 0,\\n        output_padding: _size_3_t = 0,\\n        groups: int = 1,\\n        bias: bool = True,\\n        dilation: _size_3_t = 1,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        kernel_size = _triple(kernel_size)\\n        stride = _triple(stride)\\n        padding = _triple(padding)\\n        dilation = _triple(dilation)\\n        output_padding = _triple(output_padding)\\n        super().__init__(\\n            in_channels,\\n            out_channels,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            True,\\n            output_padding,\\n            groups,\\n            bias,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n\\n    def forward(self, input: Tensor, output_size: Optional[List[int]] = None) -> Tensor:\\n        if self.padding_mode != \\\"zeros\\\":\\n            raise ValueError(\\n                \\\"Only `zeros` padding mode is supported for ConvTranspose3d\\\"\\n            )\\n\\n        assert isinstance(self.padding, tuple)\\n        # One cannot replace List by Tuple or Sequence in \\\"_output_padding\\\" because\\n        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.\\n        num_spatial_dims = 3\\n        output_padding = self._output_padding(\\n            input,\\n            output_size,\\n            self.stride,  # type: ignore[arg-type]\\n            self.padding,  # type: ignore[arg-type]\\n            self.kernel_size,  # type: ignore[arg-type]\\n            num_spatial_dims,\\n            self.dilation,  # type: ignore[arg-type]\\n        )\\n\\n        return F.conv_transpose3d(\\n            input,\\n            self.weight,\\n            self.bias,\\n            self.stride,\\n            self.padding,\\n            output_padding,\\n            self.groups,\\n            self.dilation,\\n        )\\n\\n\\n# TODO: Deprecate and remove the following alias `_ConvTransposeMixin`.\\n#\\n# `_ConvTransposeMixin` was a mixin that was removed.  It is meant to be used\\n# with `_ConvNd` to construct actual module classes that implements conv\\n# transpose ops:\\n#\\n#   class MyConvTranspose(_ConvNd, _ConvTransposeMixin):\\n#       ...\\n#\\n# In PyTorch, it has been replaced by `_ConvTransposeNd`, which is a proper\\n# subclass of `_ConvNd`.  However, some user code in the wild still (incorrectly)\\n# use the internal class `_ConvTransposeMixin`.  Hence, we provide this alias\\n# for BC, because it is cheap and easy for us to do so, even though that\\n# `_ConvTransposeNd` is really not a mixin anymore (but multiple inheritance as\\n# above would still work).\\nclass _ConvTransposeMixin(_ConvTransposeNd):\\n    @deprecated(\\n        \\\"`_ConvTransposeMixin` is a deprecated internal class. \\\"\\n        \\\"Please consider using public APIs.\\\",\\n        category=FutureWarning,\\n    )\\n    def __init__(self, *args, **kwargs):\\n        super().__init__(*args, **kwargs)\\n\\n\\n# TODO: Conv2dLocal\\n# TODO: Conv2dMap\\n# TODO: ConvTranspose2dMap\\n\\n\\nclass _LazyConvXdMixin(LazyModuleMixin):\\n    groups: int\\n    transposed: bool\\n    in_channels: int\\n    out_channels: int\\n    kernel_size: Tuple[int, ...]\\n    weight: UninitializedParameter\\n    bias: UninitializedParameter\\n\\n    def reset_parameters(self) -> None:\\n        # has_uninitialized_params is defined in parent class and it is using a protocol on self\\n        if not self.has_uninitialized_params() and self.in_channels != 0:  # type: ignore[misc]\\n            # \\\"type:ignore[..]\\\" is required because mypy thinks that \\\"reset_parameters\\\" is undefined\\n            # in super class. Turns out that it is defined in _ConvND which is inherited by any class\\n            # that also inherits _LazyConvXdMixin\\n            super().reset_parameters()  # type: ignore[misc]\\n\\n    # Signature of \\\"initialize_parameters\\\" is incompatible with the definition in supertype LazyModuleMixin\\n    def initialize_parameters(self, input: Tensor, *args, **kwargs) -> None:  # type: ignore[override]\\n        # defined by parent class but using a protocol\\n        if self.has_uninitialized_params():  # type: ignore[misc]\\n            self.in_channels = self._get_in_channels(input)\\n            if self.in_channels % self.groups != 0:\\n                raise ValueError(\\\"in_channels must be divisible by groups\\\")\\n            assert isinstance(self.weight, UninitializedParameter)\\n            if self.transposed:\\n                self.weight.materialize(\\n                    (\\n                        self.in_channels,\\n                        self.out_channels // self.groups,\\n                        *self.kernel_size,\\n                    )\\n                )\\n            else:\\n                self.weight.materialize(\\n                    (\\n                        self.out_channels,\\n                        self.in_channels // self.groups,\\n                        *self.kernel_size,\\n                    )\\n                )\\n            if self.bias is not None:\\n                assert isinstance(self.bias, UninitializedParameter)\\n                self.bias.materialize((self.out_channels,))\\n            self.reset_parameters()\\n\\n    # Function to extract in_channels from first input.\\n    def _get_in_channels(self, input: Tensor) -> int:\\n        num_spatial_dims = self._get_num_spatial_dims()\\n        num_dims_no_batch = num_spatial_dims + 1  # +1 for channels dim\\n        num_dims_batch = num_dims_no_batch + 1\\n        if input.dim() not in (num_dims_no_batch, num_dims_batch):\\n            raise RuntimeError(\\n                f\\\"Expected {num_dims_no_batch}D (unbatched) or {num_dims_batch}D (batched) input \\\"\\n                f\\\"to {self.__class__.__name__}, but \\\"\\n                f\\\"got input of size: {input.shape}\\\"\\n            )\\n        return input.shape[1] if input.dim() == num_dims_batch else input.shape[0]\\n\\n    # Function to return the number of spatial dims expected for inputs to the module.\\n    # This is expected to be implemented by subclasses.\\n    def _get_num_spatial_dims(self) -> int:\\n        raise NotImplementedError\\n\\n\\n# LazyConv1d defines weight as a Tensor but derived class defines it as UnitializeParameter\\nclass LazyConv1d(_LazyConvXdMixin, Conv1d):  # type: ignore[misc]\\n    r\\\"\\\"\\\"A :class:`torch.nn.Conv1d` module with lazy initialization of the ``in_channels`` argument.\\n\\n    The ``in_channels`` argument of the :class:`Conv1d` is inferred from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight` and `bias`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of\\n            the input. Default: 0\\n        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,\\n            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``\\n        dilation (int or tuple, optional): Spacing between kernel\\n            elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input\\n            channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the\\n            output. Default: ``True``\\n\\n    .. seealso:: :class:`torch.nn.Conv1d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`\\n    \\\"\\\"\\\"\\n\\n    # super class define this variable as None. \\\"type: ignore[..] is required\\n    # since we are redefining the variable.\\n    cls_to_become = Conv1d  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        out_channels: int,\\n        kernel_size: _size_1_t,\\n        stride: _size_1_t = 1,\\n        padding: _size_1_t = 0,\\n        dilation: _size_1_t = 1,\\n        groups: int = 1,\\n        bias: bool = True,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            0,\\n            0,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            groups,\\n            # bias is hardcoded to False to avoid creating tensor\\n            # that will soon be overwritten.\\n            False,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_channels = out_channels\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def _get_num_spatial_dims(self) -> int:\\n        return 1\\n\\n\\n# LazyConv2d defines weight as a Tensor but derived class defines it as UnitializeParameter\\nclass LazyConv2d(_LazyConvXdMixin, Conv2d):  # type: ignore[misc]\\n    r\\\"\\\"\\\"A :class:`torch.nn.Conv2d` module with lazy initialization of the ``in_channels`` argument.\\n\\n    The ``in_channels`` argument of the :class:`Conv2d` that is inferred from the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight` and `bias`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of\\n            the input. Default: 0\\n        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,\\n            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``\\n        dilation (int or tuple, optional): Spacing between kernel\\n            elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input\\n            channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the\\n            output. Default: ``True``\\n\\n    .. seealso:: :class:`torch.nn.Conv2d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`\\n    \\\"\\\"\\\"\\n\\n    # super class define this variable as None. \\\"type: ignore[..] is required\\n    # since we are redefining the variable.\\n    cls_to_become = Conv2d  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        out_channels: int,\\n        kernel_size: _size_2_t,\\n        stride: _size_2_t = 1,\\n        padding: _size_2_t = 0,\\n        dilation: _size_2_t = 1,\\n        groups: int = 1,\\n        bias: bool = True,\\n        padding_mode: str = \\\"zeros\\\",  # TODO: refine this type\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            0,\\n            0,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            groups,\\n            # bias is hardcoded to False to avoid creating tensor\\n            # that will soon be overwritten.\\n            False,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_channels = out_channels\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def _get_num_spatial_dims(self) -> int:\\n        return 2\\n\\n\\n# LazyConv3d defines weight as a Tensor but derived class defines it as UnitializeParameter\\nclass LazyConv3d(_LazyConvXdMixin, Conv3d):  # type: ignore[misc]\\n    r\\\"\\\"\\\"A :class:`torch.nn.Conv3d` module with lazy initialization of the ``in_channels`` argument.\\n\\n    The ``in_channels`` argument of the :class:`Conv3d` that is inferred from\\n    the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight` and `bias`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): Zero-padding added to both sides of\\n            the input. Default: 0\\n        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,\\n            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``\\n        dilation (int or tuple, optional): Spacing between kernel\\n            elements. Default: 1\\n        groups (int, optional): Number of blocked connections from input\\n            channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the\\n            output. Default: ``True``\\n\\n    .. seealso:: :class:`torch.nn.Conv3d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`\\n    \\\"\\\"\\\"\\n\\n    # super class define this variable as None. \\\"type: ignore[..] is required\\n    # since we are redefining the variable.\\n    cls_to_become = Conv3d  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        out_channels: int,\\n        kernel_size: _size_3_t,\\n        stride: _size_3_t = 1,\\n        padding: _size_3_t = 0,\\n        dilation: _size_3_t = 1,\\n        groups: int = 1,\\n        bias: bool = True,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            0,\\n            0,\\n            kernel_size,\\n            stride,\\n            padding,\\n            dilation,\\n            groups,\\n            # bias is hardcoded to False to avoid creating tensor\\n            # that will soon be overwritten.\\n            False,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_channels = out_channels\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def _get_num_spatial_dims(self) -> int:\\n        return 3\\n\\n\\n# LazyConvTranspose1d defines weight as a Tensor but derived class defines it as UnitializeParameter\\nclass LazyConvTranspose1d(_LazyConvXdMixin, ConvTranspose1d):  # type: ignore[misc]\\n    r\\\"\\\"\\\"A :class:`torch.nn.ConvTranspose1d` module with lazy initialization of the ``in_channels`` argument.\\n\\n    The ``in_channels`` argument of the :class:`ConvTranspose1d` that is inferred from\\n    the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight` and `bias`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding\\n            will be added to both sides of the input. Default: 0\\n        output_padding (int or tuple, optional): Additional size added to one side\\n            of the output shape. Default: 0\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n\\n    .. seealso:: :class:`torch.nn.ConvTranspose1d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`\\n    \\\"\\\"\\\"\\n\\n    # super class define this variable as None. \\\"type: ignore[..] is required\\n    # since we are redefining the variable.\\n    cls_to_become = ConvTranspose1d  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        out_channels: int,\\n        kernel_size: _size_1_t,\\n        stride: _size_1_t = 1,\\n        padding: _size_1_t = 0,\\n        output_padding: _size_1_t = 0,\\n        groups: int = 1,\\n        bias: bool = True,\\n        dilation: _size_1_t = 1,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            0,\\n            0,\\n            kernel_size,\\n            stride,\\n            padding,\\n            output_padding,\\n            groups,\\n            # bias is hardcoded to False to avoid creating tensor\\n            # that will soon be overwritten.\\n            False,\\n            dilation,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_channels = out_channels\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def _get_num_spatial_dims(self) -> int:\\n        return 1\\n\\n\\n# LazyConvTranspose2d defines weight as a Tensor but derived class defines it as UnitializeParameter\\nclass LazyConvTranspose2d(_LazyConvXdMixin, ConvTranspose2d):  # type: ignore[misc]\\n    r\\\"\\\"\\\"A :class:`torch.nn.ConvTranspose2d` module with lazy initialization of the ``in_channels`` argument.\\n\\n    The ``in_channels`` argument of the :class:`ConvTranspose2d` is inferred from\\n    the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight` and `bias`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding\\n            will be added to both sides of each dimension in the input. Default: 0\\n        output_padding (int or tuple, optional): Additional size added to one side\\n            of each dimension in the output shape. Default: 0\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n\\n    .. seealso:: :class:`torch.nn.ConvTranspose2d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`\\n    \\\"\\\"\\\"\\n\\n    # super class define this variable as None. \\\"type: ignore[..] is required\\n    # since we are redefining the variable.\\n    cls_to_become = ConvTranspose2d  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        out_channels: int,\\n        kernel_size: _size_2_t,\\n        stride: _size_2_t = 1,\\n        padding: _size_2_t = 0,\\n        output_padding: _size_2_t = 0,\\n        groups: int = 1,\\n        bias: bool = True,\\n        dilation: int = 1,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            0,\\n            0,\\n            kernel_size,\\n            stride,\\n            padding,\\n            output_padding,\\n            groups,\\n            # bias is hardcoded to False to avoid creating tensor\\n            # that will soon be overwritten.\\n            False,\\n            dilation,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_channels = out_channels\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def _get_num_spatial_dims(self) -> int:\\n        return 2\\n\\n\\n# LazyConvTranspose3d defines weight as a Tensor but derived class defines it as UnitializeParameter\\nclass LazyConvTranspose3d(_LazyConvXdMixin, ConvTranspose3d):  # type: ignore[misc]\\n    r\\\"\\\"\\\"A :class:`torch.nn.ConvTranspose3d` module with lazy initialization of the ``in_channels`` argument.\\n\\n    The ``in_channels`` argument of the :class:`ConvTranspose3d` is inferred from\\n    the ``input.size(1)``.\\n    The attributes that will be lazily initialized are `weight` and `bias`.\\n\\n    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation\\n    on lazy modules and their limitations.\\n\\n    Args:\\n        out_channels (int): Number of channels produced by the convolution\\n        kernel_size (int or tuple): Size of the convolving kernel\\n        stride (int or tuple, optional): Stride of the convolution. Default: 1\\n        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding\\n            will be added to both sides of each dimension in the input. Default: 0\\n        output_padding (int or tuple, optional): Additional size added to one side\\n            of each dimension in the output shape. Default: 0\\n        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\\n        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\\n        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\\n\\n    .. seealso:: :class:`torch.nn.ConvTranspose3d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`\\n    \\\"\\\"\\\"\\n\\n    # super class define this variable as None. \\\"type: ignore[..] is required\\n    # since we are redefining the variable.\\n    cls_to_become = ConvTranspose3d  # type: ignore[assignment]\\n\\n    def __init__(\\n        self,\\n        out_channels: int,\\n        kernel_size: _size_3_t,\\n        stride: _size_3_t = 1,\\n        padding: _size_3_t = 0,\\n        output_padding: _size_3_t = 0,\\n        groups: int = 1,\\n        bias: bool = True,\\n        dilation: _size_3_t = 1,\\n        padding_mode: str = \\\"zeros\\\",\\n        device=None,\\n        dtype=None,\\n    ) -> None:\\n        factory_kwargs = {\\\"device\\\": device, \\\"dtype\\\": dtype}\\n        super().__init__(\\n            0,\\n            0,\\n            kernel_size,\\n            stride,\\n            padding,\\n            output_padding,\\n            groups,\\n            # bias is hardcoded to False to avoid creating tensor\\n            # that will soon be overwritten.\\n            False,\\n            dilation,\\n            padding_mode,\\n            **factory_kwargs,\\n        )\\n        self.weight = UninitializedParameter(**factory_kwargs)\\n        self.out_channels = out_channels\\n        if bias:\\n            self.bias = UninitializedParameter(**factory_kwargs)\\n\\n    def _get_num_spatial_dims(self) -> int:\\n        return 3\\n\\n\\nfrom .module import Module  # usort: skip\\nfrom .linear import Bilinear, Identity, LazyLinear, Linear  # usort: skip\\nfrom .activation import (\\n    CELU,\\n    ELU,\\n    GELU,\\n    GLU,\\n    Hardshrink,\\n    Hardsigmoid,\\n    Hardswish,\\n    Hardtanh,\\n    LeakyReLU,\\n    LogSigmoid,\\n    LogSoftmax,\\n    Mish,\\n    MultiheadAttention,\\n    PReLU,\\n    ReLU,\\n    ReLU6,\\n    RReLU,\\n    SELU,\\n    Sigmoid,\\n    SiLU,\\n    Softmax,\\n    Softmax2d,\\n    Softmin,\\n    Softplus,\\n    Softshrink,\\n    Softsign,\\n    Tanh,\\n    Tanhshrink,\\n    Threshold,\\n)\\nfrom .adaptive import AdaptiveLogSoftmaxWithLoss\\nfrom .batchnorm import (\\n    BatchNorm1d,\\n    BatchNorm2d,\\n    BatchNorm3d,\\n    LazyBatchNorm1d,\\n    LazyBatchNorm2d,\\n    LazyBatchNorm3d,\\n    SyncBatchNorm,\\n)\\nfrom .channelshuffle import ChannelShuffle\\nfrom .container import (\\n    Container,\\n    ModuleDict,\\n    ModuleList,\\n    ParameterDict,\\n    ParameterList,\\n    Sequential,\\n)\\nfrom .conv import (\\n    Conv1d,\\n    Conv2d,\\n    Conv3d,\\n    ConvTranspose1d,\\n    ConvTranspose2d,\\n    ConvTranspose3d,\\n    LazyConv1d,\\n    LazyConv2d,\\n    LazyConv3d,\\n    LazyConvTranspose1d,\\n    LazyConvTranspose2d,\\n    LazyConvTranspose3d,\\n)\\nfrom .distance import CosineSimilarity, PairwiseDistance\\nfrom .dropout import (\\n    AlphaDropout,\\n    Dropout,\\n    Dropout1d,\\n    Dropout2d,\\n    Dropout3d,\\n    FeatureAlphaDropout,\\n)\\nfrom .flatten import Flatten, Unflatten\\nfrom .fold import Fold, Unfold\\nfrom .instancenorm import (\\n    InstanceNorm1d,\\n    InstanceNorm2d,\\n    InstanceNorm3d,\\n    LazyInstanceNorm1d,\\n    LazyInstanceNorm2d,\\n    LazyInstanceNorm3d,\\n)\\nfrom .loss import (\\n    BCELoss,\\n    BCEWithLogitsLoss,\\n    CosineEmbeddingLoss,\\n    CrossEntropyLoss,\\n    CTCLoss,\\n    GaussianNLLLoss,\\n    HingeEmbeddingLoss,\\n    HuberLoss,\\n    KLDivLoss,\\n    L1Loss,\\n    MarginRankingLoss,\\n    MSELoss,\\n    MultiLabelMarginLoss,\\n    MultiLabelSoftMarginLoss,\\n    MultiMarginLoss,\\n    NLLLoss,\\n    NLLLoss2d,\\n    PoissonNLLLoss,\\n    SmoothL1Loss,\\n    SoftMarginLoss,\\n    TripletMarginLoss,\\n    TripletMarginWithDistanceLoss,\\n)\\nfrom .normalization import (\\n    CrossMapLRN2d,\\n    GroupNorm,\\n    LayerNorm,\\n    LocalResponseNorm,\\n    RMSNorm,\\n)\\nfrom .padding import (\\n    CircularPad1d,\\n    CircularPad2d,\\n    CircularPad3d,\\n    ConstantPad1d,\\n    ConstantPad2d,\\n    ConstantPad3d,\\n    ReflectionPad1d,\\n    ReflectionPad2d,\\n    ReflectionPad3d,\\n    ReplicationPad1d,\\n    ReplicationPad2d,\\n    ReplicationPad3d,\\n    ZeroPad1d,\\n    ZeroPad2d,\\n    ZeroPad3d,\\n)\\nfrom .pixelshuffle import PixelShuffle, PixelUnshuffle\\nfrom .pooling import (\\n    AdaptiveAvgPool1d,\\n    AdaptiveAvgPool2d,\\n    AdaptiveAvgPool3d,\\n    AdaptiveMaxPool1d,\\n    AdaptiveMaxPool2d,\\n    AdaptiveMaxPool3d,\\n    AvgPool1d,\\n    AvgPool2d,\\n    AvgPool3d,\\n    FractionalMaxPool2d,\\n    FractionalMaxPool3d,\\n    LPPool1d,\\n    LPPool2d,\\n    LPPool3d,\\n    MaxPool1d,\\n    MaxPool2d,\\n    MaxPool3d,\\n    MaxUnpool1d,\\n    MaxUnpool2d,\\n    MaxUnpool3d,\\n)\\nfrom .rnn import GRU, GRUCell, LSTM, LSTMCell, RNN, RNNBase, RNNCell, RNNCellBase\\nfrom .sparse import Embedding, EmbeddingBag\\nfrom .transformer import (\\n    Transformer,\\n    TransformerDecoder,\\n    TransformerDecoderLayer,\\n    TransformerEncoder,\\n    TransformerEncoderLayer,\\n)\\nfrom .upsampling import Upsample, UpsamplingBilinear2d, UpsamplingNearest2d\\n\\n\\n__all__ = [\\n    \\\"AdaptiveAvgPool1d\\\",\\n    \\\"AdaptiveAvgPool2d\\\",\\n    \\\"AdaptiveAvgPool3d\\\",\\n    \\\"AdaptiveLogSoftmaxWithLoss\\\",\\n    \\\"AdaptiveMaxPool1d\\\",\\n    \\\"AdaptiveMaxPool2d\\\",\\n    \\\"AdaptiveMaxPool3d\\\",\\n    \\\"AlphaDropout\\\",\\n    \\\"AvgPool1d\\\",\\n    \\\"AvgPool2d\\\",\\n    \\\"AvgPool3d\\\",\\n    \\\"BCELoss\\\",\\n    \\\"BCEWithLogitsLoss\\\",\\n    \\\"BatchNorm1d\\\",\\n    \\\"BatchNorm2d\\\",\\n    \\\"BatchNorm3d\\\",\\n    \\\"Bilinear\\\",\\n    \\\"CELU\\\",\\n    \\\"CTCLoss\\\",\\n    \\\"ChannelShuffle\\\",\\n    \\\"CircularPad1d\\\",\\n    \\\"CircularPad2d\\\",\\n    \\\"CircularPad3d\\\",\\n    \\\"ConstantPad1d\\\",\\n    \\\"ConstantPad2d\\\",\\n    \\\"ConstantPad3d\\\",\\n    \\\"Container\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"ConvTranspose1d\\\",\\n    \\\"ConvTranspose2d\\\",\\n    \\\"ConvTranspose3d\\\",\\n    \\\"CosineEmbeddingLoss\\\",\\n    \\\"CosineSimilarity\\\",\\n    \\\"CrossEntropyLoss\\\",\\n    \\\"CrossMapLRN2d\\\",\\n    \\\"Dropout\\\",\\n    \\\"Dropout1d\\\",\\n    \\\"Dropout2d\\\",\\n    \\\"Dropout3d\\\",\\n    \\\"ELU\\\",\\n    \\\"Embedding\\\",\\n    \\\"EmbeddingBag\\\",\\n    \\\"FeatureAlphaDropout\\\",\\n    \\\"Flatten\\\",\\n    \\\"Fold\\\",\\n    \\\"FractionalMaxPool2d\\\",\\n    \\\"FractionalMaxPool3d\\\",\\n    \\\"GELU\\\",\\n    \\\"GLU\\\",\\n    \\\"GRU\\\",\\n    \\\"GRUCell\\\",\\n    \\\"GaussianNLLLoss\\\",\\n    \\\"GroupNorm\\\",\\n    \\\"Hardshrink\\\",\\n    \\\"Hardsigmoid\\\",\\n    \\\"Hardswish\\\",\\n    \\\"Hardtanh\\\",\\n    \\\"HingeEmbeddingLoss\\\",\\n    \\\"HuberLoss\\\",\\n    \\\"Identity\\\",\\n    \\\"InstanceNorm1d\\\",\\n    \\\"InstanceNorm2d\\\",\\n    \\\"InstanceNorm3d\\\",\\n    \\\"KLDivLoss\\\",\\n    \\\"L1Loss\\\",\\n    \\\"LPPool1d\\\",\\n    \\\"LPPool2d\\\",\\n    \\\"LPPool3d\\\",\\n    \\\"LSTM\\\",\\n    \\\"LSTMCell\\\",\\n    \\\"LayerNorm\\\",\\n    \\\"LazyBatchNorm1d\\\",\\n    \\\"LazyBatchNorm2d\\\",\\n    \\\"LazyBatchNorm3d\\\",\\n    \\\"LazyConv1d\\\",\\n    \\\"LazyConv2d\\\",\\n    \\\"LazyConv3d\\\",\\n    \\\"LazyConvTranspose1d\\\",\\n    \\\"LazyConvTranspose2d\\\",\\n    \\\"LazyConvTranspose3d\\\",\\n    \\\"LazyInstanceNorm1d\\\",\\n    \\\"LazyInstanceNorm2d\\\",\\n    \\\"LazyInstanceNorm3d\\\",\\n    \\\"LazyLinear\\\",\\n    \\\"LeakyReLU\\\",\\n    \\\"Linear\\\",\\n    \\\"LocalResponseNorm\\\",\\n    \\\"LogSigmoid\\\",\\n    \\\"LogSoftmax\\\",\\n    \\\"MSELoss\\\",\\n    \\\"MarginRankingLoss\\\",\\n    \\\"MaxPool1d\\\",\\n    \\\"MaxPool2d\\\",\\n    \\\"MaxPool3d\\\",\\n    \\\"MaxUnpool1d\\\",\\n    \\\"MaxUnpool2d\\\",\\n    \\\"MaxUnpool3d\\\",\\n    \\\"Mish\\\",\\n    \\\"Module\\\",\\n    \\\"ModuleDict\\\",\\n    \\\"ModuleList\\\",\\n    \\\"MultiLabelMarginLoss\\\",\\n    \\\"MultiLabelSoftMarginLoss\\\",\\n    \\\"MultiMarginLoss\\\",\\n    \\\"MultiheadAttention\\\",\\n    \\\"NLLLoss\\\",\\n    \\\"NLLLoss2d\\\",\\n    \\\"PReLU\\\",\\n    \\\"PairwiseDistance\\\",\\n    \\\"ParameterDict\\\",\\n    \\\"ParameterList\\\",\\n    \\\"PixelShuffle\\\",\\n    \\\"PixelUnshuffle\\\",\\n    \\\"PoissonNLLLoss\\\",\\n    \\\"RMSNorm\\\",\\n    \\\"RNN\\\",\\n    \\\"RNNBase\\\",\\n    \\\"RNNCell\\\",\\n    \\\"RNNCellBase\\\",\\n    \\\"RReLU\\\",\\n    \\\"ReLU\\\",\\n    \\\"ReLU6\\\",\\n    \\\"ReflectionPad1d\\\",\\n    \\\"ReflectionPad2d\\\",\\n    \\\"ReflectionPad3d\\\",\\n    \\\"ReplicationPad1d\\\",\\n    \\\"ReplicationPad2d\\\",\\n    \\\"ReplicationPad3d\\\",\\n    \\\"SELU\\\",\\n    \\\"Sequential\\\",\\n    \\\"SiLU\\\",\\n    \\\"Sigmoid\\\",\\n    \\\"SmoothL1Loss\\\",\\n    \\\"SoftMarginLoss\\\",\\n    \\\"Softmax\\\",\\n    \\\"Softmax2d\\\",\\n    \\\"Softmin\\\",\\n    \\\"Softplus\\\",\\n    \\\"Softshrink\\\",\\n    \\\"Softsign\\\",\\n    \\\"SyncBatchNorm\\\",\\n    \\\"Tanh\\\",\\n    \\\"Tanhshrink\\\",\\n    \\\"Threshold\\\",\\n    \\\"Transformer\\\",\\n    \\\"TransformerDecoder\\\",\\n    \\\"TransformerDecoderLayer\\\",\\n    \\\"TransformerEncoder\\\",\\n    \\\"TransformerEncoderLayer\\\",\\n    \\\"TripletMarginLoss\\\",\\n    \\\"TripletMarginWithDistanceLoss\\\",\\n    \\\"Unflatten\\\",\\n    \\\"Unfold\\\",\\n    \\\"Upsample\\\",\\n    \\\"UpsamplingBilinear2d\\\",\\n    \\\"UpsamplingNearest2d\\\",\\n    \\\"ZeroPad1d\\\",\\n    \\\"ZeroPad2d\\\",\\n    \\\"ZeroPad3d\\\",\\n]\\n\\n# Please keep this list sorted\\nassert __all__ == sorted(__all__)\\n\\n\\nimport warnings\\nfrom collections.abc import Iterable\\nfrom typing import (\\n    Any,\\n    Callable,\\n    List,\\n    NamedTuple,\\n    Optional,\\n    overload,\\n    Tuple,\\n    TypeVar,\\n    Union,\\n)\\nfrom typing_extensions import Self\\n\\nimport torch\\nfrom torch import _VF, Tensor\\n\\n\\n__all__ = [\\n    \\\"PackedSequence\\\",\\n    \\\"invert_permutation\\\",\\n    \\\"pack_padded_sequence\\\",\\n    \\\"pad_packed_sequence\\\",\\n    \\\"pad_sequence\\\",\\n    \\\"unpad_sequence\\\",\\n    \\\"pack_sequence\\\",\\n    \\\"unpack_sequence\\\",\\n]\\n\\n_T = TypeVar(\\\"_T\\\")\\n_R = TypeVar(\\\"_R\\\")\\n\\n\\nclass PackedSequence_(NamedTuple):\\n    data: torch.Tensor\\n    batch_sizes: torch.Tensor\\n    sorted_indices: Optional[torch.Tensor]\\n    unsorted_indices: Optional[torch.Tensor]\\n\\n\\ndef bind(optional: Optional[_T], fn: Callable[[_T], _R]) -> Optional[_R]:\\n    if optional is None:\\n        return None\\n    return fn(optional)\\n\\n\\nclass PackedSequence(PackedSequence_):\\n    r\\\"\\\"\\\"Holds the data and list of :attr:`batch_sizes` of a packed sequence.\\n\\n    All RNN modules accept packed sequences as inputs.\\n\\n    Note:\\n        Instances of this class should never be created manually. They are meant\\n        to be instantiated by functions like :func:`pack_padded_sequence`.\\n\\n        Batch sizes represent the number elements at each sequence step in\\n        the batch, not the varying sequence lengths passed to\\n        :func:`pack_padded_sequence`.  For instance, given data ``abc`` and ``x``\\n        the :class:`PackedSequence` would contain data ``axbc`` with\\n        ``batch_sizes=[2,1,1]``.\\n\\n    Attributes:\\n        data (Tensor): Tensor containing packed sequence\\n        batch_sizes (Tensor): Tensor of integers holding\\n            information about the batch size at each sequence step\\n        sorted_indices (Tensor, optional): Tensor of integers holding how this\\n            :class:`PackedSequence` is constructed from sequences.\\n        unsorted_indices (Tensor, optional): Tensor of integers holding how this\\n            to recover the original sequences with correct order.\\n\\n    .. note::\\n        :attr:`data` can be on arbitrary device and of arbitrary dtype.\\n        :attr:`sorted_indices` and :attr:`unsorted_indices` must be ``torch.int64``\\n        tensors on the same device as :attr:`data`.\\n\\n        However, :attr:`batch_sizes` should always be a CPU ``torch.int64`` tensor.\\n\\n        This invariant is maintained throughout :class:`PackedSequence` class,\\n        and all functions that construct a :class:`PackedSequence` in PyTorch\\n        (i.e., they only pass in tensors conforming to this constraint).\\n    \\\"\\\"\\\"\\n\\n    def __new__(\\n        cls,\\n        data: Tensor,\\n        batch_sizes: Optional[Tensor] = None,\\n        sorted_indices: Optional[Tensor] = None,\\n        unsorted_indices: Optional[Tensor] = None,\\n    ) -> Self:\\n        return super().__new__(\\n            cls,\\n            *_packed_sequence_init_args(\\n                data, batch_sizes, sorted_indices, unsorted_indices\\n            ),\\n        )\\n\\n    # NOTE [ device and dtype of a PackedSequence ]\\n    #\\n    # See the note above in doc string (starting with \\\":attr:`data` can be on\\n    # arbitrary device...\\\").\\n    def pin_memory(self) -> Self:\\n        # Why not convert `batch_sizes`?\\n        # See NOTE [ device and dtype of a PackedSequence ]\\n        return type(self)(\\n            self.data.pin_memory(),\\n            self.batch_sizes,\\n            bind(self.sorted_indices, lambda t: t.pin_memory()),\\n            bind(self.unsorted_indices, lambda t: t.pin_memory()),\\n        )\\n\\n    @overload\\n    def to(\\n        self,\\n        dtype: torch.dtype,\\n        non_blocking: bool = ...,\\n        copy: bool = ...,\\n    ) -> Self:\\n        ...\\n\\n    @overload\\n    def to(\\n        self,\\n        device: Optional[Union[str, torch.device, int]] = ...,\\n        dtype: Optional[torch.dtype] = ...,\\n        non_blocking: bool = ...,\\n        copy: bool = ...,\\n    ) -> Self:\\n        ...\\n\\n    @overload\\n    def to(\\n        self,\\n        other: Tensor,\\n        non_blocking: bool = ...,\\n        copy: bool = ...,\\n    ) -> Self:\\n        ...\\n\\n    def to(self, *args: Any, **kwargs: Any) -> Self:\\n        r\\\"\\\"\\\"Perform dtype and/or device conversion on `self.data`.\\n\\n        It has similar signature as :meth:`torch.Tensor.to`, except optional\\n        arguments like `non_blocking` and `copy` should be passed as kwargs,\\n        not args, or they will not apply to the index tensors.\\n\\n        .. note::\\n\\n            If the ``self.data`` Tensor already has the correct :class:`torch.dtype`\\n            and :class:`torch.device`, then ``self`` is returned.\\n            Otherwise, returns a copy with the desired configuration.\\n        \\\"\\\"\\\"\\n        # Why not convert `batch_sizes`?\\n        # See NOTE [ device and dtype of a PackedSequence ]\\n        data = self.data.to(*args, **kwargs)\\n        if data is self.data:\\n            return self\\n        else:\\n            # Does not forward device or dtype arg/kwargs, device is set from data.device\\n            kwargs = dict(\\n                filter(lambda t: t[0] != \\\"device\\\" and t[0] != \\\"dtype\\\", kwargs.items())\\n            )\\n            sorted_indices = bind(\\n                self.sorted_indices, lambda t: t.to(data.device, **kwargs)\\n            )\\n            unsorted_indices = bind(\\n                self.unsorted_indices, lambda t: t.to(data.device, **kwargs)\\n            )\\n            return type(self)(data, self.batch_sizes, sorted_indices, unsorted_indices)\\n\\n    def cuda(self, *args: Any, **kwargs: Any) -> Self:\\n        # Tests to see if 'cuda' should be added to kwargs\\n        ex = torch.tensor((), dtype=self.data.dtype, device=self.data.device).to(\\n            *args, **kwargs\\n        )\\n        if ex.is_cuda:\\n            return self.to(*args, **kwargs)\\n        kwargs[\\\"device\\\"] = \\\"cuda\\\"\\n        return self.to(*args, **kwargs)\\n\\n    def cpu(self, *args: Any, **kwargs: Any) -> Self:\\n        ex = torch.tensor((), dtype=self.data.dtype, device=self.data.device).to(\\n            *args, **kwargs\\n        )\\n        if ex.device.type == \\\"cpu\\\":\\n            return self.to(*args, **kwargs)\\n        kwargs[\\\"device\\\"] = \\\"cpu\\\"\\n        return self.to(*args, **kwargs)\\n\\n    def double(self) -> Self:\\n        return self.to(dtype=torch.double)\\n\\n    def float(self) -> Self:\\n        return self.to(dtype=torch.float)\\n\\n    def half(self) -> Self:\\n        return self.to(dtype=torch.half)\\n\\n    def long(self) -> Self:\\n        return self.to(dtype=torch.long)\\n\\n    def int(self) -> Self:\\n        return self.to(dtype=torch.int)\\n\\n    def short(self) -> Self:\\n        return self.to(dtype=torch.short)\\n\\n    def char(self) -> Self:\\n        return self.to(dtype=torch.int8)\\n\\n    def byte(self) -> Self:\\n        return self.to(dtype=torch.uint8)\\n\\n    @property\\n    def is_cuda(self) -> bool:\\n        r\\\"\\\"\\\"Return true if `self.data` stored on a gpu.\\\"\\\"\\\"\\n        return self.data.is_cuda\\n\\n    def is_pinned(self) -> bool:\\n        r\\\"\\\"\\\"Return true if `self.data` stored on in pinned memory.\\\"\\\"\\\"\\n        return self.data.is_pinned()\\n\\n\\n# TorchScript doesn't support constructors on named tuples, so we use this helper\\n# method to construct PackedSequence\\ndef _packed_sequence_init_args(\\n    data: Tensor,\\n    batch_sizes: Optional[Tensor] = None,\\n    sorted_indices: Optional[Tensor] = None,\\n    unsorted_indices: Optional[Tensor] = None,\\n) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:\\n    # NB: if unsorted_indices is provided, it should be the inverse permutation\\n    # to sorted_indices. Don't assert it here because the PackedSequence ctor\\n    # should only be used internally.\\n\\n    if unsorted_indices is None:\\n        unsorted_indices = invert_permutation(sorted_indices)\\n\\n    # support being called as `PackedSequence(data, batch_sizes, sorted_indices)`\\n    if batch_sizes is not None:\\n        # TODO: Re-enable this check (.type isn't supported in TorchScript)\\n        if batch_sizes.device.type != \\\"cpu\\\":\\n            raise ValueError(\\n                \\\"batch_sizes should always be on CPU. \\\"\\n                \\\"Instances of PackedSequence should never be created manually. \\\"\\n                \\\"They should be instantiated by functions like pack_sequence \\\"\\n                \\\"and pack_padded_sequences in nn.utils.rnn. \\\"\\n                \\\"https://pytorch.org/docs/stable/nn.html#torch.nn.utils.rnn.pack_sequence\\\"\\n            )\\n        return data, batch_sizes, sorted_indices, unsorted_indices\\n\\n    # support being called as `PackedSequence((data, batch_sizes), *, sorted_indices)`\\n    else:\\n        assert isinstance(data, (list, tuple)) and len(data) == 2\\n        return data[0], data[1], sorted_indices, unsorted_indices\\n\\n\\ndef _packed_sequence_init(\\n    data: Tensor,\\n    batch_sizes: Optional[Tensor] = None,\\n    sorted_indices: Optional[Tensor] = None,\\n    unsorted_indices: Optional[Tensor] = None,\\n) -> PackedSequence:\\n    data, batch_sizes, sorted_indices, unsorted_indices = _packed_sequence_init_args(\\n        data, batch_sizes, sorted_indices, unsorted_indices\\n    )\\n    return PackedSequence(data, batch_sizes, sorted_indices, unsorted_indices)\\n\\n\\ndef invert_permutation(permutation: Optional[Tensor]) -> Optional[Tensor]:\\n    if permutation is None:\\n        return None\\n    output = torch.empty_like(permutation, memory_format=torch.legacy_contiguous_format)\\n    output.scatter_(\\n        0, permutation, torch.arange(0, permutation.numel(), device=permutation.device)\\n    )\\n    return output\\n\\n\\ndef pack_padded_sequence(\\n    input: Tensor,\\n    lengths: Union[Tensor, List[int]],\\n    batch_first: bool = False,\\n    enforce_sorted: bool = True,\\n) -> PackedSequence:\\n    r\\\"\\\"\\\"Packs a Tensor containing padded sequences of variable length.\\n\\n    :attr:`input` can be of size ``T x B x *`` (if :attr:`batch_first` is ``False``)\\n    or ``B x T x *`` (if :attr:`batch_first` is ``True``) where ``T`` is the length\\n    of the longest sequence, ``B`` is the batch size, and ``*`` is any number of dimensions\\n    (including 0).\\n\\n    For unsorted sequences, use `enforce_sorted = False`. If :attr:`enforce_sorted` is\\n    ``True``, the sequences should be sorted by length in a decreasing order, i.e.\\n    ``input[:,0]`` should be the longest sequence, and ``input[:,B-1]`` the shortest\\n    one. `enforce_sorted = True` is only necessary for ONNX export.\\n\\n    Note:\\n        This function accepts any input that has at least two dimensions. You\\n        can apply it to pack the labels, and use the output of the RNN with\\n        them to compute the loss directly. A Tensor can be retrieved from\\n        a :class:`PackedSequence` object by accessing its ``.data`` attribute.\\n\\n    Args:\\n        input (Tensor): padded batch of variable length sequences.\\n        lengths (Tensor or list(int)): list of sequence lengths of each batch\\n            element (must be on the CPU if provided as a tensor).\\n        batch_first (bool, optional): if ``True``, the input is expected in ``B x T x *``\\n            format, ``T x B x *`` otherwise.\\n        enforce_sorted (bool, optional): if ``True``, the input is expected to\\n            contain sequences sorted by length in a decreasing order. If\\n            ``False``, the input will get sorted unconditionally. Default: ``True``.\\n\\n    Returns:\\n        a :class:`PackedSequence` object\\n    \\\"\\\"\\\"\\n    if not isinstance(lengths, torch.Tensor):\\n        if torch._C._get_tracing_state():\\n            warnings.warn(\\n                \\\"pack_padded_sequence has been called with a Python list of \\\"\\n                \\\"sequence lengths. The tracer cannot track the data flow of Python \\\"\\n                \\\"values, and it will treat them as constants, likely rendering \\\"\\n                \\\"the trace incorrect for any other combination of lengths.\\\",\\n                stacklevel=2,\\n            )\\n        lengths = torch.as_tensor(lengths, dtype=torch.int64, device=\\\"cpu\\\")\\n    else:\\n        lengths = lengths.to(dtype=torch.int64)\\n\\n    if enforce_sorted:\\n        sorted_indices = None\\n    else:\\n        lengths, sorted_indices = torch.sort(lengths, descending=True)\\n        sorted_indices = sorted_indices.to(input.device)\\n        batch_dim = 0 if batch_first else 1\\n        input = input.index_select(batch_dim, sorted_indices)\\n\\n    data, batch_sizes = _VF._pack_padded_sequence(input, lengths, batch_first)\\n    return _packed_sequence_init(data, batch_sizes, sorted_indices, None)\\n\\n\\ndef pad_packed_sequence(\\n    sequence: PackedSequence,\\n    batch_first: bool = False,\\n    padding_value: float = 0.0,\\n    total_length: Optional[int] = None,\\n) -> Tuple[Tensor, Tensor]:\\n    r\\\"\\\"\\\"Pad a packed batch of variable length sequences.\\n\\n    It is an inverse operation to :func:`pack_padded_sequence`.\\n\\n    The returned Tensor's data will be of size ``T x B x *`` (if :attr:`batch_first` is ``False``)\\n    or ``B x T x *`` (if :attr:`batch_first` is ``True``) , where ``T`` is the length of the longest\\n    sequence and ``B`` is the batch size.\\n\\n    Example:\\n        >>> from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\\n        >>> seq = torch.tensor([[1, 2, 0], [3, 0, 0], [4, 5, 6]])\\n        >>> lens = [2, 1, 3]\\n        >>> packed = pack_padded_sequence(seq, lens, batch_first=True, enforce_sorted=False)\\n        >>> packed\\n        PackedSequence(data=tensor([4, 1, 3, 5, 2, 6]), batch_sizes=tensor([3, 2, 1]),\\n                       sorted_indices=tensor([2, 0, 1]), unsorted_indices=tensor([1, 2, 0]))\\n        >>> seq_unpacked, lens_unpacked = pad_packed_sequence(packed, batch_first=True)\\n        >>> seq_unpacked\\n        tensor([[1, 2, 0],\\n                [3, 0, 0],\\n                [4, 5, 6]])\\n        >>> lens_unpacked\\n        tensor([2, 1, 3])\\n\\n    .. note::\\n        :attr:`total_length` is useful to implement the\\n        ``pack sequence -> recurrent network -> unpack sequence`` pattern in a\\n        :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.\\n        See :ref:`this FAQ section <pack-rnn-unpack-with-data-parallelism>` for\\n        details.\\n\\n    Args:\\n        sequence (PackedSequence): batch to pad\\n        batch_first (bool, optional): if ``True``, the output will be in ``B x T x *``\\n            format, ``T x B x *`` otherwise.\\n        padding_value (float, optional): values for padded elements.\\n        total_length (int, optional): if not ``None``, the output will be padded to\\n            have length :attr:`total_length`. This method will throw :class:`ValueError`\\n            if :attr:`total_length` is less than the max sequence length in\\n            :attr:`sequence`.\\n\\n    Returns:\\n        Tuple of Tensor containing the padded sequence, and a Tensor\\n        containing the list of lengths of each sequence in the batch.\\n        Batch elements will be re-ordered as they were ordered originally when\\n        the batch was passed to ``pack_padded_sequence`` or ``pack_sequence``.\\n    \\\"\\\"\\\"\\n    max_seq_length = sequence.batch_sizes.size(0)\\n    if total_length is not None:\\n        if total_length < max_seq_length:\\n            raise ValueError(\\n                \\\"Expected total_length to be at least the length \\\"\\n                \\\"of the longest sequence in input, but got \\\"\\n                f\\\"total_length={total_length} and max sequence length being {max_seq_length}\\\"\\n            )\\n        max_seq_length = total_length\\n    padded_output, lengths = _VF._pad_packed_sequence(\\n        sequence.data, sequence.batch_sizes, batch_first, padding_value, max_seq_length\\n    )\\n    unsorted_indices = sequence.unsorted_indices\\n    if unsorted_indices is not None:\\n        batch_dim = 0 if batch_first else 1\\n        return (\\n            padded_output.index_select(batch_dim, unsorted_indices),\\n            lengths[unsorted_indices.cpu()],\\n        )\\n    return padded_output, lengths\\n\\n\\n# NOTE: for JIT-compatibility, we need to be more restrictive here and use specific types instead of Iterable.\\ndef pad_sequence(\\n    sequences: Union[Tensor, List[Tensor]],\\n    batch_first: bool = False,\\n    padding_value: float = 0.0,\\n    padding_side: str = \\\"right\\\",\\n) -> Tensor:\\n    r\\\"\\\"\\\"Pad a list of variable length Tensors with :attr:`padding_value`.\\n\\n    ``pad_sequence`` stacks a list of Tensors along a new dimension, and pads them\\n    to equal length. :attr:`sequences` can be list of sequences with size ``L x *``,\\n    where `L` is length of the sequence and ``*`` is any number of dimensions\\n    (including 0). If :attr:`batch_first` is ``False``, the output is of size\\n    ``T x B x *``, and ``B x T x *`` otherwise, where ``B`` is the batch size\\n    (the number of elements in :attr:`sequences`), ``T`` is the length of the longest\\n    sequence.\\n\\n    Example:\\n        >>> from torch.nn.utils.rnn import pad_sequence\\n        >>> a = torch.ones(25, 300)\\n        >>> b = torch.ones(22, 300)\\n        >>> c = torch.ones(15, 300)\\n        >>> pad_sequence([a, b, c]).size()\\n        torch.Size([25, 3, 300])\\n\\n    Note:\\n        This function returns a Tensor of size ``T x B x *`` or ``B x T x *``\\n        where `T` is the length of the longest sequence. This function assumes\\n        trailing dimensions and type of all the Tensors in sequences are same.\\n\\n    Args:\\n        sequences (list[Tensor]): list of variable length sequences.\\n        batch_first (bool, optional): if ``True``, the output will be in ``B x T x *``\\n            format, ``T x B x *`` otherwise.\\n        padding_value (float, optional): value for padded elements. Default: 0.\\n        padding_side (str, optional): the side to pad the sequences on.\\n            Default: \\\"right\\\".\\n\\n    Returns:\\n        Tensor of size ``T x B x *`` if :attr:`batch_first` is ``False``.\\n        Tensor of size ``B x T x *`` otherwise\\n    \\\"\\\"\\\"\\n    if not (torch.jit.is_tracing() or torch.jit.is_scripting()):\\n        # JIT doesn't support `Iterable`\\n        if not isinstance(sequences, Iterable):\\n            msg = (\\n                \\\"pad_sequence: Expected iterable for input sequences, but got arg of type: \\\"\\n                f\\\"{type(sequences)}\\\"\\n            )\\n            raise RuntimeError(msg)\\n\\n        # In JIT context this leads to,\\n        # RuntimeError: cannot statically infer the expected size of a list in this context\\n        sequences = tuple(sequences)  # type: ignore[assignment]\\n    else:\\n        # For JIT, we only support Union[Tensor, Tuple[Tensor]]\\n        if isinstance(sequences, torch.Tensor):\\n            sequences = sequences.unbind(0)  # type: ignore[assignment]\\n\\n    # assuming trailing dimensions and type of all the Tensors\\n    # in sequences are same and fetching those from sequences[0]\\n    return torch._C._nn.pad_sequence(\\n        sequences, batch_first, padding_value, padding_side  # type: ignore[arg-type]\\n    )\\n\\n\\ndef unpad_sequence(\\n    padded_sequences: Tensor,\\n    lengths: Tensor,\\n    batch_first: bool = False,\\n) -> List[Tensor]:\\n    r\\\"\\\"\\\"Unpad padded Tensor into a list of variable length Tensors.\\n\\n    ``unpad_sequence`` unstacks padded Tensor into a list of variable length Tensors.\\n\\n    Example:\\n        >>> from torch.nn.utils.rnn import pad_sequence, unpad_sequence\\n        >>> a = torch.ones(25, 300)\\n        >>> b = torch.ones(22, 300)\\n        >>> c = torch.ones(15, 300)\\n        >>> sequences = [a, b, c]\\n        >>> padded_sequences = pad_sequence(sequences)\\n        >>> lengths = torch.as_tensor([v.size(0) for v in sequences])\\n        >>> unpadded_sequences = unpad_sequence(padded_sequences, lengths)\\n        >>> torch.allclose(sequences[0], unpadded_sequences[0])\\n        True\\n        >>> torch.allclose(sequences[1], unpadded_sequences[1])\\n        True\\n        >>> torch.allclose(sequences[2], unpadded_sequences[2])\\n        True\\n\\n    Args:\\n        padded_sequences (Tensor): padded sequences.\\n        lengths (Tensor): length of original (unpadded) sequences.\\n        batch_first (bool, optional): whether batch dimension first or not. Default: False.\\n\\n    Returns:\\n        a list of :class:`Tensor` objects\\n    \\\"\\\"\\\"\\n    unpadded_sequences = []\\n\\n    if not batch_first:\\n        padded_sequences.transpose_(0, 1)\\n\\n    max_length = padded_sequences.shape[1]\\n    idx = torch.arange(max_length, device=lengths.device)\\n\\n    for seq, length in zip(padded_sequences, lengths):\\n        mask = idx < length\\n        unpacked_seq = seq[mask]\\n        unpadded_sequences.append(unpacked_seq)\\n\\n    return unpadded_sequences\\n\\n\\ndef pack_sequence(\\n    sequences: List[Tensor],\\n    enforce_sorted: bool = True,\\n) -> PackedSequence:\\n    r\\\"\\\"\\\"Packs a list of variable length Tensors.\\n\\n    Consecutive call of the next functions: ``pad_sequence``, ``pack_padded_sequence``.\\n\\n    ``sequences`` should be a list of Tensors of size ``L x *``, where `L` is\\n    the length of a sequence and `*` is any number of trailing dimensions,\\n    including zero.\\n\\n    For unsorted sequences, use `enforce_sorted = False`. If ``enforce_sorted``\\n    is ``True``, the sequences should be sorted in the order of decreasing length.\\n    ``enforce_sorted = True`` is only necessary for ONNX export.\\n\\n    Example:\\n        >>> from torch.nn.utils.rnn import pack_sequence\\n        >>> a = torch.tensor([1, 2, 3])\\n        >>> b = torch.tensor([4, 5])\\n        >>> c = torch.tensor([6])\\n        >>> pack_sequence([a, b, c])\\n        PackedSequence(data=tensor([1, 4, 6, 2, 5, 3]), batch_sizes=tensor([3, 2, 1]), sorted_indices=None, unsorted_indices=None)\\n\\n    Args:\\n        sequences (list[Tensor]): A list of sequences of decreasing length.\\n        enforce_sorted (bool, optional): if ``True``, checks that the input\\n            contains sequences sorted by length in a decreasing order. If\\n            ``False``, this condition is not checked. Default: ``True``.\\n\\n    Returns:\\n        a :class:`PackedSequence` object\\n    \\\"\\\"\\\"\\n    lengths = torch.as_tensor([v.size(0) for v in sequences])\\n    return pack_padded_sequence(\\n        pad_sequence(sequences), lengths, enforce_sorted=enforce_sorted\\n    )\\n\\n\\ndef unpack_sequence(packed_sequences: PackedSequence) -> List[Tensor]:\\n    r\\\"\\\"\\\"Unpack PackedSequence into a list of variable length Tensors.\\n\\n    ``packed_sequences`` should be a PackedSequence object.\\n\\n    Example:\\n        >>> from torch.nn.utils.rnn import pack_sequence, unpack_sequence\\n        >>> a = torch.tensor([1, 2, 3])\\n        >>> b = torch.tensor([4, 5])\\n        >>> c = torch.tensor([6])\\n        >>> sequences = [a, b, c]\\n        >>> print(sequences)\\n        [tensor([1, 2, 3]), tensor([4, 5]), tensor([6])]\\n        >>> packed_sequences = pack_sequence(sequences)\\n        >>> print(packed_sequences)\\n        PackedSequence(data=tensor([1, 4, 6, 2, 5, 3]), batch_sizes=tensor([3, 2, 1]), sorted_indices=None, unsorted_indices=None)\\n        >>> unpacked_sequences = unpack_sequence(packed_sequences)\\n        >>> print(unpacked_sequences)\\n        [tensor([1, 2, 3]), tensor([4, 5]), tensor([6])]\\n\\n    Args:\\n        packed_sequences (PackedSequence): A PackedSequence object.\\n\\n    Returns:\\n        a list of :class:`Tensor` objects\\n    \\\"\\\"\\\"\\n    padded_sequences, lengths = pad_packed_sequence(packed_sequences, batch_first=True)\\n    unpacked_sequences = unpad_sequence(padded_sequences, lengths, batch_first=True)\\n    return unpacked_sequences\\n\\n\\n# mypy: allow-untyped-defs\\nimport torch\\n\\n\\ndef convert_conv2d_weight_memory_format(module, memory_format):\\n    r\\\"\\\"\\\"Convert ``memory_format`` of ``nn.Conv2d.weight`` to ``memory_format``.\\n\\n    The conversion recursively applies to nested ``nn.Module``, including ``module``.\\n    Note that it only changes the memory_format, but not the semantics of each dimensions.\\n    This function is used to facilitate the computation to adopt NHWC kernels, which\\n    provides considerable speed up for fp16 data on CUDA devices with compute capability >= 7.0\\n\\n    .. note::\\n        Calling ``model.to(memory_format=torch.channels_last)`` is more aggressive\\n        than the utility function ``convert_conv2d_weight_memory_format``. Any\\n        layer with 4d weight will be affected by ``model.to``, which does not\\n        necessarily benefit from conversion to specified ``memory_format``.\\n        One place we are confident in is that NHWC(channels_last) conversion for\\n        convolution in cuDNN, as it is beneficial to run convolution in NHWC,\\n        even in cases where we have to apply permutation to input tensors.\\n\\n        Hence our strategy here is to convert only the weight of convolution to\\n        channels_last. This ensures that;\\n        1. Fast convolution kernels will be used, the benefit of which could\\n        outweigh overhead of permutation (if input is not in the same format).\\n        2. No unnecessary permutations are applied on layers that do not benefit\\n        from memory_format conversion.\\n\\n        The optimal case is that, layers between convolution layers are channels\\n        last compatible. Input tensor would be permuted to channels last when it\\n        encounters the first convolution layer and stay in that memory format.\\n        Hence following convolutions will not need to permute its input tensor.\\n\\n        In case where a channels last incompatible layer is between convolution\\n        layers, we need to permute the input tensor back to contiguous format\\n        for that layer. The input tensor will go through the remaining layers in\\n        contiguous format and be permuted to channels last when it encounters\\n        another convolution layer. There's no point in propagating that\\n        permutation to an earlier layer, as most layers are quite agnostic to\\n        ``memory_format``.\\n\\n        This claim might change when PyTorch supports fusion of permutation, as\\n        there might have been a better spot to fuse the permutation other than\\n        immediately before a convolution.\\n\\n    Args:\\n        module (nn.Module): ``nn.Conv2d`` & ``nn.ConvTranspose2d`` or container\\n                            ``nn.Module``\\n        memory_format: user specified ``memory_format``,\\n            e.g. ``torch.channels_last`` or ``torch.contiguous_format``\\n\\n    Returns:\\n        The original module with updated ``nn.Conv2d``\\n\\n    Example:\\n        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)\\n        >>> # xdoctest: +REQUIRES(env:CUBLAS_WORKSPACE_CONFIG)\\n        >>> input = torch.randint(1, 10, (2, 8, 4, 4), dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> model = nn.Sequential(\\n        >>>     nn.Conv2d(8, 4, 3)).cuda().half()\\n        >>> # This is identical to:\\n        >>> # nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last)\\n        >>> model = nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last)\\n        >>> out = model(input)\\n    \\\"\\\"\\\"\\n    # TODO: expand this to `_ConvNd` when channels_last support is extended\\n    # beyond only 4d tensors.\\n    if isinstance(module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)):\\n        weight_data = (\\n            module.weight.detach().clone().contiguous(memory_format=memory_format)\\n        )\\n        module.weight.data = weight_data.resize_(\\n            weight_data.size(), memory_format=memory_format\\n        )\\n    for child in module.children():\\n        convert_conv2d_weight_memory_format(child, memory_format)\\n    return module\\n\\n\\ndef convert_conv3d_weight_memory_format(module, memory_format):\\n    r\\\"\\\"\\\"Convert ``memory_format`` of ``nn.Conv3d.weight`` to ``memory_format``\\n    The conversion recursively applies to nested ``nn.Module``, including ``module``.\\n    Note that it only changes the memory_format, but not the semantics of each dimensions.\\n    This function is used to facilitate the computation to adopt NHWC kernels, which\\n    provides considerable speed up for fp16 data on CUDA devices with compute capability >= 7.0\\n\\n    .. note::\\n        Calling ``model.to(memory_format=torch.channels_last_3d)`` is more aggressive\\n        than the utility function ``convert_conv3d_weight_memory_format``. Any\\n        layer with 4d weight will be affected by ``model.to``, which does not\\n        necessarily benefit from conversion to specified ``memory_format``.\\n        One place we are confident in is that NDHWC(channels_last_3d) conversion for\\n        convolution in cuDNN, as it is beneficial to run convolution in NDHWC,\\n        even in cases where we have to apply permutation to input tensors.\\n\\n        Hence our strategy here is to convert only the weight of convolution to\\n        channels_last_3d. This ensures that;\\n        1. Fast convolution kernels will be used, the benefit of which could\\n        outweigh overhead of permutation (if input is not in the same format).\\n        2. No unnecessary permutations are applied on layers that do not benefit\\n        from memory_format conversion.\\n\\n        The optimal case is that, layers between convolution layers are channels\\n        last compatible. Input tensor would be permuted to channels last when it\\n        encounters the first convolution layer and stay in that memory format.\\n        Hence following convolutions will not need to permute its input tensor.\\n\\n        In case where a channels last incompatible layer is between convolution\\n        layers, we need to permute the input tensor back to contiguous format\\n        for that layer. The input tensor will go through the remaining layers in\\n        contiguous format and be permuted to channels last when it encounters\\n        another convolution layer. There's no point in propagating that\\n        permutation to an earlier layer, as most layers are quite agnostic to\\n        ``memory_format``.\\n\\n        This claim might change when PyTorch supports fusion of permutation, as\\n        there might have been a better spot to fuse the permutation other than\\n        immediately before a convolution.\\n\\n    Args:\\n        module (nn.Module): ``nn.Conv3d`` & ``nn.ConvTranspose3d`` or container\\n                            ``nn.Module``\\n        memory_format: user specified ``memory_format``,\\n            e.g. ``torch.channels_last`` or ``torch.contiguous_format``\\n\\n    Returns:\\n        The original module with updated ``nn.Conv3d``\\n\\n    Example:\\n        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)\\n        >>> # xdoctest: +REQUIRES(env:CUBLAS_WORKSPACE_CONFIG)\\n        >>> input = torch.randint(1, 10, (2, 8, 4, 4, 4), dtype=torch.float16, device=\\\"cuda\\\")\\n        >>> model = nn.Sequential(\\n        >>>     nn.Conv3d(8, 4, 3)).cuda().half()\\n        >>> # This is identical to:\\n        >>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)\\n        >>> model = nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)\\n        >>> out = model(input)\\n    \\\"\\\"\\\"\\n\\n    # TODO: expand this to `_ConvNd` when channels_last support is extended\\n    # beyond only 4d tensors.\\n    if isinstance(module, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)):\\n        weight_data = (\\n            module.weight.detach().clone().contiguous(memory_format=memory_format)\\n        )\\n        module.weight.data = weight_data.resize_(\\n            weight_data.size(), memory_format=memory_format\\n        )\\n    for child in module.children():\\n        convert_conv3d_weight_memory_format(child, memory_format)\\n    return module\\n\\n\\n# mypy: allow-untyped-defs\\nr\\\"\\\"\\\"Weight Normalization from https://arxiv.org/abs/1602.07868.\\\"\\\"\\\"\\nfrom typing import Any, TypeVar\\nfrom typing_extensions import deprecated\\n\\nfrom torch import _weight_norm, norm_except_dim\\nfrom torch.nn.modules import Module\\nfrom torch.nn.parameter import Parameter, UninitializedParameter\\n\\n\\n__all__ = [\\\"WeightNorm\\\", \\\"weight_norm\\\", \\\"remove_weight_norm\\\"]\\n\\n\\nclass WeightNorm:\\n    name: str\\n    dim: int\\n\\n    def __init__(self, name: str, dim: int) -> None:\\n        if dim is None:\\n            dim = -1\\n        self.name = name\\n        self.dim = dim\\n\\n    # TODO Make return type more specific\\n    def compute_weight(self, module: Module) -> Any:\\n        g = getattr(module, self.name + \\\"_g\\\")\\n        v = getattr(module, self.name + \\\"_v\\\")\\n        return _weight_norm(v, g, self.dim)\\n\\n    @staticmethod\\n    @deprecated(\\n        \\\"`torch.nn.utils.weight_norm` is deprecated \\\"\\n        \\\"in favor of `torch.nn.utils.parametrizations.weight_norm`.\\\",\\n        category=FutureWarning,\\n    )\\n    def apply(module, name: str, dim: int) -> \\\"WeightNorm\\\":\\n        for hook in module._forward_pre_hooks.values():\\n            if isinstance(hook, WeightNorm) and hook.name == name:\\n                raise RuntimeError(\\n                    f\\\"Cannot register two weight_norm hooks on the same parameter {name}\\\"\\n                )\\n\\n        if dim is None:\\n            dim = -1\\n\\n        fn = WeightNorm(name, dim)\\n\\n        weight = getattr(module, name)\\n        if isinstance(weight, UninitializedParameter):\\n            raise ValueError(\\n                \\\"The module passed to `WeightNorm` can't have uninitialized parameters. \\\"\\n                \\\"Make sure to run the dummy forward before applying weight normalization\\\"\\n            )\\n        # remove w from parameter list\\n        del module._parameters[name]\\n\\n        # add g and v as new parameters and express w as g/||v|| * v\\n        module.register_parameter(\\n            name + \\\"_g\\\", Parameter(norm_except_dim(weight, 2, dim).data)\\n        )\\n        module.register_parameter(name + \\\"_v\\\", Parameter(weight.data))\\n        setattr(module, name, fn.compute_weight(module))\\n\\n        # recompute weight before every forward()\\n        module.register_forward_pre_hook(fn)\\n\\n        return fn\\n\\n    def remove(self, module: Module) -> None:\\n        weight = self.compute_weight(module)\\n        delattr(module, self.name)\\n        del module._parameters[self.name + \\\"_g\\\"]\\n        del module._parameters[self.name + \\\"_v\\\"]\\n        setattr(module, self.name, Parameter(weight.data))\\n\\n    def __call__(self, module: Module, inputs: Any) -> None:\\n        setattr(module, self.name, self.compute_weight(module))\\n\\n\\nT_module = TypeVar(\\\"T_module\\\", bound=Module)\\n\\n\\ndef weight_norm(module: T_module, name: str = \\\"weight\\\", dim: int = 0) -> T_module:\\n    r\\\"\\\"\\\"Apply weight normalization to a parameter in the given module.\\n\\n    .. math::\\n         \\\\mathbf{w} = g \\\\dfrac{\\\\mathbf{v}}{\\\\|\\\\mathbf{v}\\\\|}\\n\\n    Weight normalization is a reparameterization that decouples the magnitude\\n    of a weight tensor from its direction. This replaces the parameter specified\\n    by :attr:`name` (e.g. ``'weight'``) with two parameters: one specifying the magnitude\\n    (e.g. ``'weight_g'``) and one specifying the direction (e.g. ``'weight_v'``).\\n    Weight normalization is implemented via a hook that recomputes the weight\\n    tensor from the magnitude and direction before every :meth:`~Module.forward`\\n    call.\\n\\n    By default, with ``dim=0``, the norm is computed independently per output\\n    channel/plane. To compute a norm over the entire weight tensor, use\\n    ``dim=None``.\\n\\n    See https://arxiv.org/abs/1602.07868\\n\\n    .. warning::\\n\\n        This function is deprecated.  Use :func:`torch.nn.utils.parametrizations.weight_norm`\\n        which uses the modern parametrization API.  The new ``weight_norm`` is compatible\\n        with ``state_dict`` generated from old ``weight_norm``.\\n\\n        Migration guide:\\n\\n        * The magnitude (``weight_g``) and direction (``weight_v``) are now expressed\\n          as ``parametrizations.weight.original0`` and ``parametrizations.weight.original1``\\n          respectively.  If this is bothering you, please comment on\\n          https://github.com/pytorch/pytorch/issues/102999\\n\\n        * To remove the weight normalization reparametrization, use\\n          :func:`torch.nn.utils.parametrize.remove_parametrizations`.\\n\\n        * The weight is no longer recomputed once at module forward; instead, it will\\n          be recomputed on every access.  To restore the old behavior, use\\n          :func:`torch.nn.utils.parametrize.cached` before invoking the module\\n          in question.\\n\\n    Args:\\n        module (Module): containing module\\n        name (str, optional): name of weight parameter\\n        dim (int, optional): dimension over which to compute the norm\\n\\n    Returns:\\n        The original module with the weight norm hook\\n\\n    Example::\\n\\n        >>> m = weight_norm(nn.Linear(20, 40), name='weight')\\n        >>> m\\n        Linear(in_features=20, out_features=40, bias=True)\\n        >>> m.weight_g.size()\\n        torch.Size([40, 1])\\n        >>> m.weight_v.size()\\n        torch.Size([40, 20])\\n\\n    \\\"\\\"\\\"\\n    WeightNorm.apply(module, name, dim)\\n    return module\\n\\n\\ndef remove_weight_norm(module: T_module, name: str = \\\"weight\\\") -> T_module:\\n    r\\\"\\\"\\\"Remove the weight normalization reparameterization from a module.\\n\\n    Args:\\n        module (Module): containing module\\n        name (str, optional): name of weight parameter\\n\\n    Example:\\n        >>> m = weight_norm(nn.Linear(20, 40))\\n        >>> remove_weight_norm(m)\\n    \\\"\\\"\\\"\\n    for k, hook in module._forward_pre_hooks.items():\\n        if isinstance(hook, WeightNorm) and hook.name == name:\\n            hook.remove(module)\\n            del module._forward_pre_hooks[k]\\n            return module\\n\\n    raise ValueError(f\\\"weight_norm of '{name}' not found in {module}\\\")\\n\\n\\nfrom typing import Iterable, Optional\\n\\nimport torch\\n\\n\\ndef parameters_to_vector(parameters: Iterable[torch.Tensor]) -> torch.Tensor:\\n    r\\\"\\\"\\\"Flatten an iterable of parameters into a single vector.\\n\\n    Args:\\n        parameters (Iterable[Tensor]): an iterable of Tensors that are the\\n            parameters of a model.\\n\\n    Returns:\\n        The parameters represented by a single vector\\n    \\\"\\\"\\\"\\n    # Flag for the device where the parameter is located\\n    param_device = None\\n\\n    vec = []\\n    for param in parameters:\\n        # Ensure the parameters are located in the same device\\n        param_device = _check_param_device(param, param_device)\\n\\n        vec.append(param.view(-1))\\n    return torch.cat(vec)\\n\\n\\ndef vector_to_parameters(vec: torch.Tensor, parameters: Iterable[torch.Tensor]) -> None:\\n    r\\\"\\\"\\\"Copy slices of a vector into an iterable of parameters.\\n\\n    Args:\\n        vec (Tensor): a single vector representing the parameters of a model.\\n        parameters (Iterable[Tensor]): an iterable of Tensors that are the\\n            parameters of a model.\\n    \\\"\\\"\\\"\\n    # Ensure vec of type Tensor\\n    if not isinstance(vec, torch.Tensor):\\n        raise TypeError(f\\\"expected torch.Tensor, but got: {torch.typename(vec)}\\\")\\n    # Flag for the device where the parameter is located\\n    param_device = None\\n\\n    # Pointer for slicing the vector for each parameter\\n    pointer = 0\\n    for param in parameters:\\n        # Ensure the parameters are located in the same device\\n        param_device = _check_param_device(param, param_device)\\n\\n        # The length of the parameter\\n        num_param = param.numel()\\n        # Slice the vector, reshape it, and replace the old data of the parameter\\n        param.data = vec[pointer : pointer + num_param].view_as(param).data\\n\\n        # Increment the pointer\\n        pointer += num_param\\n\\n\\ndef _check_param_device(param: torch.Tensor, old_param_device: Optional[int]) -> int:\\n    r\\\"\\\"\\\"Check if the parameters are located on the same device.\\n\\n    Currently, the conversion between model parameters and single vector form is not supported\\n    for multiple allocations, e.g. parameters in different GPUs/PrivateUse1s, or mixture of CPU/GPU/PrivateUse1.\\n\\n    Args:\\n        param ([Tensor]): a Tensor of a parameter of a model\\n        old_param_device (int): the device where the first parameter of a\\n                                model is allocated.\\n\\n    Returns:\\n        old_param_device (int): report device for the first time\\n    \\\"\\\"\\\"\\n    # Meet the first parameter\\n    support_device_types = [\\\"cuda\\\", torch._C._get_privateuse1_backend_name()]\\n    if old_param_device is None:\\n        old_param_device = (\\n            param.get_device() if param.device.type in support_device_types else -1\\n        )\\n    else:\\n        warn = False\\n        if (\\n            param.device.type in support_device_types\\n        ):  # Check if in same GPU/PrivateUse1\\n            warn = param.get_device() != old_param_device\\n        else:  # Check if in CPU\\n            warn = old_param_device != -1\\n        if warn:\\n            raise TypeError(\\n                \\\"Found two parameters on different devices, \\\"\\n                \\\"this is currently not supported.\\\"\\n            )\\n    return old_param_device\\n\\n\\n# This source code is licensed under the BSD-style license found in the\\n# LICENSE file in the root directory of this source tree.\\n\\nfrom typing import Dict, Iterable, List, Tuple\\n\\nimport torch\\n\\n\\n_MISSING: torch.Tensor = object()  # type: ignore[assignment]\\n\\n\\ndef set_tensor(module: \\\"torch.nn.Module\\\", name: str, tensor: torch.Tensor) -> None:\\n    if not isinstance(module, torch.nn.Module):\\n        raise TypeError(f\\\"{module} is not an instance of torch.nn.Module\\\")\\n    if not isinstance(tensor, torch.Tensor) and tensor is not None:\\n        raise TypeError(f\\\"{tensor} is not an instance of torch.Tensor\\\")\\n    if \\\".\\\" in name:\\n        raise KeyError('tensor name can\\\\'t contain \\\".\\\"')\\n    if name == \\\"\\\":\\n        raise KeyError('tensor name can\\\\'t be empty string \\\"\\\"')\\n    if name in module._parameters:\\n        module._parameters[name] = tensor  # type: ignore[assignment]\\n    elif name in module._buffers:\\n        module._buffers[name] = tensor\\n    else:\\n        setattr(module, name, tensor)\\n\\n\\ndef swap_tensor(\\n    module: \\\"torch.nn.Module\\\",\\n    name: str,\\n    tensor: torch.Tensor,\\n    allow_missing: bool = False,\\n) -> torch.Tensor:\\n    if not isinstance(module, torch.nn.Module):\\n        raise TypeError(f\\\"{module} is not an instance of torch.nn.Module\\\")\\n    if (\\n        tensor is not _MISSING\\n        and not isinstance(tensor, torch.Tensor)\\n        and tensor is not None\\n    ):\\n        raise TypeError(f\\\"{tensor} is not an instance of torch.Tensor\\\")\\n    if \\\".\\\" in name:\\n        raise KeyError('tensor name can\\\\'t contain \\\".\\\"')\\n    if name == \\\"\\\":\\n        raise KeyError('tensor name can\\\\'t be empty string \\\"\\\"')\\n\\n    orig_tensor: torch.Tensor\\n    if name in module._parameters:\\n        orig_tensor = module._parameters[name]  # type: ignore[assignment]\\n        if tensor is not _MISSING:\\n            module._parameters[name] = tensor  # type: ignore[assignment]\\n        else:\\n            del module._parameters[name]\\n    elif name in module._buffers:\\n        orig_tensor = module._buffers[name]  # type: ignore[assignment]\\n        if tensor is not _MISSING:\\n            module._buffers[name] = tensor\\n        else:\\n            del module._buffers[name]\\n    else:\\n        if hasattr(module, name):\\n            orig_tensor = getattr(module, name)\\n        else:\\n            if not allow_missing:\\n                raise AttributeError(f\\\"{module._get_name()} has no attribute `{name}`\\\")\\n            orig_tensor = _MISSING\\n        if (\\n            orig_tensor is not _MISSING\\n            and not isinstance(orig_tensor, torch.Tensor)\\n            and orig_tensor is not None\\n        ):\\n            raise TypeError(\\n                f\\\"attribute `{name}`: {orig_tensor} is not an instance of torch.Tensor\\\"\\n            )\\n        if tensor is not _MISSING:\\n            setattr(module, name, tensor)\\n        elif hasattr(module, name):\\n            delattr(module, name)\\n    return orig_tensor\\n\\n\\ndef swap_submodule(\\n    module: \\\"torch.nn.Module\\\",\\n    name: str,\\n    submodule: \\\"torch.nn.Module\\\",\\n) -> \\\"torch.nn.Module\\\":\\n    if not isinstance(module, torch.nn.Module):\\n        raise TypeError(f\\\"{module} is not an instance of torch.nn.Module\\\")\\n    if not isinstance(submodule, torch.nn.Module):\\n        raise TypeError(f\\\"{submodule} is not an instance of torch.nn.Module\\\")\\n    if \\\".\\\" in name:\\n        raise KeyError('submodule name can\\\\'t contain \\\".\\\"')\\n    if name == \\\"\\\":\\n        raise KeyError('submodule name can\\\\'t be empty string \\\"\\\"')\\n    if name not in module._modules:\\n        raise KeyError(f\\\"submodule {name} does not exist\\\")\\n\\n    orig_submodule = module._modules[name]\\n    if not isinstance(orig_submodule, torch.nn.Module):\\n        raise TypeError(f\\\"{name} attribute is not an instance of torch.nn.Module\\\")\\n    module._modules[name] = submodule\\n    return orig_submodule\\n\\n\\nclass NamedMemberAccessor:\\n    \\\"\\\"\\\"\\n    A class that provides a way to access the submodules and parameters/buffers of a module.\\n\\n    It provides caching mechanism to speed up submodule lookups.\\n    This is useful for functional programming to manipulate the module state.\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, module: \\\"torch.nn.Module\\\") -> None:\\n        self.module = module\\n        self.memo: Dict[str, torch.nn.Module] = {}\\n\\n    # Nested attribute access\\n\\n    def get_submodule(self, name: str) -> \\\"torch.nn.Module\\\":\\n        \\\"\\\"\\\"\\n        Return the submodule specified by the given path.\\n\\n        For example, to get the submodule mod.layer1.conv1,\\n        use accessor.get_submodule(\\\"layer1.conv1\\\")\\n\\n        Compare to mod.get_submodule(\\\"layer1.conv1\\\"), this method will cache the\\n        intermediate submodule access to speed up future lookups.\\n        \\\"\\\"\\\"\\n        if not name:\\n            return self.module\\n\\n        if name in self.memo:\\n            return self.memo[name]\\n        else:\\n            prefix, dot, attr = name.rpartition(\\\".\\\")\\n            if dot:\\n                module = self.get_submodule(prefix)\\n            else:\\n                module = self.module\\n            try:\\n                submodule = getattr(module, attr)\\n            except AttributeError as ex:\\n                raise AttributeError(\\n                    f\\\"{module._get_name()} has no attribute `{attr}`\\\"\\n                ) from ex\\n            if not isinstance(submodule, torch.nn.Module):\\n                raise TypeError(  # noqa: B904\\n                    f\\\"submodule `{name}`: {submodule} is not an instance of torch.nn.Module\\\"\\n                )\\n            self.memo[name] = submodule\\n            return submodule\\n\\n    def swap_submodule(self, path: str, value: \\\"torch.nn.Module\\\") -> \\\"torch.nn.Module\\\":\\n        \\\"\\\"\\\"\\n        Swap the submodule specified by the given ``path`` to ``value``.\\n\\n        For example, to swap the attribute mod.layer1.conv1 use\\n        ``accessor.swap_submodule(\\\"layer1.conv1\\\", conv2)``.\\n        \\\"\\\"\\\"\\n        prefix, _, attr = path.rpartition(\\\".\\\")\\n        return swap_submodule(self.get_submodule(prefix), attr, value)\\n\\n    def get_tensor(self, name: str) -> torch.Tensor:\\n        \\\"\\\"\\\"\\n        Get the tensor specified by the given path to value.\\n\\n        For example, to get the attribute mod.layer1.conv1.weight,\\n        use accessor.get_tensor('layer1.conv1.weight')\\n\\n        Compare to mod.get_parameter(\\\"layer1.conv1.weight\\\"), this method will\\n        cache the intermediate submodule access to speed up future lookups.\\n        \\\"\\\"\\\"\\n        prefix, _, attr = name.rpartition(\\\".\\\")\\n        submodule = self.get_submodule(prefix)\\n        try:\\n            tensor = getattr(submodule, attr)\\n        except AttributeError as ex:\\n            raise AttributeError(\\n                f\\\"{submodule._get_name()} has no attribute `{name}`\\\"\\n            ) from ex\\n        if not isinstance(tensor, torch.Tensor) and tensor is not None:\\n            raise TypeError(f\\\"{tensor} is not an instance of torch.Tensor\\\")\\n        return tensor  # type: ignore[return-value]\\n\\n    def set_tensor(self, name: str, value: torch.Tensor) -> None:\\n        \\\"\\\"\\\"\\n        Set the attribute specified by the given path to value.\\n\\n        For example, to set the attribute mod.layer1.conv1.weight,\\n        use accessor.set_tensor(\\\"layer1.conv1.weight\\\", value)\\n        \\\"\\\"\\\"\\n        prefix, _, attr = name.rpartition(\\\".\\\")\\n        set_tensor(self.get_submodule(prefix), attr, value)\\n\\n    def del_tensor(self, name: str) -> None:\\n        \\\"\\\"\\\"\\n        Delete the attribute specified by the given path.\\n\\n        For example, to delete the attribute mod.layer1.conv1.weight,\\n        use accessor.del_tensor(\\\"layer1.conv1.weight\\\")\\n        \\\"\\\"\\\"\\n        prefix, _, attr = name.rpartition(\\\".\\\")\\n        submodule = self.get_submodule(prefix)\\n        try:\\n            delattr(submodule, attr)\\n        except AttributeError as ex:\\n            raise AttributeError(\\n                f\\\"{submodule._get_name()} has no attribute `{name}`\\\"\\n            ) from ex\\n\\n    def swap_tensor(\\n        self, name: str, value: torch.Tensor, allow_missing: bool = False\\n    ) -> torch.Tensor:\\n        \\\"\\\"\\\"\\n        Swap the attribute specified by the given path to value.\\n\\n        For example, to swap the attribute mod.layer1.conv1.weight,\\n        use accessor.swap_tensor(\\\"layer1.conv1.weight\\\", value)\\n        \\\"\\\"\\\"\\n        prefix, _, attr = name.rpartition(\\\".\\\")\\n        return swap_tensor(\\n            self.get_submodule(prefix), attr, value, allow_missing=allow_missing\\n        )\\n\\n    # Batched operations\\n\\n    def get_tensors(self, names: Iterable[str]) -> List[torch.Tensor]:\\n        \\\"\\\"\\\"\\n        Get the tensors specified by the given paths.\\n\\n        For example, to get the attributes mod.layer1.conv1.weight and\\n        mod.layer1.conv1.bias, use accessor.get_tensors([\\\"layer1.conv1.weight\\\",\\n        \\\"layer1.conv1.bias\\\"])\\n        \\\"\\\"\\\"\\n        return [self.get_tensor(name) for name in names]\\n\\n    def set_tensors(self, names: Iterable[str], values: Iterable[torch.Tensor]) -> None:\\n        \\\"\\\"\\\"\\n        Set the attributes specified by the given paths to values.\\n\\n        For example, to set the attributes mod.layer1.conv1.weight and\\n        mod.layer1.conv1.bias, use accessor.set_tensors([\\\"layer1.conv1.weight\\\",\\n        \\\"layer1.conv1.bias\\\"], [weight, bias])\\n        \\\"\\\"\\\"\\n        if not isinstance(names, (list, tuple)):\\n            names = list(names)\\n        if not isinstance(values, (list, tuple)):\\n            values = list(values)\\n        assert len(names) == len(values), \\\"names and values must have the same length\\\"\\n\\n        for name, value in zip(names, values):\\n            self.set_tensor(name, value)\\n\\n    def set_tensors_dict(self, named_tensors: Dict[str, torch.Tensor]) -> None:\\n        \\\"\\\"\\\"\\n        Set the attributes specified by the given paths to values.\\n\\n        For example, to set the attributes mod.layer1.conv1.weight and\\n        mod.layer1.conv1.bias, use accessor.set_tensors_dict({\\n            \\\"layer1.conv1.weight\\\": weight,\\n            \\\"layer1.conv1.bias\\\": bias,\\n        })\\n        \\\"\\\"\\\"\\n        for name, value in named_tensors.items():\\n            self.set_tensor(name, value)\\n\\n    def del_tensors(self, names: Iterable[str]) -> None:\\n        \\\"\\\"\\\"\\n        Delete the attributes specified by the given paths.\\n\\n        For example, to delete the attributes mod.layer1.conv1.weight and\\n        mod.layer1.conv1.bias, use accessor.del_tensors([\\\"layer1.conv1.weight\\\",\\n        \\\"layer1.conv1.bias\\\"])\\n        \\\"\\\"\\\"\\n        for name in names:\\n            self.del_tensor(name)\\n\\n    def swap_tensors(\\n        self,\\n        names: Iterable[str],\\n        values: Iterable[torch.Tensor],\\n        allow_missing: bool = False,\\n    ) -> List[torch.Tensor]:\\n        \\\"\\\"\\\"\\n        Swap the attributes specified by the given paths to values.\\n\\n        For example, to swap the attributes mod.layer1.conv1.weight and\\n        mod.layer1.conv1.bias, use accessor.swap_tensors([\\\"layer1.conv1.weight\\\",\\n        \\\"layer1.conv1.bias\\\"], [weight, bias])\\n        \\\"\\\"\\\"\\n        if not isinstance(names, (list, tuple)):\\n            names = list(names)\\n        if not isinstance(values, (list, tuple)):\\n            values = list(values)\\n        assert len(names) == len(values), \\\"names and values must have the same length\\\"\\n\\n        return [\\n            self.swap_tensor(name, value, allow_missing=allow_missing)\\n            for name, value in zip(names, values)\\n        ]\\n\\n    def swap_tensors_dict(\\n        self, named_tensors: Dict[str, torch.Tensor], allow_missing: bool = False\\n    ) -> Tuple[Dict[str, torch.Tensor], List[str]]:\\n        \\\"\\\"\\\"\\n        Swap the attributes specified by the given paths to values.\\n\\n        For example, to swap the attributes mod.layer1.conv1.weight and\\n        mod.layer1.conv1.bias, use accessor.swap_tensors_dict({\\n            \\\"layer1.conv1.weight\\\": weight,\\n            \\\"layer1.conv1.bias\\\": bias,\\n        })\\n        \\\"\\\"\\\"\\n        orig_named_tensors = {}\\n        missing_keys = []\\n        try:\\n            for name, tensor in named_tensors.items():\\n                orig_tensor = self.swap_tensor(name, tensor, allow_missing=True)\\n                if orig_tensor is _MISSING:\\n                    missing_keys.append(name)\\n                orig_named_tensors[name] = orig_tensor\\n        except Exception:\\n            # Swap back if any exception occurs\\n            for name, orig_tensor in orig_named_tensors.items():\\n                self.swap_tensor(name, orig_tensor, allow_missing=True)\\n            raise\\n        if missing_keys and not allow_missing:\\n            # Swap back if any key is missing when allow_missing is False\\n            for name, orig_tensor in orig_named_tensors.items():\\n                self.swap_tensor(name, orig_tensor, allow_missing=True)\\n            raise RuntimeError(f\\\"Missing key(s): {', '.join(map(repr, missing_keys))}.\\\")\\n        return orig_named_tensors, missing_keys\\n\\n    def check_keys(self, keys: Iterable[str]) -> Tuple[List[str], List[str]]:\\n        \\\"\\\"\\\"Check that the given keys are valid.\\\"\\\"\\\"\\n        keys = set(keys)\\n        valid_keys = {name for name, _ in self.named_tensors(remove_duplicate=False)}\\n        missing_keys = valid_keys - keys\\n        unexpected_keys = keys - valid_keys\\n        return sorted(missing_keys), sorted(unexpected_keys)\\n\\n    # Shortcut methods\\n\\n    def named_parameters(\\n        self,\\n        remove_duplicate: bool = True,\\n    ) -> Iterable[Tuple[str, torch.Tensor]]:\\n        \\\"\\\"\\\"Iterate over all the parameters in the module.\\\"\\\"\\\"\\n        yield from self.module.named_parameters(remove_duplicate=remove_duplicate)\\n\\n    def named_buffers(\\n        self,\\n        remove_duplicate: bool = True,\\n    ) -> Iterable[Tuple[str, torch.Tensor]]:\\n        \\\"\\\"\\\"Iterate over all the buffers in the module.\\\"\\\"\\\"\\n        yield from self.module.named_buffers(remove_duplicate=remove_duplicate)\\n\\n    def named_tensors(\\n        self,\\n        remove_duplicate: bool = True,\\n    ) -> Iterable[Tuple[str, torch.Tensor]]:\\n        \\\"\\\"\\\"Iterate over all the tensors in the module.\\\"\\\"\\\"\\n        yield from self.module.named_parameters(remove_duplicate=remove_duplicate)\\n        yield from self.module.named_buffers(remove_duplicate=remove_duplicate)\\n\\n    def named_modules(\\n        self,\\n        remove_duplicate: bool = True,\\n    ) -> Iterable[Tuple[str, \\\"torch.nn.Module\\\"]]:\\n        \\\"\\\"\\\"Iterate over all the modules in the module.\\\"\\\"\\\"\\n        yield from self.module.named_modules(remove_duplicate=remove_duplicate)\\n\\n\\n# mypy: allow-untyped-defs\\nr\\\"\\\"\\\"Pruning methods.\\\"\\\"\\\"\\nimport numbers\\nfrom abc import ABC, abstractmethod\\nfrom collections.abc import Iterable\\nfrom typing import Tuple\\n\\nimport torch\\n\\n\\nclass BasePruningMethod(ABC):\\n    r\\\"\\\"\\\"Abstract base class for creation of new pruning techniques.\\n\\n    Provides a skeleton for customization requiring the overriding of methods\\n    such as :meth:`compute_mask` and :meth:`apply`.\\n    \\\"\\\"\\\"\\n\\n    _tensor_name: str\\n\\n    def __call__(self, module, inputs):\\n        r\\\"\\\"\\\"Multiply the mask into original tensor and store the result.\\n\\n        Multiplies the mask (stored in ``module[name + '_mask']``)\\n        into the original tensor (stored in ``module[name + '_orig']``)\\n        and stores the result into ``module[name]`` by using :meth:`apply_mask`.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            inputs: not used.\\n        \\\"\\\"\\\"\\n        setattr(module, self._tensor_name, self.apply_mask(module))\\n\\n    @abstractmethod\\n    def compute_mask(self, t, default_mask):\\n        r\\\"\\\"\\\"Compute and returns a mask for the input tensor ``t``.\\n\\n        Starting from a base ``default_mask`` (which should be a mask of ones\\n        if the tensor has not been pruned yet), generate a random mask to\\n        apply on top of the ``default_mask`` according to the specific pruning\\n        method recipe.\\n\\n        Args:\\n            t (torch.Tensor): tensor representing the importance scores of the\\n            parameter to prune.\\n            default_mask (torch.Tensor): Base mask from previous pruning\\n            iterations, that need to be respected after the new mask is\\n            applied. Same dims as ``t``.\\n\\n        Returns:\\n            mask (torch.Tensor): mask to apply to ``t``, of same dims as ``t``\\n        \\\"\\\"\\\"\\n\\n    def apply_mask(self, module):\\n        r\\\"\\\"\\\"Simply handles the multiplication between the parameter being pruned and the generated mask.\\n\\n        Fetches the mask and the original tensor from the module\\n        and returns the pruned version of the tensor.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n\\n        Returns:\\n            pruned_tensor (torch.Tensor): pruned version of the input tensor\\n        \\\"\\\"\\\"\\n        # to carry out the multiplication, the mask needs to have been computed,\\n        # so the pruning method must know what tensor it's operating on\\n        assert (\\n            self._tensor_name is not None\\n        ), f\\\"Module {module} has to be pruned\\\"  # this gets set in apply()\\n        mask = getattr(module, self._tensor_name + \\\"_mask\\\")\\n        orig = getattr(module, self._tensor_name + \\\"_orig\\\")\\n        pruned_tensor = mask.to(dtype=orig.dtype) * orig\\n        return pruned_tensor\\n\\n    @classmethod\\n    def apply(cls, module, name, *args, importance_scores=None, **kwargs):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n            args: arguments passed on to a subclass of\\n                :class:`BasePruningMethod`\\n            importance_scores (torch.Tensor): tensor of importance scores (of\\n                same shape as module parameter) used to compute mask for pruning.\\n                The values in this tensor indicate the importance of the\\n                corresponding elements in the parameter being pruned.\\n                If unspecified or None, the parameter will be used in its place.\\n            kwargs: keyword arguments passed on to a subclass of a\\n                :class:`BasePruningMethod`\\n        \\\"\\\"\\\"\\n\\n        def _get_composite_method(cls, module, name, *args, **kwargs):\\n            # Check if a pruning method has already been applied to\\n            # `module[name]`. If so, store that in `old_method`.\\n            old_method = None\\n            found = 0\\n            # there should technically be only 1 hook with hook.name == name\\n            # assert this using `found`\\n            hooks_to_remove = []\\n            for k, hook in module._forward_pre_hooks.items():\\n                # if it exists, take existing thing, remove hook, then\\n                # go through normal thing\\n                if isinstance(hook, BasePruningMethod) and hook._tensor_name == name:\\n                    old_method = hook\\n                    hooks_to_remove.append(k)\\n                    found += 1\\n            assert (\\n                found <= 1\\n            ), f\\\"Avoid adding multiple pruning hooks to the\\\\\\n                same tensor {name} of module {module}. Use a PruningContainer.\\\"\\n\\n            for k in hooks_to_remove:\\n                del module._forward_pre_hooks[k]\\n\\n            # Apply the new pruning method, either from scratch or on top of\\n            # the previous one.\\n            method = cls(*args, **kwargs)  # new pruning\\n            # Have the pruning method remember what tensor it's been applied to\\n            method._tensor_name = name\\n\\n            # combine `methods` with `old_method`, if `old_method` exists\\n            if old_method is not None:  # meaning that there was a hook\\n                # if the hook is already a pruning container, just add the\\n                # new pruning method to the container\\n                if isinstance(old_method, PruningContainer):\\n                    old_method.add_pruning_method(method)\\n                    method = old_method  # rename old_method --> method\\n\\n                # if the hook is simply a single pruning method, create a\\n                # container, add the old pruning method and the new one\\n                elif isinstance(old_method, BasePruningMethod):\\n                    container = PruningContainer(old_method)\\n                    # Have the pruning method remember the name of its tensor\\n                    # setattr(container, '_tensor_name', name)\\n                    container.add_pruning_method(method)\\n                    method = container  # rename container --> method\\n            return method\\n\\n        method = _get_composite_method(cls, module, name, *args, **kwargs)\\n        # at this point we have no forward_pre_hooks but we could have an\\n        # active reparametrization of the tensor if another pruning method\\n        # had been applied (in which case `method` would be a PruningContainer\\n        # and not a simple pruning method).\\n\\n        # Pruning is to be applied to the module's tensor named `name`,\\n        # starting from the state it is found in prior to this iteration of\\n        # pruning. The pruning mask is calculated based on importances scores.\\n\\n        orig = getattr(module, name)\\n        if importance_scores is not None:\\n            assert (\\n                importance_scores.shape == orig.shape\\n            ), f\\\"importance_scores should have the same shape as parameter                 {name} of {module}\\\"\\n        else:\\n            importance_scores = orig\\n\\n        # If this is the first time pruning is applied, take care of moving\\n        # the original tensor to a new parameter called name + '_orig' and\\n        # and deleting the original parameter\\n        if not isinstance(method, PruningContainer):\\n            # copy `module[name]` to `module[name + '_orig']`\\n            module.register_parameter(name + \\\"_orig\\\", orig)\\n            # temporarily delete `module[name]`\\n            del module._parameters[name]\\n            default_mask = torch.ones_like(orig)  # temp\\n        # If this is not the first time pruning is applied, all of the above\\n        # has been done before in a previous pruning iteration, so we're good\\n        # to go\\n        else:\\n            default_mask = (\\n                getattr(module, name + \\\"_mask\\\")\\n                .detach()\\n                .clone(memory_format=torch.contiguous_format)\\n            )\\n\\n        # Use try/except because if anything goes wrong with the mask\\n        # computation etc., you'd want to roll back.\\n        try:\\n            # get the final mask, computed according to the specific method\\n            mask = method.compute_mask(importance_scores, default_mask=default_mask)\\n            # reparameterize by saving mask to `module[name + '_mask']`...\\n            module.register_buffer(name + \\\"_mask\\\", mask)\\n            # ... and the new pruned tensor to `module[name]`\\n            setattr(module, name, method.apply_mask(module))\\n            # associate the pruning method to the module via a hook to\\n            # compute the function before every forward() (compile by run)\\n            module.register_forward_pre_hook(method)\\n\\n        except Exception as e:\\n            if not isinstance(method, PruningContainer):\\n                orig = getattr(module, name + \\\"_orig\\\")\\n                module.register_parameter(name, orig)\\n                del module._parameters[name + \\\"_orig\\\"]\\n            raise e\\n\\n        return method\\n\\n    def prune(self, t, default_mask=None, importance_scores=None):\\n        r\\\"\\\"\\\"Compute and returns a pruned version of input tensor ``t``.\\n\\n        According to the pruning rule specified in :meth:`compute_mask`.\\n\\n        Args:\\n            t (torch.Tensor): tensor to prune (of same dimensions as\\n                ``default_mask``).\\n            importance_scores (torch.Tensor): tensor of importance scores (of\\n                same shape as ``t``) used to compute mask for pruning ``t``.\\n                The values in this tensor indicate the importance of the\\n                corresponding elements in the ``t`` that is being pruned.\\n                If unspecified or None, the tensor ``t`` will be used in its place.\\n            default_mask (torch.Tensor, optional): mask from previous pruning\\n                iteration, if any. To be considered when determining what\\n                portion of the tensor that pruning should act on. If None,\\n                default to a mask of ones.\\n\\n        Returns:\\n            pruned version of tensor ``t``.\\n        \\\"\\\"\\\"\\n        if importance_scores is not None:\\n            assert (\\n                importance_scores.shape == t.shape\\n            ), \\\"importance_scores should have the same shape as tensor t\\\"\\n        else:\\n            importance_scores = t\\n        default_mask = default_mask if default_mask is not None else torch.ones_like(t)\\n        return t * self.compute_mask(importance_scores, default_mask=default_mask)\\n\\n    def remove(self, module):\\n        r\\\"\\\"\\\"Remove the pruning reparameterization from a module.\\n\\n        The pruned parameter named ``name`` remains permanently pruned,\\n        and the parameter named ``name+'_orig'`` is removed from the parameter list.\\n        Similarly, the buffer named ``name+'_mask'`` is removed from the buffers.\\n\\n        Note:\\n            Pruning itself is NOT undone or reversed!\\n        \\\"\\\"\\\"\\n        # before removing pruning from a tensor, it has to have been applied\\n        assert (\\n            self._tensor_name is not None\\n        ), f\\\"Module {module} has to be pruned            before pruning can be removed\\\"  # this gets set in apply()\\n\\n        # to update module[name] to latest trained weights\\n        weight = self.apply_mask(module)  # masked weights\\n\\n        # delete and reset\\n        if hasattr(module, self._tensor_name):\\n            delattr(module, self._tensor_name)\\n        orig = module._parameters[self._tensor_name + \\\"_orig\\\"]\\n        orig.data = weight.data\\n        del module._parameters[self._tensor_name + \\\"_orig\\\"]\\n        del module._buffers[self._tensor_name + \\\"_mask\\\"]\\n        setattr(module, self._tensor_name, orig)\\n\\n\\nclass PruningContainer(BasePruningMethod):\\n    \\\"\\\"\\\"Container holding a sequence of pruning methods for iterative pruning.\\n\\n    Keeps track of the order in which pruning methods are applied and handles\\n    combining successive pruning calls.\\n\\n    Accepts as argument an instance of a BasePruningMethod or an iterable of\\n    them.\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, *args):\\n        self._pruning_methods: Tuple[BasePruningMethod, ...] = ()\\n        if not isinstance(args, Iterable):  # only 1 item\\n            self._tensor_name = args._tensor_name\\n            self.add_pruning_method(args)\\n        elif len(args) == 1:  # only 1 item in a tuple\\n            self._tensor_name = args[0]._tensor_name\\n            self.add_pruning_method(args[0])\\n        else:  # manual construction from list or other iterable (or no args)\\n            for method in args:\\n                self.add_pruning_method(method)\\n\\n    def add_pruning_method(self, method):\\n        r\\\"\\\"\\\"Add a child pruning ``method`` to the container.\\n\\n        Args:\\n            method (subclass of BasePruningMethod): child pruning method\\n                to be added to the container.\\n        \\\"\\\"\\\"\\n        # check that we're adding a pruning method to the container\\n        if not isinstance(method, BasePruningMethod) and method is not None:\\n            raise TypeError(f\\\"{type(method)} is not a BasePruningMethod subclass\\\")\\n        elif method is not None and self._tensor_name != method._tensor_name:\\n            raise ValueError(\\n                \\\"Can only add pruning methods acting on \\\"\\n                f\\\"the parameter named '{self._tensor_name}' to PruningContainer {self}.\\\"\\n                + f\\\" Found '{method._tensor_name}'\\\"\\n            )\\n        # if all checks passed, add to _pruning_methods tuple\\n        self._pruning_methods += (method,)  # type: ignore[operator]\\n\\n    def __len__(self):\\n        return len(self._pruning_methods)\\n\\n    def __iter__(self):\\n        return iter(self._pruning_methods)\\n\\n    def __getitem__(self, idx):\\n        return self._pruning_methods[idx]\\n\\n    def compute_mask(self, t, default_mask):\\n        r\\\"\\\"\\\"Apply the latest ``method`` by computing the new partial masks and returning its combination with the ``default_mask``.\\n\\n        The new partial mask should be computed on the entries or channels\\n        that were not zeroed out by the ``default_mask``.\\n        Which portions of the tensor ``t`` the new mask will be calculated from\\n        depends on the ``PRUNING_TYPE`` (handled by the type handler):\\n\\n        * for 'unstructured', the mask will be computed from the raveled\\n          list of nonmasked entries;\\n\\n        * for 'structured', the mask will be computed from the nonmasked\\n          channels in the tensor;\\n\\n        * for 'global', the mask will be computed across all entries.\\n\\n        Args:\\n            t (torch.Tensor): tensor representing the parameter to prune\\n                (of same dimensions as ``default_mask``).\\n            default_mask (torch.Tensor): mask from previous pruning iteration.\\n\\n        Returns:\\n            mask (torch.Tensor): new mask that combines the effects\\n            of the ``default_mask`` and the new mask from the current\\n            pruning ``method`` (of same dimensions as ``default_mask`` and\\n            ``t``).\\n        \\\"\\\"\\\"\\n\\n        def _combine_masks(method, t, mask):\\n            r\\\"\\\"\\\"Combine the masks from all pruning methods and returns a new mask.\\n\\n            Args:\\n                method (a BasePruningMethod subclass): pruning method\\n                    currently being applied.\\n                t (torch.Tensor): tensor representing the parameter to prune\\n                    (of same dimensions as mask).\\n                mask (torch.Tensor): mask from previous pruning iteration\\n\\n            Returns:\\n                new_mask (torch.Tensor): new mask that combines the effects\\n                    of the old mask and the new mask from the current\\n                    pruning method (of same dimensions as mask and t).\\n            \\\"\\\"\\\"\\n            new_mask = mask  # start off from existing mask\\n            new_mask = new_mask.to(dtype=t.dtype)\\n\\n            # compute a slice of t onto which the new pruning method will operate\\n            if method.PRUNING_TYPE == \\\"unstructured\\\":\\n                # prune entries of t where the mask is 1\\n                slc = mask == 1\\n\\n            # for struct pruning, exclude channels that have already been\\n            # entirely pruned\\n            elif method.PRUNING_TYPE == \\\"structured\\\":\\n                if not hasattr(method, \\\"dim\\\"):\\n                    raise AttributeError(\\n                        \\\"Pruning methods of PRUNING_TYPE \\\"\\n                        '\\\"structured\\\" need to have the attribute `dim` defined.'\\n                    )\\n\\n                # find the channels to keep by removing the ones that have been\\n                # zeroed out already (i.e. where sum(entries) == 0)\\n                n_dims = t.dim()  # \\\"is this a 2D tensor? 3D? ...\\\"\\n                dim = method.dim\\n                # convert negative indexing\\n                if dim < 0:\\n                    dim = n_dims + dim\\n                # if dim is still negative after subtracting it from n_dims\\n                if dim < 0:\\n                    raise IndexError(\\n                        f\\\"Index is out of bounds for tensor with dimensions {n_dims}\\\"\\n                    )\\n                # find channels along dim = dim that aren't already tots 0ed out\\n                keep_channel = mask.sum(dim=[d for d in range(n_dims) if d != dim]) != 0\\n                # create slice to identify what to prune\\n                slc = [slice(None)] * n_dims\\n                slc[dim] = keep_channel\\n\\n            elif method.PRUNING_TYPE == \\\"global\\\":\\n                n_dims = len(t.shape)  # \\\"is this a 2D tensor? 3D? ...\\\"\\n                slc = [slice(None)] * n_dims\\n\\n            else:\\n                raise ValueError(f\\\"Unrecognized PRUNING_TYPE {method.PRUNING_TYPE}\\\")\\n\\n            # compute the new mask on the unpruned slice of the tensor t\\n            partial_mask = method.compute_mask(t[slc], default_mask=mask[slc])\\n            new_mask[slc] = partial_mask.to(dtype=new_mask.dtype)\\n\\n            return new_mask\\n\\n        method = self._pruning_methods[-1]\\n        mask = _combine_masks(method, t, default_mask)\\n        return mask\\n\\n\\nclass Identity(BasePruningMethod):\\n    r\\\"\\\"\\\"Utility pruning method that does not prune any units but generates the pruning parametrization with a mask of ones.\\\"\\\"\\\"\\n\\n    PRUNING_TYPE = \\\"unstructured\\\"\\n\\n    def compute_mask(self, t, default_mask):\\n        mask = default_mask\\n        return mask\\n\\n    @classmethod\\n    def apply(cls, module, name):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n        \\\"\\\"\\\"\\n        return super().apply(module, name)\\n\\n\\nclass RandomUnstructured(BasePruningMethod):\\n    r\\\"\\\"\\\"Prune (currently unpruned) units in a tensor at random.\\n\\n    Args:\\n        name (str): parameter name within ``module`` on which pruning\\n            will act.\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n    \\\"\\\"\\\"\\n\\n    PRUNING_TYPE = \\\"unstructured\\\"\\n\\n    def __init__(self, amount):\\n        # Check range of validity of pruning amount\\n        _validate_pruning_amount_init(amount)\\n        self.amount = amount\\n\\n    def compute_mask(self, t, default_mask):\\n        # Check that the amount of units to prune is not > than the number of\\n        # parameters in t\\n        tensor_size = t.nelement()\\n        # Compute number of units to prune: amount if int,\\n        # else amount * tensor_size\\n        nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size)\\n        # This should raise an error if the number of units to prune is larger\\n        # than the number of units in the tensor\\n        _validate_pruning_amount(nparams_toprune, tensor_size)\\n\\n        mask = default_mask.clone(memory_format=torch.contiguous_format)\\n\\n        if nparams_toprune != 0:  # k=0 not supported by torch.kthvalue\\n            prob = torch.rand_like(t)\\n            topk = torch.topk(prob.view(-1), k=nparams_toprune)\\n            mask.view(-1)[topk.indices] = 0\\n\\n        return mask\\n\\n    @classmethod\\n    def apply(cls, module, name, amount):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n            amount (int or float): quantity of parameters to prune.\\n                If ``float``, should be between 0.0 and 1.0 and represent the\\n                fraction of parameters to prune. If ``int``, it represents the\\n                absolute number of parameters to prune.\\n        \\\"\\\"\\\"\\n        return super().apply(module, name, amount=amount)\\n\\n\\nclass L1Unstructured(BasePruningMethod):\\n    r\\\"\\\"\\\"Prune (currently unpruned) units in a tensor by zeroing out the ones with the lowest L1-norm.\\n\\n    Args:\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n    \\\"\\\"\\\"\\n\\n    PRUNING_TYPE = \\\"unstructured\\\"\\n\\n    def __init__(self, amount):\\n        # Check range of validity of pruning amount\\n        _validate_pruning_amount_init(amount)\\n        self.amount = amount\\n\\n    def compute_mask(self, t, default_mask):\\n        # Check that the amount of units to prune is not > than the number of\\n        # parameters in t\\n        tensor_size = t.nelement()\\n        # Compute number of units to prune: amount if int,\\n        # else amount * tensor_size\\n        nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size)\\n        # This should raise an error if the number of units to prune is larger\\n        # than the number of units in the tensor\\n        _validate_pruning_amount(nparams_toprune, tensor_size)\\n\\n        mask = default_mask.clone(memory_format=torch.contiguous_format)\\n\\n        if nparams_toprune != 0:  # k=0 not supported by torch.kthvalue\\n            # largest=True --> top k; largest=False --> bottom k\\n            # Prune the smallest k\\n            topk = torch.topk(torch.abs(t).view(-1), k=nparams_toprune, largest=False)\\n            # topk will have .indices and .values\\n            mask.view(-1)[topk.indices] = 0\\n\\n        return mask\\n\\n    @classmethod\\n    def apply(cls, module, name, amount, importance_scores=None):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n            amount (int or float): quantity of parameters to prune.\\n                If ``float``, should be between 0.0 and 1.0 and represent the\\n                fraction of parameters to prune. If ``int``, it represents the\\n                absolute number of parameters to prune.\\n            importance_scores (torch.Tensor): tensor of importance scores (of same\\n                shape as module parameter) used to compute mask for pruning.\\n                The values in this tensor indicate the importance of the corresponding\\n                elements in the parameter being pruned.\\n                If unspecified or None, the module parameter will be used in its place.\\n        \\\"\\\"\\\"\\n        return super().apply(\\n            module, name, amount=amount, importance_scores=importance_scores\\n        )\\n\\n\\nclass RandomStructured(BasePruningMethod):\\n    r\\\"\\\"\\\"Prune entire (currently unpruned) channels in a tensor at random.\\n\\n    Args:\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n        dim (int, optional): index of the dim along which we define\\n            channels to prune. Default: -1.\\n    \\\"\\\"\\\"\\n\\n    PRUNING_TYPE = \\\"structured\\\"\\n\\n    def __init__(self, amount, dim=-1):\\n        # Check range of validity of amount\\n        _validate_pruning_amount_init(amount)\\n        self.amount = amount\\n        self.dim = dim\\n\\n    def compute_mask(self, t, default_mask):\\n        r\\\"\\\"\\\"Compute and returns a mask for the input tensor ``t``.\\n\\n        Starting from a base ``default_mask`` (which should be a mask of ones\\n        if the tensor has not been pruned yet), generate a random mask to\\n        apply on top of the ``default_mask`` by randomly zeroing out channels\\n        along the specified dim of the tensor.\\n\\n        Args:\\n            t (torch.Tensor): tensor representing the parameter to prune\\n            default_mask (torch.Tensor): Base mask from previous pruning\\n                iterations, that need to be respected after the new mask is\\n                applied. Same dims as ``t``.\\n\\n        Returns:\\n            mask (torch.Tensor): mask to apply to ``t``, of same dims as ``t``\\n\\n        Raises:\\n            IndexError: if ``self.dim >= len(t.shape)``\\n        \\\"\\\"\\\"\\n        # Check that tensor has structure (i.e. more than 1 dimension) such\\n        # that the concept of \\\"channels\\\" makes sense\\n        _validate_structured_pruning(t)\\n\\n        # Check that self.dim is a valid dim to index t, else raise IndexError\\n        _validate_pruning_dim(t, self.dim)\\n\\n        # Check that the amount of channels to prune is not > than the number of\\n        # channels in t along the dim to prune\\n        tensor_size = t.shape[self.dim]\\n        # Compute number of units to prune: amount if int,\\n        # else amount * tensor_size\\n        nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size)\\n        # This should raise an error if the number of units to prune is larger\\n        # than the number of units in the tensor\\n        _validate_pruning_amount(nparams_toprune, tensor_size)\\n\\n        # Compute binary mask by initializing it to all 0s and then filling in\\n        # 1s wherever topk.indices indicates, along self.dim.\\n        # mask has the same shape as tensor t\\n        def make_mask(t, dim, nchannels, nchannels_toprune):\\n            # generate a random number in [0, 1] to associate to each channel\\n            prob = torch.rand(nchannels)\\n            # generate mask for each channel by 0ing out the channels that\\n            # got assigned the k = nchannels_toprune lowest values in prob\\n            threshold = torch.kthvalue(prob, k=nchannels_toprune).values\\n            channel_mask = prob > threshold\\n\\n            mask = torch.zeros_like(t)\\n            slc = [slice(None)] * len(t.shape)\\n            slc[dim] = channel_mask\\n            mask[slc] = 1\\n            return mask\\n\\n        if nparams_toprune == 0:  # k=0 not supported by torch.kthvalue\\n            mask = default_mask\\n        else:\\n            # apply the new structured mask on top of prior (potentially\\n            # unstructured) mask\\n            mask = make_mask(t, self.dim, tensor_size, nparams_toprune)\\n            mask *= default_mask.to(dtype=mask.dtype)\\n        return mask\\n\\n    @classmethod\\n    def apply(cls, module, name, amount, dim=-1):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n            amount (int or float): quantity of parameters to prune.\\n                If ``float``, should be between 0.0 and 1.0 and represent the\\n                fraction of parameters to prune. If ``int``, it represents the\\n                absolute number of parameters to prune.\\n            dim (int, optional): index of the dim along which we define\\n                channels to prune. Default: -1.\\n        \\\"\\\"\\\"\\n        return super().apply(module, name, amount=amount, dim=dim)\\n\\n\\nclass LnStructured(BasePruningMethod):\\n    r\\\"\\\"\\\"Prune entire (currently unpruned) channels in a tensor based on their L\\\\ ``n``-norm.\\n\\n    Args:\\n        amount (int or float): quantity of channels to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n        n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid\\n            entries for argument ``p`` in :func:`torch.norm`.\\n        dim (int, optional): index of the dim along which we define\\n            channels to prune. Default: -1.\\n    \\\"\\\"\\\"\\n\\n    PRUNING_TYPE = \\\"structured\\\"\\n\\n    def __init__(self, amount, n, dim=-1):\\n        # Check range of validity of amount\\n        _validate_pruning_amount_init(amount)\\n        self.amount = amount\\n        self.n = n\\n        self.dim = dim\\n\\n    def compute_mask(self, t, default_mask):\\n        r\\\"\\\"\\\"Compute and returns a mask for the input tensor ``t``.\\n\\n        Starting from a base ``default_mask`` (which should be a mask of ones\\n        if the tensor has not been pruned yet), generate a mask to apply on\\n        top of the ``default_mask`` by zeroing out the channels along the\\n        specified dim with the lowest L\\\\ ``n``-norm.\\n\\n        Args:\\n            t (torch.Tensor): tensor representing the parameter to prune\\n            default_mask (torch.Tensor): Base mask from previous pruning\\n                iterations, that need to be respected after the new mask is\\n                applied.  Same dims as ``t``.\\n\\n        Returns:\\n            mask (torch.Tensor): mask to apply to ``t``, of same dims as ``t``\\n\\n        Raises:\\n            IndexError: if ``self.dim >= len(t.shape)``\\n        \\\"\\\"\\\"\\n        # Check that tensor has structure (i.e. more than 1 dimension) such\\n        # that the concept of \\\"channels\\\" makes sense\\n        _validate_structured_pruning(t)\\n        # Check that self.dim is a valid dim to index t, else raise IndexError\\n        _validate_pruning_dim(t, self.dim)\\n\\n        # Check that the amount of channels to prune is not > than the number of\\n        # channels in t along the dim to prune\\n        tensor_size = t.shape[self.dim]\\n        # Compute number of units to prune: amount if int,\\n        # else amount * tensor_size\\n        nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size)\\n        nparams_tokeep = tensor_size - nparams_toprune\\n        # This should raise an error if the number of units to prune is larger\\n        # than the number of units in the tensor\\n        _validate_pruning_amount(nparams_toprune, tensor_size)\\n\\n        # Structured pruning prunes entire channels so we need to know the\\n        # L_n norm along each channel to then find the topk based on this\\n        # metric\\n        norm = _compute_norm(t, self.n, self.dim)\\n        # largest=True --> top k; largest=False --> bottom k\\n        # Keep the largest k channels along dim=self.dim\\n        topk = torch.topk(norm, k=nparams_tokeep, largest=True)\\n        # topk will have .indices and .values\\n\\n        # Compute binary mask by initializing it to all 0s and then filling in\\n        # 1s wherever topk.indices indicates, along self.dim.\\n        # mask has the same shape as tensor t\\n        def make_mask(t, dim, indices):\\n            # init mask to 0\\n            mask = torch.zeros_like(t)\\n            # e.g.: slc = [None, None, None], if len(t.shape) = 3\\n            slc = [slice(None)] * len(t.shape)\\n            # replace a None at position=dim with indices\\n            # e.g.: slc = [None, None, [0, 2, 3]] if dim=2 & indices=[0,2,3]\\n            slc[dim] = indices\\n            # use slc to slice mask and replace all its entries with 1s\\n            # e.g.: mask[:, :, [0, 2, 3]] = 1\\n            mask[slc] = 1\\n            return mask\\n\\n        if nparams_toprune == 0:  # k=0 not supported by torch.kthvalue\\n            mask = default_mask\\n        else:\\n            mask = make_mask(t, self.dim, topk.indices)\\n            mask *= default_mask.to(dtype=mask.dtype)\\n\\n        return mask\\n\\n    @classmethod\\n    def apply(cls, module, name, amount, n, dim, importance_scores=None):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n            amount (int or float): quantity of parameters to prune.\\n                If ``float``, should be between 0.0 and 1.0 and represent the\\n                fraction of parameters to prune. If ``int``, it represents the\\n                absolute number of parameters to prune.\\n            n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid\\n                entries for argument ``p`` in :func:`torch.norm`.\\n            dim (int): index of the dim along which we define channels to\\n                prune.\\n            importance_scores (torch.Tensor): tensor of importance scores (of same\\n                shape as module parameter) used to compute mask for pruning.\\n                The values in this tensor indicate the importance of the corresponding\\n                elements in the parameter being pruned.\\n                If unspecified or None, the module parameter will be used in its place.\\n        \\\"\\\"\\\"\\n        return super().apply(\\n            module,\\n            name,\\n            amount=amount,\\n            n=n,\\n            dim=dim,\\n            importance_scores=importance_scores,\\n        )\\n\\n\\nclass CustomFromMask(BasePruningMethod):\\n    PRUNING_TYPE = \\\"global\\\"\\n\\n    def __init__(self, mask):\\n        self.mask = mask\\n\\n    def compute_mask(self, t, default_mask):\\n        assert default_mask.shape == self.mask.shape\\n        mask = default_mask * self.mask.to(dtype=default_mask.dtype)\\n        return mask\\n\\n    @classmethod\\n    def apply(cls, module, name, mask):\\n        r\\\"\\\"\\\"Add pruning on the fly and reparametrization of a tensor.\\n\\n        Adds the forward pre-hook that enables pruning on the fly and\\n        the reparametrization of a tensor in terms of the original tensor\\n        and the pruning mask.\\n\\n        Args:\\n            module (nn.Module): module containing the tensor to prune\\n            name (str): parameter name within ``module`` on which pruning\\n                will act.\\n        \\\"\\\"\\\"\\n        return super().apply(module, name, mask=mask)\\n\\n\\ndef identity(module, name):\\n    r\\\"\\\"\\\"Apply pruning reparametrization without pruning any units.\\n\\n    Applies pruning reparametrization to the tensor corresponding to the\\n    parameter called ``name`` in ``module`` without actually pruning any\\n    units. Modifies module in place (and also return the modified module)\\n    by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Note:\\n        The mask is a tensor of ones.\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune.\\n        name (str): parameter name within ``module`` on which pruning\\n                will act.\\n\\n    Returns:\\n        module (nn.Module): modified (i.e. pruned) version of the input module\\n\\n    Examples:\\n        >>> # xdoctest: +SKIP\\n        >>> m = prune.identity(nn.Linear(2, 3), 'bias')\\n        >>> print(m.bias_mask)\\n        tensor([1., 1., 1.])\\n    \\\"\\\"\\\"\\n    Identity.apply(module, name)\\n    return module\\n\\n\\ndef random_unstructured(module, name, amount):\\n    r\\\"\\\"\\\"Prune tensor by removing random (currently unpruned) units.\\n\\n    Prunes tensor corresponding to parameter called ``name`` in ``module``\\n    by removing the specified ``amount`` of (currently unpruned) units\\n    selected at random.\\n    Modifies module in place (and also return the modified module) by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune\\n        name (str): parameter name within ``module`` on which pruning\\n                will act.\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n\\n    Returns:\\n        module (nn.Module): modified (i.e. pruned) version of the input module\\n\\n    Examples:\\n        >>> # xdoctest: +SKIP\\n        >>> m = prune.random_unstructured(nn.Linear(2, 3), 'weight', amount=1)\\n        >>> torch.sum(m.weight_mask == 0)\\n        tensor(1)\\n\\n    \\\"\\\"\\\"\\n    RandomUnstructured.apply(module, name, amount)\\n    return module\\n\\n\\ndef l1_unstructured(module, name, amount, importance_scores=None):\\n    r\\\"\\\"\\\"Prune tensor by removing units with the lowest L1-norm.\\n\\n    Prunes tensor corresponding to parameter called ``name`` in ``module``\\n    by removing the specified `amount` of (currently unpruned) units with the\\n    lowest L1-norm.\\n    Modifies module in place (and also return the modified module)\\n    by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune\\n        name (str): parameter name within ``module`` on which pruning\\n                will act.\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n        importance_scores (torch.Tensor): tensor of importance scores (of same\\n            shape as module parameter) used to compute mask for pruning.\\n            The values in this tensor indicate the importance of the corresponding\\n            elements in the parameter being pruned.\\n            If unspecified or None, the module parameter will be used in its place.\\n\\n    Returns:\\n        module (nn.Module): modified (i.e. pruned) version of the input module\\n\\n    Examples:\\n        >>> # xdoctest: +SKIP\\n        >>> m = prune.l1_unstructured(nn.Linear(2, 3), 'weight', amount=0.2)\\n        >>> m.state_dict().keys()\\n        odict_keys(['bias', 'weight_orig', 'weight_mask'])\\n    \\\"\\\"\\\"\\n    L1Unstructured.apply(\\n        module, name, amount=amount, importance_scores=importance_scores\\n    )\\n    return module\\n\\n\\ndef random_structured(module, name, amount, dim):\\n    r\\\"\\\"\\\"Prune tensor by removing random channels along the specified dimension.\\n\\n    Prunes tensor corresponding to parameter called ``name`` in ``module``\\n    by removing the specified ``amount`` of (currently unpruned) channels\\n    along the specified ``dim`` selected at random.\\n    Modifies module in place (and also return the modified module)\\n    by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune\\n        name (str): parameter name within ``module`` on which pruning\\n                will act.\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n        dim (int): index of the dim along which we define channels to prune.\\n\\n    Returns:\\n        module (nn.Module): modified (i.e. pruned) version of the input module\\n\\n    Examples:\\n        >>> # xdoctest: +SKIP\\n        >>> m = prune.random_structured(\\n        ...     nn.Linear(5, 3), 'weight', amount=3, dim=1\\n        ... )\\n        >>> columns_pruned = int(sum(torch.sum(m.weight, dim=0) == 0))\\n        >>> print(columns_pruned)\\n        3\\n    \\\"\\\"\\\"\\n    RandomStructured.apply(module, name, amount, dim)\\n    return module\\n\\n\\ndef ln_structured(module, name, amount, n, dim, importance_scores=None):\\n    r\\\"\\\"\\\"Prune tensor by removing channels with the lowest L\\\\ ``n``-norm along the specified dimension.\\n\\n    Prunes tensor corresponding to parameter called ``name`` in ``module``\\n    by removing the specified ``amount`` of (currently unpruned) channels\\n    along the specified ``dim`` with the lowest L\\\\ ``n``-norm.\\n    Modifies module in place (and also return the modified module)\\n    by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune\\n        name (str): parameter name within ``module`` on which pruning\\n                will act.\\n        amount (int or float): quantity of parameters to prune.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n        n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid\\n            entries for argument ``p`` in :func:`torch.norm`.\\n        dim (int): index of the dim along which we define channels to prune.\\n        importance_scores (torch.Tensor): tensor of importance scores (of same\\n            shape as module parameter) used to compute mask for pruning.\\n            The values in this tensor indicate the importance of the corresponding\\n            elements in the parameter being pruned.\\n            If unspecified or None, the module parameter will be used in its place.\\n\\n    Returns:\\n        module (nn.Module): modified (i.e. pruned) version of the input module\\n\\n    Examples:\\n        >>> from torch.nn.utils import prune\\n        >>> m = prune.ln_structured(\\n        ...     nn.Conv2d(5, 3, 2), 'weight', amount=0.3, dim=1, n=float('-inf')\\n        ... )\\n    \\\"\\\"\\\"\\n    LnStructured.apply(\\n        module, name, amount, n, dim, importance_scores=importance_scores\\n    )\\n    return module\\n\\n\\ndef global_unstructured(parameters, pruning_method, importance_scores=None, **kwargs):\\n    r\\\"\\\"\\\"\\n    Globally prunes tensors corresponding to all parameters in ``parameters`` by applying the specified ``pruning_method``.\\n\\n    Modifies modules in place by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Args:\\n        parameters (Iterable of (module, name) tuples): parameters of\\n            the model to prune in a global fashion, i.e. by aggregating all\\n            weights prior to deciding which ones to prune. module must be of\\n            type :class:`nn.Module`, and name must be a string.\\n        pruning_method (function): a valid pruning function from this module,\\n            or a custom one implemented by the user that satisfies the\\n            implementation guidelines and has ``PRUNING_TYPE='unstructured'``.\\n        importance_scores (dict): a dictionary mapping (module, name) tuples to\\n            the corresponding parameter's importance scores tensor. The tensor\\n            should be the same shape as the parameter, and is used for computing\\n            mask for pruning.\\n            If unspecified or None, the parameter will be used in place of its\\n            importance scores.\\n        kwargs: other keyword arguments such as:\\n            amount (int or float): quantity of parameters to prune across the\\n            specified parameters.\\n            If ``float``, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If ``int``, it represents the\\n            absolute number of parameters to prune.\\n\\n    Raises:\\n        TypeError: if ``PRUNING_TYPE != 'unstructured'``\\n\\n    Note:\\n        Since global structured pruning doesn't make much sense unless the\\n        norm is normalized by the size of the parameter, we now limit the\\n        scope of global pruning to unstructured methods.\\n\\n    Examples:\\n        >>> from torch.nn.utils import prune\\n        >>> from collections import OrderedDict\\n        >>> net = nn.Sequential(OrderedDict([\\n        ...     ('first', nn.Linear(10, 4)),\\n        ...     ('second', nn.Linear(4, 1)),\\n        ... ]))\\n        >>> parameters_to_prune = (\\n        ...     (net.first, 'weight'),\\n        ...     (net.second, 'weight'),\\n        ... )\\n        >>> prune.global_unstructured(\\n        ...     parameters_to_prune,\\n        ...     pruning_method=prune.L1Unstructured,\\n        ...     amount=10,\\n        ... )\\n        >>> print(sum(torch.nn.utils.parameters_to_vector(net.buffers()) == 0))\\n        tensor(10)\\n\\n    \\\"\\\"\\\"\\n    # ensure parameters is a list or generator of tuples\\n    if not isinstance(parameters, Iterable):\\n        raise TypeError(\\\"global_unstructured(): parameters is not an Iterable\\\")\\n\\n    importance_scores = importance_scores if importance_scores is not None else {}\\n    if not isinstance(importance_scores, dict):\\n        raise TypeError(\\\"global_unstructured(): importance_scores must be of type dict\\\")\\n\\n    # flatten importance scores to consider them all at once in global pruning\\n    relevant_importance_scores = torch.nn.utils.parameters_to_vector(\\n        [\\n            importance_scores.get((module, name), getattr(module, name))\\n            for (module, name) in parameters\\n        ]\\n    )\\n    # similarly, flatten the masks (if they exist), or use a flattened vector\\n    # of 1s of the same dimensions as t\\n    default_mask = torch.nn.utils.parameters_to_vector(\\n        [\\n            getattr(module, name + \\\"_mask\\\", torch.ones_like(getattr(module, name)))\\n            for (module, name) in parameters\\n        ]\\n    )\\n\\n    # use the canonical pruning methods to compute the new mask, even if the\\n    # parameter is now a flattened out version of `parameters`\\n    container = PruningContainer()\\n    container._tensor_name = \\\"temp\\\"  # to make it match that of `method`\\n    method = pruning_method(**kwargs)\\n    method._tensor_name = \\\"temp\\\"  # to make it match that of `container`\\n    if method.PRUNING_TYPE != \\\"unstructured\\\":\\n        raise TypeError(\\n            'Only \\\"unstructured\\\" PRUNING_TYPE supported for '\\n            f\\\"the `pruning_method`. Found method {pruning_method} of type {method.PRUNING_TYPE}\\\"\\n        )\\n\\n    container.add_pruning_method(method)\\n\\n    # use the `compute_mask` method from `PruningContainer` to combine the\\n    # mask computed by the new method with the pre-existing mask\\n    final_mask = container.compute_mask(relevant_importance_scores, default_mask)\\n\\n    # Pointer for slicing the mask to match the shape of each parameter\\n    pointer = 0\\n    for module, name in parameters:\\n        param = getattr(module, name)\\n        # The length of the parameter\\n        num_param = param.numel()\\n        # Slice the mask, reshape it\\n        param_mask = final_mask[pointer : pointer + num_param].view_as(param)\\n        # Assign the correct pre-computed mask to each parameter and add it\\n        # to the forward_pre_hooks like any other pruning method\\n        custom_from_mask(module, name, mask=param_mask)\\n\\n        # Increment the pointer to continue slicing the final_mask\\n        pointer += num_param\\n\\n\\ndef custom_from_mask(module, name, mask):\\n    r\\\"\\\"\\\"Prune tensor corresponding to parameter called ``name`` in ``module`` by applying the pre-computed mask in ``mask``.\\n\\n    Modifies module in place (and also return the modified module) by:\\n\\n    1) adding a named buffer called ``name+'_mask'`` corresponding to the\\n       binary mask applied to the parameter ``name`` by the pruning method.\\n    2) replacing the parameter ``name`` by its pruned version, while the\\n       original (unpruned) parameter is stored in a new parameter named\\n       ``name+'_orig'``.\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune\\n        name (str): parameter name within ``module`` on which pruning\\n            will act.\\n        mask (Tensor): binary mask to be applied to the parameter.\\n\\n    Returns:\\n        module (nn.Module): modified (i.e. pruned) version of the input module\\n\\n    Examples:\\n        >>> from torch.nn.utils import prune\\n        >>> m = prune.custom_from_mask(\\n        ...     nn.Linear(5, 3), name='bias', mask=torch.tensor([0, 1, 0])\\n        ... )\\n        >>> print(m.bias_mask)\\n        tensor([0., 1., 0.])\\n\\n    \\\"\\\"\\\"\\n    CustomFromMask.apply(module, name, mask)\\n    return module\\n\\n\\ndef remove(module, name):\\n    r\\\"\\\"\\\"Remove the pruning reparameterization from a module and the pruning method from the forward hook.\\n\\n    The pruned parameter named ``name`` remains permanently pruned, and the parameter\\n    named ``name+'_orig'`` is removed from the parameter list. Similarly,\\n    the buffer named ``name+'_mask'`` is removed from the buffers.\\n\\n    Note:\\n        Pruning itself is NOT undone or reversed!\\n\\n    Args:\\n        module (nn.Module): module containing the tensor to prune\\n        name (str): parameter name within ``module`` on which pruning\\n            will act.\\n\\n    Examples:\\n        >>> m = random_unstructured(nn.Linear(5, 7), name='weight', amount=0.2)\\n        >>> m = remove(m, name='weight')\\n    \\\"\\\"\\\"\\n    for k, hook in module._forward_pre_hooks.items():\\n        if isinstance(hook, BasePruningMethod) and hook._tensor_name == name:\\n            hook.remove(module)\\n            del module._forward_pre_hooks[k]\\n            return module\\n\\n    raise ValueError(\\n        f\\\"Parameter '{name}' of module {module} has to be pruned before pruning can be removed\\\"\\n    )\\n\\n\\ndef is_pruned(module):\\n    r\\\"\\\"\\\"Check if a module is pruned by looking for pruning pre-hooks.\\n\\n    Check whether ``module`` is pruned by looking for\\n    ``forward_pre_hooks`` in its modules that inherit from the\\n    :class:`BasePruningMethod`.\\n\\n    Args:\\n        module (nn.Module): object that is either pruned or unpruned\\n\\n    Returns:\\n        binary answer to whether ``module`` is pruned.\\n\\n    Examples:\\n        >>> from torch.nn.utils import prune\\n        >>> m = nn.Linear(5, 7)\\n        >>> print(prune.is_pruned(m))\\n        False\\n        >>> prune.random_unstructured(m, name='weight', amount=0.2)\\n        >>> print(prune.is_pruned(m))\\n        True\\n    \\\"\\\"\\\"\\n    for _, submodule in module.named_modules():\\n        for hook in submodule._forward_pre_hooks.values():\\n            if isinstance(hook, BasePruningMethod):\\n                return True\\n    return False\\n\\n\\ndef _validate_pruning_amount_init(amount):\\n    r\\\"\\\"\\\"Validate helper to check the range of amount at init.\\n\\n    Args:\\n        amount (int or float): quantity of parameters to prune.\\n            If float, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If int, it represents the\\n            absolute number of parameters to prune.\\n\\n    Raises:\\n        ValueError: if amount is a float not in [0, 1], or if it's a negative\\n            integer.\\n        TypeError: if amount is neither a float nor an integer.\\n\\n    Note:\\n        This does not take into account the number of parameters in the\\n        tensor to be pruned, which is known only at prune.\\n    \\\"\\\"\\\"\\n    if not isinstance(amount, numbers.Real):\\n        raise TypeError(f\\\"Invalid type for amount: {amount}. Must be int or float.\\\")\\n\\n    if (isinstance(amount, numbers.Integral) and amount < 0) or (\\n        not isinstance(amount, numbers.Integral)  # so it's a float\\n        and (float(amount) > 1.0 or float(amount) < 0.0)\\n    ):\\n        raise ValueError(\\n            f\\\"amount={amount} should either be a float in the range [0, 1] or a non-negative integer\\\"\\n        )\\n\\n\\ndef _validate_pruning_amount(amount, tensor_size):\\n    r\\\"\\\"\\\"Validate that the pruning amount is meaningful wrt to the size of the data.\\n\\n    Validation helper to check that the amount of parameters to prune\\n    is meaningful wrt to the size of the data (`tensor_size`).\\n\\n    Args:\\n        amount (int or float): quantity of parameters to prune.\\n            If float, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If int, it represents the\\n            absolute number of parameters to prune.\\n        tensor_size (int): absolute number of parameters in the tensor\\n            to prune.\\n    \\\"\\\"\\\"\\n    # TODO: consider removing this check and allowing users to specify\\n    # a number of units to prune that is greater than the number of units\\n    # left to prune. In this case, the tensor will just be fully pruned.\\n\\n    if isinstance(amount, numbers.Integral) and amount > tensor_size:\\n        raise ValueError(\\n            f\\\"amount={amount} should be smaller than the number of parameters to prune={tensor_size}\\\"\\n        )\\n\\n\\ndef _validate_structured_pruning(t):\\n    r\\\"\\\"\\\"Validate that the tensor to be pruned is at least 2-Dimensional.\\n\\n    Validation helper to check that the tensor to be pruned is multi-\\n    dimensional, such that the concept of \\\"channels\\\" is well-defined.\\n\\n    Args:\\n        t (torch.Tensor): tensor representing the parameter to prune\\n\\n    Raises:\\n        ValueError: if the tensor `t` is not at least 2D.\\n    \\\"\\\"\\\"\\n    shape = t.shape\\n    if len(shape) <= 1:\\n        raise ValueError(\\n            \\\"Structured pruning can only be applied to \\\"\\n            \\\"multidimensional tensors. Found tensor of shape \\\"\\n            f\\\"{shape} with {len(shape)} dims\\\"\\n        )\\n\\n\\ndef _compute_nparams_toprune(amount, tensor_size):\\n    r\\\"\\\"\\\"Convert the pruning amount from a percentage to absolute value.\\n\\n    Since amount can be expressed either in absolute value or as a\\n    percentage of the number of units/channels in a tensor, this utility\\n    function converts the percentage to absolute value to standardize\\n    the handling of pruning.\\n\\n    Args:\\n        amount (int or float): quantity of parameters to prune.\\n            If float, should be between 0.0 and 1.0 and represent the\\n            fraction of parameters to prune. If int, it represents the\\n            absolute number of parameters to prune.\\n        tensor_size (int): absolute number of parameters in the tensor\\n            to prune.\\n\\n    Returns:\\n        int: the number of units to prune in the tensor\\n    \\\"\\\"\\\"\\n    # incorrect type already checked in _validate_pruning_amount_init\\n    if isinstance(amount, numbers.Integral):\\n        return amount\\n    else:\\n        return round(amount * tensor_size)\\n\\n\\ndef _validate_pruning_dim(t, dim):\\n    r\\\"\\\"\\\"Validate that the pruning dimension is within the bounds of the tensor dimension.\\n\\n    Args:\\n        t (torch.Tensor): tensor representing the parameter to prune\\n        dim (int): index of the dim along which we define channels to prune\\n    \\\"\\\"\\\"\\n    if dim >= t.dim():\\n        raise IndexError(f\\\"Invalid index {dim} for tensor of size {t.shape}\\\")\\n\\n\\ndef _compute_norm(t, n, dim):\\n    r\\\"\\\"\\\"Compute the L_n-norm of a tensor along all dimensions except for the specified dimension.\\n\\n    The L_n-norm will be computed across all entries in tensor `t` along all dimension\\n    except for the one identified by dim.\\n    Example: if `t` is of shape, say, 3x2x4 and dim=2 (the last dim),\\n    then norm will have Size [4], and each entry will represent the\\n    `L_n`-norm computed using the 3x2=6 entries for each of the 4 channels.\\n\\n    Args:\\n        t (torch.Tensor): tensor representing the parameter to prune\\n        n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid\\n            entries for argument p in torch.norm\\n        dim (int): dim identifying the channels to prune\\n\\n    Returns:\\n        norm (torch.Tensor): L_n norm computed across all dimensions except\\n            for `dim`. By construction, `norm.shape = t.shape[-1]`.\\n    \\\"\\\"\\\"\\n    # dims = all axes, except for the one identified by `dim`\\n    dims = list(range(t.dim()))\\n    # convert negative indexing\\n    if dim < 0:\\n        dim = dims[dim]\\n    dims.remove(dim)\\n\\n    norm = torch.norm(t, p=n, dim=dims)\\n    return norm\\n\\n\\n# mypy: allow-untyped-defs\\n\\\"\\\"\\\"Spectral Normalization from https://arxiv.org/abs/1802.05957.\\\"\\\"\\\"\\nfrom typing import Any, Optional, TypeVar\\n\\nimport torch\\nimport torch.nn.functional as F\\nfrom torch.nn.modules import Module\\n\\n\\n__all__ = [\\n    \\\"SpectralNorm\\\",\\n    \\\"SpectralNormLoadStateDictPreHook\\\",\\n    \\\"SpectralNormStateDictHook\\\",\\n    \\\"spectral_norm\\\",\\n    \\\"remove_spectral_norm\\\",\\n]\\n\\n\\nclass SpectralNorm:\\n    # Invariant before and after each forward call:\\n    #   u = F.normalize(W @ v)\\n    # NB: At initialization, this invariant is not enforced\\n\\n    _version: int = 1\\n    # At version 1:\\n    #   made  `W` not a buffer,\\n    #   added `v` as a buffer, and\\n    #   made eval mode use `W = u @ W_orig @ v` rather than the stored `W`.\\n    name: str\\n    dim: int\\n    n_power_iterations: int\\n    eps: float\\n\\n    def __init__(\\n        self,\\n        name: str = \\\"weight\\\",\\n        n_power_iterations: int = 1,\\n        dim: int = 0,\\n        eps: float = 1e-12,\\n    ) -> None:\\n        self.name = name\\n        self.dim = dim\\n        if n_power_iterations <= 0:\\n            raise ValueError(\\n                \\\"Expected n_power_iterations to be positive, but \\\"\\n                f\\\"got n_power_iterations={n_power_iterations}\\\"\\n            )\\n        self.n_power_iterations = n_power_iterations\\n        self.eps = eps\\n\\n    def reshape_weight_to_matrix(self, weight: torch.Tensor) -> torch.Tensor:\\n        weight_mat = weight\\n        if self.dim != 0:\\n            # permute dim to front\\n            weight_mat = weight_mat.permute(\\n                self.dim, *[d for d in range(weight_mat.dim()) if d != self.dim]\\n            )\\n        height = weight_mat.size(0)\\n        return weight_mat.reshape(height, -1)\\n\\n    def compute_weight(self, module: Module, do_power_iteration: bool) -> torch.Tensor:\\n        # NB: If `do_power_iteration` is set, the `u` and `v` vectors are\\n        #     updated in power iteration **in-place**. This is very important\\n        #     because in `DataParallel` forward, the vectors (being buffers) are\\n        #     broadcast from the parallelized module to each module replica,\\n        #     which is a new module object created on the fly. And each replica\\n        #     runs its own spectral norm power iteration. So simply assigning\\n        #     the updated vectors to the module this function runs on will cause\\n        #     the update to be lost forever. And the next time the parallelized\\n        #     module is replicated, the same randomly initialized vectors are\\n        #     broadcast and used!\\n        #\\n        #     Therefore, to make the change propagate back, we rely on two\\n        #     important behaviors (also enforced via tests):\\n        #       1. `DataParallel` doesn't clone storage if the broadcast tensor\\n        #          is already on correct device; and it makes sure that the\\n        #          parallelized module is already on `device[0]`.\\n        #       2. If the out tensor in `out=` kwarg has correct shape, it will\\n        #          just fill in the values.\\n        #     Therefore, since the same power iteration is performed on all\\n        #     devices, simply updating the tensors in-place will make sure that\\n        #     the module replica on `device[0]` will update the _u vector on the\\n        #     parallelized module (by shared storage).\\n        #\\n        #    However, after we update `u` and `v` in-place, we need to **clone**\\n        #    them before using them to normalize the weight. This is to support\\n        #    backproping through two forward passes, e.g., the common pattern in\\n        #    GAN training: loss = D(real) - D(fake). Otherwise, engine will\\n        #    complain that variables needed to do backward for the first forward\\n        #    (i.e., the `u` and `v` vectors) are changed in the second forward.\\n        weight = getattr(module, self.name + \\\"_orig\\\")\\n        u = getattr(module, self.name + \\\"_u\\\")\\n        v = getattr(module, self.name + \\\"_v\\\")\\n        weight_mat = self.reshape_weight_to_matrix(weight)\\n\\n        if do_power_iteration:\\n            with torch.no_grad():\\n                for _ in range(self.n_power_iterations):\\n                    # Spectral norm of weight equals to `u^T W v`, where `u` and `v`\\n                    # are the first left and right singular vectors.\\n                    # This power iteration produces approximations of `u` and `v`.\\n                    v = F.normalize(\\n                        torch.mv(weight_mat.t(), u), dim=0, eps=self.eps, out=v\\n                    )\\n                    u = F.normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps, out=u)\\n                if self.n_power_iterations > 0:\\n                    # See above on why we need to clone\\n                    u = u.clone(memory_format=torch.contiguous_format)\\n                    v = v.clone(memory_format=torch.contiguous_format)\\n\\n        sigma = torch.dot(u, torch.mv(weight_mat, v))\\n        weight = weight / sigma\\n        return weight\\n\\n    def remove(self, module: Module) -> None:\\n        with torch.no_grad():\\n            weight = self.compute_weight(module, do_power_iteration=False)\\n        delattr(module, self.name)\\n        delattr(module, self.name + \\\"_u\\\")\\n        delattr(module, self.name + \\\"_v\\\")\\n        delattr(module, self.name + \\\"_orig\\\")\\n        module.register_parameter(self.name, torch.nn.Parameter(weight.detach()))\\n\\n    def __call__(self, module: Module, inputs: Any) -> None:\\n        setattr(\\n            module,\\n            self.name,\\n            self.compute_weight(module, do_power_iteration=module.training),\\n        )\\n\\n    def _solve_v_and_rescale(self, weight_mat, u, target_sigma):\\n        # Tries to returns a vector `v` s.t. `u = F.normalize(W @ v)`\\n        # (the invariant at top of this class) and `u @ W @ v = sigma`.\\n        # This uses pinverse in case W^T W is not invertible.\\n        v = torch.linalg.multi_dot(\\n            [weight_mat.t().mm(weight_mat).pinverse(), weight_mat.t(), u.unsqueeze(1)]\\n        ).squeeze(1)\\n        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))\\n\\n    @staticmethod\\n    def apply(\\n        module: Module, name: str, n_power_iterations: int, dim: int, eps: float\\n    ) -> \\\"SpectralNorm\\\":\\n        for hook in module._forward_pre_hooks.values():\\n            if isinstance(hook, SpectralNorm) and hook.name == name:\\n                raise RuntimeError(\\n                    f\\\"Cannot register two spectral_norm hooks on the same parameter {name}\\\"\\n                )\\n\\n        fn = SpectralNorm(name, n_power_iterations, dim, eps)\\n        weight = module._parameters[name]\\n        if weight is None:\\n            raise ValueError(\\n                f\\\"`SpectralNorm` cannot be applied as parameter `{name}` is None\\\"\\n            )\\n        if isinstance(weight, torch.nn.parameter.UninitializedParameter):\\n            raise ValueError(\\n                \\\"The module passed to `SpectralNorm` can't have uninitialized parameters. \\\"\\n                \\\"Make sure to run the dummy forward before applying spectral normalization\\\"\\n            )\\n\\n        with torch.no_grad():\\n            weight_mat = fn.reshape_weight_to_matrix(weight)\\n\\n            h, w = weight_mat.size()\\n            # randomly initialize `u` and `v`\\n            u = F.normalize(weight.new_empty(h).normal_(0, 1), dim=0, eps=fn.eps)\\n            v = F.normalize(weight.new_empty(w).normal_(0, 1), dim=0, eps=fn.eps)\\n\\n        delattr(module, fn.name)\\n        module.register_parameter(fn.name + \\\"_orig\\\", weight)\\n        # We still need to assign weight back as fn.name because all sorts of\\n        # things may assume that it exists, e.g., when initializing weights.\\n        # However, we can't directly assign as it could be an nn.Parameter and\\n        # gets added as a parameter. Instead, we register weight.data as a plain\\n        # attribute.\\n        setattr(module, fn.name, weight.data)\\n        module.register_buffer(fn.name + \\\"_u\\\", u)\\n        module.register_buffer(fn.name + \\\"_v\\\", v)\\n\\n        module.register_forward_pre_hook(fn)\\n        module._register_state_dict_hook(SpectralNormStateDictHook(fn))\\n        module._register_load_state_dict_pre_hook(SpectralNormLoadStateDictPreHook(fn))\\n        return fn\\n\\n\\n# This is a top level class because Py2 pickle doesn't like inner class nor an\\n# instancemethod.\\nclass SpectralNormLoadStateDictPreHook:\\n    # See docstring of SpectralNorm._version on the changes to spectral_norm.\\n    def __init__(self, fn) -> None:\\n        self.fn = fn\\n\\n    # For state_dict with version None, (assuming that it has gone through at\\n    # least one training forward), we have\\n    #\\n    #    u = F.normalize(W_orig @ v)\\n    #    W = W_orig / sigma, where sigma = u @ W_orig @ v\\n    #\\n    # To compute `v`, we solve `W_orig @ x = u`, and let\\n    #    v = x / (u @ W_orig @ x) * (W / W_orig).\\n    def __call__(\\n        self,\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ) -> None:\\n        fn = self.fn\\n        version = local_metadata.get(\\\"spectral_norm\\\", {}).get(\\n            fn.name + \\\".version\\\", None\\n        )\\n        if version is None or version < 1:\\n            weight_key = prefix + fn.name\\n            if (\\n                version is None\\n                and all(weight_key + s in state_dict for s in (\\\"_orig\\\", \\\"_u\\\", \\\"_v\\\"))\\n                and weight_key not in state_dict\\n            ):\\n                # Detect if it is the updated state dict and just missing metadata.\\n                # This could happen if the users are crafting a state dict themselves,\\n                # so we just pretend that this is the newest.\\n                return\\n            has_missing_keys = False\\n            for suffix in (\\\"_orig\\\", \\\"\\\", \\\"_u\\\"):\\n                key = weight_key + suffix\\n                if key not in state_dict:\\n                    has_missing_keys = True\\n                    if strict:\\n                        missing_keys.append(key)\\n            if has_missing_keys:\\n                return\\n            with torch.no_grad():\\n                weight_orig = state_dict[weight_key + \\\"_orig\\\"]\\n                weight = state_dict.pop(weight_key)\\n                sigma = (weight_orig / weight).mean()\\n                weight_mat = fn.reshape_weight_to_matrix(weight_orig)\\n                u = state_dict[weight_key + \\\"_u\\\"]\\n                v = fn._solve_v_and_rescale(weight_mat, u, sigma)\\n                state_dict[weight_key + \\\"_v\\\"] = v\\n\\n\\n# This is a top level class because Py2 pickle doesn't like inner class nor an\\n# instancemethod.\\nclass SpectralNormStateDictHook:\\n    # See docstring of SpectralNorm._version on the changes to spectral_norm.\\n    def __init__(self, fn) -> None:\\n        self.fn = fn\\n\\n    def __call__(self, module, state_dict, prefix, local_metadata) -> None:\\n        if \\\"spectral_norm\\\" not in local_metadata:\\n            local_metadata[\\\"spectral_norm\\\"] = {}\\n        key = self.fn.name + \\\".version\\\"\\n        if key in local_metadata[\\\"spectral_norm\\\"]:\\n            raise RuntimeError(f\\\"Unexpected key in metadata['spectral_norm']: {key}\\\")\\n        local_metadata[\\\"spectral_norm\\\"][key] = self.fn._version\\n\\n\\nT_module = TypeVar(\\\"T_module\\\", bound=Module)\\n\\n\\ndef spectral_norm(\\n    module: T_module,\\n    name: str = \\\"weight\\\",\\n    n_power_iterations: int = 1,\\n    eps: float = 1e-12,\\n    dim: Optional[int] = None,\\n) -> T_module:\\n    r\\\"\\\"\\\"Apply spectral normalization to a parameter in the given module.\\n\\n    .. math::\\n        \\\\mathbf{W}_{SN} = \\\\dfrac{\\\\mathbf{W}}{\\\\sigma(\\\\mathbf{W})},\\n        \\\\sigma(\\\\mathbf{W}) = \\\\max_{\\\\mathbf{h}: \\\\mathbf{h} \\\\ne 0} \\\\dfrac{\\\\|\\\\mathbf{W} \\\\mathbf{h}\\\\|_2}{\\\\|\\\\mathbf{h}\\\\|_2}\\n\\n    Spectral normalization stabilizes the training of discriminators (critics)\\n    in Generative Adversarial Networks (GANs) by rescaling the weight tensor\\n    with spectral norm :math:`\\\\sigma` of the weight matrix calculated using\\n    power iteration method. If the dimension of the weight tensor is greater\\n    than 2, it is reshaped to 2D in power iteration method to get spectral\\n    norm. This is implemented via a hook that calculates spectral norm and\\n    rescales weight before every :meth:`~Module.forward` call.\\n\\n    See `Spectral Normalization for Generative Adversarial Networks`_ .\\n\\n    .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957\\n\\n    Args:\\n        module (nn.Module): containing module\\n        name (str, optional): name of weight parameter\\n        n_power_iterations (int, optional): number of power iterations to\\n            calculate spectral norm\\n        eps (float, optional): epsilon for numerical stability in\\n            calculating norms\\n        dim (int, optional): dimension corresponding to number of outputs,\\n            the default is ``0``, except for modules that are instances of\\n            ConvTranspose{1,2,3}d, when it is ``1``\\n\\n    Returns:\\n        The original module with the spectral norm hook\\n\\n    .. note::\\n        This function has been reimplemented as\\n        :func:`torch.nn.utils.parametrizations.spectral_norm` using the new\\n        parametrization functionality in\\n        :func:`torch.nn.utils.parametrize.register_parametrization`. Please use\\n        the newer version. This function will be deprecated in a future version\\n        of PyTorch.\\n\\n    Example::\\n\\n        >>> m = spectral_norm(nn.Linear(20, 40))\\n        >>> m\\n        Linear(in_features=20, out_features=40, bias=True)\\n        >>> m.weight_u.size()\\n        torch.Size([40])\\n\\n    \\\"\\\"\\\"\\n    if dim is None:\\n        if isinstance(\\n            module,\\n            (\\n                torch.nn.ConvTranspose1d,\\n                torch.nn.ConvTranspose2d,\\n                torch.nn.ConvTranspose3d,\\n            ),\\n        ):\\n            dim = 1\\n        else:\\n            dim = 0\\n    SpectralNorm.apply(module, name, n_power_iterations, dim, eps)\\n    return module\\n\\n\\ndef remove_spectral_norm(module: T_module, name: str = \\\"weight\\\") -> T_module:\\n    r\\\"\\\"\\\"Remove the spectral normalization reparameterization from a module.\\n\\n    Args:\\n        module (Module): containing module\\n        name (str, optional): name of weight parameter\\n\\n    Example:\\n        >>> m = spectral_norm(nn.Linear(40, 10))\\n        >>> remove_spectral_norm(m)\\n    \\\"\\\"\\\"\\n    for k, hook in module._forward_pre_hooks.items():\\n        if isinstance(hook, SpectralNorm) and hook.name == name:\\n            hook.remove(module)\\n            del module._forward_pre_hooks[k]\\n            break\\n    else:\\n        raise ValueError(f\\\"spectral_norm of '{name}' not found in {module}\\\")\\n\\n    for k, hook in module._state_dict_hooks.items():\\n        if isinstance(hook, SpectralNormStateDictHook) and hook.fn.name == name:\\n            del module._state_dict_hooks[k]\\n            break\\n\\n    for k, hook in module._load_state_dict_pre_hooks.items():\\n        if isinstance(hook, SpectralNormLoadStateDictPreHook) and hook.fn.name == name:\\n            del module._load_state_dict_pre_hooks[k]\\n            break\\n\\n    return module\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Any, Dict, Optional, Set, Tuple, Union\\nfrom typing_extensions import deprecated\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch.nn.utils._named_member_accessor import NamedMemberAccessor\\n\\n\\n__all__ = [\\\"functional_call\\\"]\\n\\n\\ndef _untie_named_tensors_map(\\n    module: \\\"torch.nn.Module\\\",\\n    parameters_and_buffers: Dict[str, Tensor],\\n) -> Dict[str, Tensor]:\\n    \\\"\\\"\\\"\\n    Unties all tied tensors in the module to parameters_and_buffers.\\n\\n    This function returns a new untied_parameters_and_buffers dictionary and leave the original\\n    untied_parameters_and_buffers dictionary unchanged. It adds new (missing) keys for tied tensors\\n    in the module to untied_parameters_and_buffers. The value of the new key is the user-given value\\n    in the original parameters_and_buffers dictionary.\\n\\n    If there are more than one user-given values for the same tied tensor, it will raise an error.\\n\\n    For example, if the module has two tied weights self.foo and self.tied_foo and the user passes\\n    {'foo': foo_value, ...}, this will return {'foo': foo_value, 'tied_foo': foo_value, ...}. If the\\n    user passes {'foo': foo_value, 'tied_foo': tied_foo_value, ...}, it will raise an error. If the\\n    user passes {'foo': foo_value, 'tied_foo': foo_value, ...}, it will not raise an error.\\n\\n    Args:\\n        module (torch.nn.Module): the module to determine which tensors are tied.\\n        parameters_and_buffers (Dict[str, Tensor]): a map of {name: tensor} for reparamaterizing the module.\\n\\n    Returns:\\n        A new untied version of the parameters_and_buffers dictionary.\\n\\n    Raises:\\n        ValueError: if there are more than one user-given values for the same tied tensor.\\n    \\\"\\\"\\\"\\n    # A map of {name: tensor} for all tensors (including tied ones) in the module.\\n    all_named_tensors: Dict[str, Tensor] = {}\\n    all_named_tensors.update(module.named_parameters(remove_duplicate=False))\\n    all_named_tensors.update(module.named_buffers(remove_duplicate=False))\\n\\n    # A map of {tensor: set(all_tied_names)} for all tensor names in the module.\\n    tensor_to_tied_names_map: Dict[Tensor, Set[str]] = {}\\n    for name, tensor in all_named_tensors.items():\\n        if tensor not in tensor_to_tied_names_map:\\n            tensor_to_tied_names_map[tensor] = set()\\n        tensor_to_tied_names_map[tensor].add(name)\\n\\n    # A map of {tied_name: set(all_tied_names)} for all tensor names in the module.\\n    # If a name is not tied, it will not be in this map.\\n    tied_names_map: Dict[str, Set[str]] = {}\\n    for tied_names in tensor_to_tied_names_map.values():\\n        if len(tied_names) > 1:\\n            for tied_name in tied_names:\\n                tied_names_map[tied_name] = tied_names\\n\\n    # Make sure the user didn't pass multiple values for the same tied tensor.\\n    given_names = set(parameters_and_buffers.keys())\\n    # same as given_names.intersection(tied_names_map.keys()) but dynamo can't\\n    # handle that\\n    given_names_for_tied_tensors: set[str] = set()\\n    for name in given_names:\\n        if name in tied_names_map:\\n            given_names_for_tied_tensors.add(name)\\n\\n    for given_name in given_names_for_tied_tensors:\\n        tied_names = tied_names_map[given_name]\\n        if (\\n            # Detect if there are multiple keys present for the same tied tensor.\\n            len(tied_names.intersection(given_names_for_tied_tensors)) > 1\\n            # Only raise an error if the user passed multiple values for the same tied tensor.\\n            # If all given values are the same, don't raise.\\n            and len({parameters_and_buffers[tied_name] for tied_name in tied_names})\\n            != 1\\n        ):\\n            raise ValueError(\\n                f\\\"functional_call got multiple values for keys {sorted(tied_names)}, \\\"\\n                f\\\"which are tied. Consider using tie_weights=False\\\"\\n            )\\n\\n    # Untie the given named tensor map\\n    # Make a copy for not modifying the original dict\\n    untied_parameters_and_buffers = parameters_and_buffers.copy()\\n    for given_name in given_names_for_tied_tensors:\\n        for tied_name in tied_names_map[given_name]:\\n            untied_parameters_and_buffers[tied_name] = parameters_and_buffers[\\n                given_name\\n            ]\\n    return untied_parameters_and_buffers\\n\\n\\nclass _ReparametrizeModule:\\n    def __init__(\\n        self,\\n        module: \\\"torch.nn.Module\\\",\\n        parameters_and_buffers: Dict[str, Tensor],\\n        tie_weights: bool = False,\\n        strict: bool = False,\\n        stack_weights: bool = False,\\n    ):\\n        self.parameters_and_buffers = parameters_and_buffers\\n        self.stack_weights = stack_weights\\n\\n        if tie_weights:\\n            self.untied_parameters_and_buffers = _untie_named_tensors_map(\\n                module, parameters_and_buffers\\n            )\\n        else:\\n            self.untied_parameters_and_buffers = parameters_and_buffers\\n\\n        self.accessor = NamedMemberAccessor(module)\\n        if strict:\\n            missing_keys, unexpected_keys = self.accessor.check_keys(\\n                self.untied_parameters_and_buffers\\n            )\\n            error_msgs = []\\n            if len(unexpected_keys) > 0:\\n                error_msgs.append(\\n                    f\\\"Unexpected key(s): {', '.join(map(repr, unexpected_keys))}.\\\"\\n                )\\n            if len(missing_keys) > 0:\\n                error_msgs.append(\\n                    f\\\"Missing key(s): {', '.join(map(repr, missing_keys))}.\\\"\\n                )\\n            if len(error_msgs) > 0:\\n                raise RuntimeError(\\n                    \\\"Error(s) in reparametrizing for {}:\\\\n\\\\t{}\\\".format(\\n                        module._get_name(), \\\"\\\\n\\\\t\\\".join(error_msgs)\\n                    )\\n                )\\n\\n    def __enter__(self):\\n        self.orig_parameters_and_buffers, _ = self.accessor.swap_tensors_dict(\\n            self.untied_parameters_and_buffers, allow_missing=True\\n        )\\n\\n    def __exit__(self, exception_type, exception_value, traceback):\\n        if self.stack_weights:\\n            # When stacking is enabled, we will restore the weights in LIFO order.\\n            self.orig_parameters_and_buffers = dict(\\n                reversed(self.orig_parameters_and_buffers.items())\\n            )\\n        new_parameters_and_buffers, _ = self.accessor.swap_tensors_dict(\\n            self.orig_parameters_and_buffers, allow_missing=True\\n        )\\n        # Sometimes the module is not completely stateless and has some in-place modifications on\\n        # the _parameters and _buffers dictionaries.\\n        # Write the changed parameters and buffers back to the original dict.\\n        self.parameters_and_buffers.update(\\n            {\\n                k: new_parameters_and_buffers[k]\\n                for k in self.parameters_and_buffers\\n                if k in new_parameters_and_buffers\\n            }\\n        )\\n\\n\\ndef _reparametrize_module(\\n    module: \\\"torch.nn.Module\\\",\\n    parameters_and_buffers: Dict[str, Tensor],\\n    *,\\n    tie_weights: bool = False,\\n    strict: bool = False,\\n    stack_weights: bool = False,\\n) -> _ReparametrizeModule:\\n    return _ReparametrizeModule(\\n        module,\\n        parameters_and_buffers,\\n        tie_weights=tie_weights,\\n        strict=strict,\\n        stack_weights=stack_weights,\\n    )\\n\\n\\n@deprecated(\\n    \\\"`torch.nn.utils.stateless.functional_call` is deprecated as of PyTorch 2.0 \\\"\\n    \\\"and will be removed in a future version of PyTorch. \\\"\\n    \\\"Please use `torch.func.functional_call` instead which is a drop-in replacement.\\\",\\n    category=FutureWarning,\\n)\\ndef functional_call(\\n    module: \\\"torch.nn.Module\\\",\\n    parameters_and_buffers: Dict[str, Tensor],\\n    args: Union[Any, Tuple],\\n    kwargs: Optional[Dict[str, Any]] = None,\\n    *,\\n    tie_weights: bool = True,\\n    strict: bool = False,\\n):\\n    r\\\"\\\"\\\"Perform a functional call on the module by replacing the module parameters and buffers with the provided ones.\\n\\n    .. warning::\\n\\n        This API is deprecated as of PyTorch 2.0 and will be removed in a future\\n        version of PyTorch. Please use :func:`torch.func.functional_call` instead,\\n        which is a drop-in replacement for this API.\\n\\n    .. note:: If the module has active parametrizations, passing a value in the\\n        :attr:`parameters_and_buffers` argument with the name set to the regular parameter\\n        name will completely disable the parametrization.\\n        If you want to apply the parametrization function to the value passed\\n        please set the key as ``{submodule_name}.parametrizations.{parameter_name}.original``.\\n\\n    .. note:: If the module performs in-place operations on parameters/buffers, these will be reflected\\n        in the `parameters_and_buffers` input.\\n\\n        Example::\\n\\n            >>> a = {'foo': torch.zeros(())}\\n            >>> # xdoctest: +SKIP\\n            >>> mod = Foo()  # does self.foo = self.foo + 1\\n            >>> print(mod.foo)  # tensor(0.)\\n            >>> functional_call(mod, a, torch.ones(()))\\n            >>> print(mod.foo)  # tensor(0.)\\n            >>> print(a['foo'])  # tensor(1.)\\n\\n    .. note:: If the module has tied weights, whether or not functional_call respects the tying is determined by the\\n        tie_weights flag.\\n\\n        Example::\\n\\n            >>> a = {'foo': torch.zeros(())}\\n            >>> # xdoctest: +SKIP\\n            >>> mod = Foo()  # has both self.foo and self.foo_tied which are tied. Returns x + self.foo + self.foo_tied\\n            >>> print(mod.foo)  # tensor(1.)\\n            >>> mod(torch.zeros(()))  # tensor(2.)\\n            >>> functional_call(mod, a, torch.zeros(()))  # tensor(0.) since it will change self.foo_tied too\\n            >>> functional_call(mod, a, torch.zeros(()), tie_weights=False)  # tensor(1.)--self.foo_tied is not updated\\n            >>> new_a = {'foo': torch.zeros(()), 'foo_tied': torch.zeros(())}\\n            >>> functional_call(mod, new_a, torch.zeros()) # tensor(0.)\\n\\n    Args:\\n        module (torch.nn.Module): the module to call\\n        parameters_and_buffers (dict of str and Tensor): the parameters that will be used in\\n            the module call.\\n        args (Any or tuple): arguments to be passed to the module call. If not a tuple, considered a single argument.\\n        kwargs (dict): keyword arguments to be passed to the module call\\n        tie_weights (bool, optional): If True, then parameters and buffers tied in the original model will be treated as\\n            tied in the reparamaterized version. Therefore, if True and different values are passed for the tied\\n            parameters and buffers, it will error. If False, it will not respect the originally tied parameters and\\n            buffers unless the values passed for both weights are the same. Default: True.\\n        strict (bool, optional): If True, then the parameters and buffers passed in must match the parameters and\\n            buffers in the original module. Therefore, if True and there are any missing or unexpected keys, it will\\n            error. Default: False.\\n\\n    Returns:\\n        Any: the result of calling ``module``.\\n    \\\"\\\"\\\"\\n    return _functional_call(\\n        module,\\n        parameters_and_buffers,\\n        args,\\n        kwargs,\\n        tie_weights=tie_weights,\\n        strict=strict,\\n    )\\n\\n\\ndef _functional_call(\\n    module: \\\"torch.nn.Module\\\",\\n    parameters_and_buffers: Dict[str, Tensor],\\n    args: Union[Any, Tuple],\\n    kwargs: Optional[Dict[str, Any]] = None,\\n    *,\\n    tie_weights: bool = True,\\n    strict: bool = False,\\n):\\n    # TODO allow kwargs such as unsafe and others for parametrization\\n    if (\\n        torch.jit.is_tracing()\\n        or torch.jit.is_scripting()\\n        or isinstance(\\n            module,\\n            (\\n                torch.jit.RecursiveScriptModule,\\n                torch.jit.ScriptModule,\\n                torch.jit.ScriptFunction,\\n            ),\\n        )\\n    ):\\n        raise RuntimeError(\\\"The stateless API can't be used with Jitted modules\\\")\\n    if isinstance(module, torch.nn.DataParallel):\\n        raise RuntimeError(\\n            \\\"The stateless API can't be used with nn.DataParallel module\\\"\\n        )\\n    if kwargs is None:\\n        kwargs = {}\\n    if not isinstance(args, tuple):\\n        args = (args,)\\n    with _reparametrize_module(\\n        module, parameters_and_buffers, tie_weights=tie_weights, strict=strict\\n    ):\\n        return module(*args, **kwargs)\\n\\n\\n# mypy: allow-untyped-defs\\nimport functools\\n\\nimport torch\\nfrom torch.nn.utils._expanded_weights.expanded_weights_impl import ExpandedWeight\\nfrom torch.utils import _pytree as pytree\\n\\n\\n# dependency on `functional_call` means that this can't be exposed in utils\\n# without creating circular dependency\\ndef call_for_per_sample_grads(\\n    module,\\n    *,\\n    batch_size=None,\\n    loss_reduction=\\\"sum\\\",\\n    batch_first=True,\\n):\\n    r\\\"\\\"\\\"\\n    Return a forward function for a module, populating grad_sample with per sample gradients on backward invocation.\\n\\n    Args:\\n        module: The ``nn.Module`` to get per sample gradients with respect to. All trainable\\n          parameters will compute per sample gradients, located in a ``grad_sample``\\n          field when ``backward`` is invoked\\n        batch_size: The batch size of the input. If None is passed, all tensor arguments in args and kwargs must have\\n          the same batch size, which is the size of the first dimension. Otherwise, it must be passed manually.\\n          Default: None\\n        loss_reduction: Indicates if the loss reduction (for aggregating the gradients) is a sum or a mean operation. If\\n          \\\"mean\\\", per sample gradients will be scaled by the batch size to offset the crossbatch interaction from\\n          running mean across a batch. Must be \\\"mean\\\" or \\\"sum\\\". Default: \\\"sum\\\"\\n        batch_first: Indicates if the batch dimension is the first dimension. If True, the batch dimension is the first\\n          dimension. If False, it's the second dimension. Default: True.\\n\\n    Examples::\\n        >>> # xdoctest: +SKIP\\n        >>> model = nn.Linear(4, 3)\\n        >>> batched_input = torch.randn(5, 4)  # batch size of 5\\n        >>> res = call_for_per_sample_grads(model)(batched_input).sum()\\n        >>> res.backward()\\n        >>> assert model.weight.shape == (3, 4)\\n        >>> assert model.weight.grad_sample.shape == (5, 3, 4)\\n        >>> assert model.weight.grad is None\\n        >>> assert model.bias.shape == (3,)\\n        >>> assert model.bias.grad_sample.shape == (5, 3)\\n        >>> assert model.bias.grad is None\\n\\n    An example using \\\"mean\\\" loss reduction. The grad_sample fields will be scaled by batch_size from what they would be\\n    if we ran the same code with loss_reduction=\\\"sum\\\". This is because the mean at the end will scale all\\n    grad_outputs by 1 / batch_size from cross batch interaction.\\n        >>> model = nn.Linear(4, 3)\\n        >>> batched_input = torch.randn(5, 4)  # batch size of 5\\n        >>> res = call_for_per_sample_grads(model, 5, loss_reduction=\\\"mean\\\")(batched_input).mean()\\n        >>> res.backward()\\n\\n    Note::\\n        Does not work with any `nn.RNN`, including `nn.GRU` or `nn.LSTM`. Please use custom\\n        rewrites that wrap an `nn.Linear` module. See Opacus for an example\\n    \\\"\\\"\\\"\\n\\n    def maybe_build_expanded_weight(og_tensor, batch_size):\\n        if og_tensor.requires_grad:\\n            return ExpandedWeight(og_tensor, batch_size, loss_reduction)\\n        else:\\n            return og_tensor\\n\\n    def compute_batch_size(*args, **kwargs):\\n        args_and_kwargs = pytree.arg_tree_leaves(*args, **kwargs)\\n        batch_size = None\\n        for arg in args_and_kwargs:\\n            if not isinstance(arg, torch.Tensor):\\n                continue\\n\\n            arg_batch_size = arg.shape[0] if batch_first else arg.shape[1]\\n            if batch_size is not None and batch_size != arg_batch_size:\\n                raise RuntimeError(\\n                    \\\"When computing batch size, found at least one input with batch size \\\"\\n                    f\\\"{batch_size} and one with batch size {arg_batch_size}. Please specify it \\\"\\n                    \\\"explicitly using the batch size kwarg in call_for_per_sample_grads\\\"\\n                )\\n            batch_size = arg_batch_size\\n        if batch_size is None:\\n            raise RuntimeError(\\n                \\\"Unable to find a tensor in the passed args and kwargs. They may not be pytree-able \\\"\\n                \\\"and so ExpandedWeights cannot compute the batch size from the inputs. Please specify \\\"\\n                \\\"it explicitly\\\"\\n            )\\n        return batch_size\\n\\n    if loss_reduction not in [\\\"sum\\\", \\\"mean\\\"]:\\n        raise RuntimeError(\\n            f\\\"Expected loss_reduction argument to be sum or mean, got {loss_reduction}\\\"\\n        )\\n\\n    if not isinstance(module, torch.nn.Module):\\n        raise RuntimeError(\\n            f\\\"Module passed must be nn.Module, got {type(module).__name__}\\\"\\n        )\\n    if not (batch_size is None or isinstance(batch_size, int)):\\n        raise RuntimeError(\\n            f\\\"Batch size passed must be None or an integer, got {type(batch_size).__name__}\\\"\\n        )\\n    if batch_size is not None and batch_size < 1:\\n        raise RuntimeError(f\\\"Batch size must be positive, got {batch_size}\\\")\\n    for weight in module.parameters():\\n        if hasattr(weight, \\\"grad_sample\\\") and weight.grad_sample is not None:  # type: ignore[attr-defined]\\n            raise RuntimeError(\\n                \\\"Current Expanded Weights accumulates the gradients, which will be incorrect for multiple \\\"\\n                f\\\"calls without clearing gradients. Please clear out the grad_sample parameter of {weight} or \\\"\\n                \\\"post an issue to pytorch/pytorch to prioritize correct behavior\\\"\\n            )\\n\\n    @functools.wraps(module.forward)\\n    def wrapper(*args, **kwargs):\\n        wrapper_batch_size = batch_size\\n        if wrapper_batch_size is None:\\n            wrapper_batch_size = compute_batch_size(*args, **kwargs)\\n\\n        params = {\\n            name: maybe_build_expanded_weight(value, wrapper_batch_size)\\n            for (name, value) in module.named_parameters()\\n        }\\n        return torch.func.functional_call(module, params, args, kwargs)\\n\\n    return wrapper\\n\\n\\n# mypy: allow-untyped-decorators\\n# mypy: allow-untyped-defs\\nimport functools\\nfrom typing import cast, Dict, Iterable, List, Optional, Tuple, Union\\nfrom typing_extensions import deprecated\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch.utils._foreach_utils import (\\n    _device_has_foreach_support,\\n    _group_tensors_by_device_and_dtype,\\n    _has_foreach_support,\\n)\\n\\n\\n__all__ = [\\\"clip_grad_norm_\\\", \\\"clip_grad_norm\\\", \\\"clip_grad_value_\\\"]\\n\\n\\n_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]\\n\\n\\ndef _no_grad(func):\\n    \\\"\\\"\\\"\\n    This wrapper is needed to avoid a circular import when using @torch.no_grad on the exposed functions\\n    clip_grad_norm_ and clip_grad_value_ themselves.\\n    \\\"\\\"\\\"\\n\\n    def _no_grad_wrapper(*args, **kwargs):\\n        with torch.no_grad():\\n            return func(*args, **kwargs)\\n\\n    functools.update_wrapper(_no_grad_wrapper, func)\\n    return _no_grad_wrapper\\n\\n\\n@_no_grad\\ndef clip_grad_norm_(\\n    parameters: _tensor_or_tensors,\\n    max_norm: float,\\n    norm_type: float = 2.0,\\n    error_if_nonfinite: bool = False,\\n    foreach: Optional[bool] = None,\\n) -> torch.Tensor:\\n    r\\\"\\\"\\\"Clip the gradient norm of an iterable of parameters.\\n\\n    The norm is computed over the norms of the individual gradients of all parameters,\\n    as if the norms of the individual gradients were concatenated into a single vector.\\n    Gradients are modified in-place.\\n\\n    Args:\\n        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a\\n            single Tensor that will have gradients normalized\\n        max_norm (float): max norm of the gradients\\n        norm_type (float): type of the used p-norm. Can be ``'inf'`` for\\n            infinity norm.\\n        error_if_nonfinite (bool): if True, an error is thrown if the total\\n            norm of the gradients from :attr:`parameters` is ``nan``,\\n            ``inf``, or ``-inf``. Default: False (will switch to True in the future)\\n        foreach (bool): use the faster foreach-based implementation.\\n            If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently\\n            fall back to the slow implementation for other device types.\\n            Default: ``None``\\n\\n    Returns:\\n        Total norm of the parameter gradients (viewed as a single vector).\\n    \\\"\\\"\\\"\\n    if isinstance(parameters, torch.Tensor):\\n        parameters = [parameters]\\n    grads = [p.grad for p in parameters if p.grad is not None]\\n    max_norm = float(max_norm)\\n    norm_type = float(norm_type)\\n    if len(grads) == 0:\\n        return torch.tensor(0.0)\\n    first_device = grads[0].device\\n    grouped_grads: Dict[\\n        Tuple[torch.device, torch.dtype], Tuple[List[List[Tensor]], List[int]]\\n    ] = _group_tensors_by_device_and_dtype(\\n        [grads]\\n    )  # type: ignore[assignment]\\n\\n    norms: List[Tensor] = []\\n    for (device, _), ([device_grads], _) in grouped_grads.items():  # type: ignore[assignment]\\n        if (foreach is None and _has_foreach_support(device_grads, device)) or (\\n            foreach and _device_has_foreach_support(device)\\n        ):\\n            norms.extend(torch._foreach_norm(device_grads, norm_type))\\n        elif foreach:\\n            raise RuntimeError(\\n                f\\\"foreach=True was passed, but can't use the foreach API on {device.type} tensors\\\"\\n            )\\n        else:\\n            norms.extend([torch.linalg.vector_norm(g, norm_type) for g in device_grads])\\n\\n    total_norm = torch.linalg.vector_norm(\\n        torch.stack([norm.to(first_device) for norm in norms]), norm_type\\n    )\\n\\n    if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):\\n        raise RuntimeError(\\n            f\\\"The total norm of order {norm_type} for gradients from \\\"\\n            \\\"`parameters` is non-finite, so it cannot be clipped. To disable \\\"\\n            \\\"this error and scale the gradients by the non-finite norm anyway, \\\"\\n            \\\"set `error_if_nonfinite=False`\\\"\\n        )\\n    clip_coef = max_norm / (total_norm + 1e-6)\\n    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so\\n    # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization\\n    # when the gradients do not reside in CPU memory.\\n    clip_coef_clamped = torch.clamp(clip_coef, max=1.0)\\n    for (device, _), ([device_grads], _) in grouped_grads.items():  # type: ignore[assignment]\\n        if (foreach is None and _has_foreach_support(device_grads, device)) or (\\n            foreach and _device_has_foreach_support(device)\\n        ):\\n            torch._foreach_mul_(device_grads, clip_coef_clamped.to(device))\\n        elif foreach:\\n            raise RuntimeError(\\n                f\\\"foreach=True was passed, but can't use the foreach API on {device.type} tensors\\\"\\n            )\\n        else:\\n            clip_coef_clamped_device = clip_coef_clamped.to(device)\\n            for g in device_grads:\\n                g.mul_(clip_coef_clamped_device)\\n\\n    return total_norm\\n\\n\\n@deprecated(\\n    \\\"`torch.nn.utils.clip_grad_norm` is now deprecated \\\"\\n    \\\"in favor of `torch.nn.utils.clip_grad_norm_`.\\\",\\n    category=FutureWarning,\\n)\\ndef clip_grad_norm(\\n    parameters: _tensor_or_tensors,\\n    max_norm: float,\\n    norm_type: float = 2.0,\\n    error_if_nonfinite: bool = False,\\n    foreach: Optional[bool] = None,\\n) -> torch.Tensor:\\n    r\\\"\\\"\\\"Clip the gradient norm of an iterable of parameters.\\n\\n    .. warning::\\n        This method is now deprecated in favor of\\n        :func:`torch.nn.utils.clip_grad_norm_`.\\n    \\\"\\\"\\\"\\n    return clip_grad_norm_(parameters, max_norm, norm_type, error_if_nonfinite, foreach)\\n\\n\\n@_no_grad\\ndef clip_grad_value_(\\n    parameters: _tensor_or_tensors,\\n    clip_value: float,\\n    foreach: Optional[bool] = None,\\n) -> None:\\n    r\\\"\\\"\\\"Clip the gradients of an iterable of parameters at specified value.\\n\\n    Gradients are modified in-place.\\n\\n    Args:\\n        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a\\n            single Tensor that will have gradients normalized\\n        clip_value (float): maximum allowed value of the gradients.\\n            The gradients are clipped in the range\\n            :math:`\\\\left[\\\\text{-clip\\\\_value}, \\\\text{clip\\\\_value}\\\\right]`\\n        foreach (bool): use the faster foreach-based implementation\\n            If ``None``, use the foreach implementation for CUDA and CPU native tensors and\\n            silently fall back to the slow implementation for other device types.\\n            Default: ``None``\\n    \\\"\\\"\\\"\\n    if isinstance(parameters, torch.Tensor):\\n        parameters = [parameters]\\n    clip_value = float(clip_value)\\n\\n    grads = [p.grad for p in parameters if p.grad is not None]\\n    grouped_grads = _group_tensors_by_device_and_dtype([grads])\\n\\n    for (device, _), ([grads], _) in grouped_grads.items():  # type: ignore[assignment]\\n        if (\\n            foreach is None\\n            and _has_foreach_support(cast(List[Tensor], grads), device=device)\\n        ) or (foreach and _device_has_foreach_support(device)):\\n            torch._foreach_clamp_min_(cast(List[Tensor], grads), -clip_value)\\n            torch._foreach_clamp_max_(cast(List[Tensor], grads), clip_value)\\n        elif foreach:\\n            raise RuntimeError(\\n                f\\\"foreach=True was passed, but can't use the foreach API on {device.type} tensors\\\"\\n            )\\n        else:\\n            for grad in grads:\\n                cast(Tensor, grad).clamp_(min=-clip_value, max=clip_value)\\n\\n\\n# mypy: allow-untyped-decorators\\n# mypy: allow-untyped-defs\\nimport collections\\nimport copyreg\\nfrom contextlib import contextmanager\\nfrom copy import deepcopy\\nfrom typing import Dict, Optional, Sequence, Tuple, Union\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch.__future__ import get_swap_module_params_on_conversion\\nfrom torch.nn.modules.container import Module, ModuleDict, ModuleList\\nfrom torch.nn.parameter import Parameter\\nfrom torch.utils._python_dispatch import is_traceable_wrapper_subclass\\n\\n\\n__all__ = [\\n    \\\"cached\\\",\\n    \\\"ParametrizationList\\\",\\n    \\\"register_parametrization\\\",\\n    \\\"is_parametrized\\\",\\n    \\\"remove_parametrizations\\\",\\n    \\\"type_before_parametrizations\\\",\\n    \\\"transfer_parametrizations_and_params\\\",\\n]\\n\\n_cache_enabled = 0\\n_cache: Dict[Tuple[int, str], Optional[Tensor]] = {}\\n\\n\\n@contextmanager\\ndef cached():\\n    r\\\"\\\"\\\"Context manager that enables the caching system within parametrizations registered with :func:`register_parametrization`.\\n\\n    The value of the parametrized objects is computed and cached the first time\\n    they are required when this context manager is active. The cached values are\\n    discarded when leaving the context manager.\\n\\n    This is useful when using a parametrized parameter more than once in the forward pass.\\n    An example of this is when parametrizing the recurrent kernel of an RNN or when\\n    sharing weights.\\n\\n    The simplest way to activate the cache is by wrapping the forward pass of the neural network\\n\\n    .. code-block:: python\\n\\n        import torch.nn.utils.parametrize as P\\n        ...\\n        with P.cached():\\n            output = model(inputs)\\n\\n    in training and evaluation. One may also wrap the parts of the modules that use\\n    several times the parametrized tensors. For example, the loop of an RNN with a\\n    parametrized recurrent kernel:\\n\\n    .. code-block:: python\\n\\n        with P.cached():\\n            for x in xs:\\n                out_rnn = self.rnn_cell(x, out_rnn)\\n    \\\"\\\"\\\"\\n    global _cache\\n    global _cache_enabled\\n    _cache_enabled += 1\\n    try:\\n        yield\\n    finally:\\n        _cache_enabled -= 1\\n        if not _cache_enabled:\\n            _cache = {}\\n\\n\\ndef _register_parameter_or_buffer(module, name, X):\\n    if isinstance(X, Parameter):\\n        module.register_parameter(name, X)\\n    else:\\n        module.register_buffer(name, X)\\n\\n\\ndef _maybe_set(dest: Tensor, src: Tensor) -> None:\\n    should_swap = (\\n        get_swap_module_params_on_conversion() or is_traceable_wrapper_subclass(dest)\\n    )\\n    if should_swap:\\n        if isinstance(dest, Parameter) and not isinstance(src, Parameter):\\n            src = Parameter(src, requires_grad=dest.requires_grad)\\n        torch.utils.swap_tensors(dest, src)\\n    else:\\n        dest.set_(src)  # type: ignore[call-overload]\\n\\n\\nclass ParametrizationList(ModuleList):\\n    r\\\"\\\"\\\"A sequential container that holds and manages the original parameters or buffers of a parametrized :class:`torch.nn.Module`.\\n\\n    It is the type of ``module.parametrizations[tensor_name]`` when ``module[tensor_name]``\\n    has been parametrized with :func:`register_parametrization`.\\n\\n    If the first registered parametrization has a ``right_inverse`` that returns one tensor or\\n    does not have a ``right_inverse`` (in which case we assume that ``right_inverse`` is the identity),\\n    it will hold the tensor under the name ``original``.\\n    If it has a ``right_inverse`` that returns more than one tensor, these will be registered as\\n    ``original0``, ``original1``, ...\\n\\n    .. warning::\\n        This class is used internally by :func:`register_parametrization`. It is documented\\n        here for completeness. It shall not be instantiated by the user.\\n\\n    Args:\\n        modules (sequence): sequence of modules representing the parametrizations\\n        original (Parameter or Tensor): parameter or buffer that is parametrized\\n        unsafe (bool): a boolean flag that denotes whether the parametrization\\n            may change the dtype and shape of the tensor. Default: `False`\\n            Warning: the parametrization is not checked for consistency upon registration.\\n            Enable this flag at your own risk.\\n    \\\"\\\"\\\"\\n\\n    original: Tensor\\n    unsafe: bool\\n\\n    def __init__(\\n        self,\\n        modules: Sequence[Module],\\n        original: Union[Tensor, Parameter],\\n        unsafe: bool = False,\\n    ) -> None:\\n        # We require this because we need to treat differently the first parametrization\\n        # This should never throw, unless this class is used from the outside\\n        if len(modules) == 0:\\n            raise ValueError(\\\"ParametrizationList requires one or more modules.\\\")\\n\\n        super().__init__(modules)\\n        self.unsafe = unsafe\\n\\n        # In plain words:\\n        # module.weight must keep its dtype and shape.\\n        # Furthermore, if there is no right_inverse or the right_inverse returns a tensor,\\n        # this should be of the same dtype as the original tensor\\n        #\\n        # We check that the following invariants hold:\\n        #    X = module.weight\\n        #    Y = param.right_inverse(X)\\n        #    assert isinstance(Y, Tensor) or\\n        #           (isinstance(Y, collections.abc.Sequence) and all(isinstance(t, Tensor) for t in Y))\\n        #    Z = param(Y) if isinstance(Y, Tensor) else param(*Y)\\n        #    # Consistency checks\\n        #    assert X.dtype == Z.dtype and X.shape == Z.shape\\n        #    # If it has one input, this allows to be able to use set_ to be able to\\n        #    # move data to/from the original tensor without changing its id (which is what the\\n        #    # optimizer uses to track parameters)\\n        #    if isinstance(Y, Tensor)\\n        #      assert X.dtype == Y.dtype\\n        # Below we use original = X, new = Y\\n\\n        original_shape = original.shape\\n        original_dtype = original.dtype\\n\\n        # Compute new\\n        with torch.no_grad():\\n            new = original\\n            for module in reversed(self):  # type: ignore[call-overload]\\n                if hasattr(module, \\\"right_inverse\\\"):\\n                    try:\\n                        new = module.right_inverse(new)\\n                    except NotImplementedError:\\n                        pass\\n                # else, or if it throws, we assume that right_inverse is the identity\\n\\n        if not isinstance(new, Tensor) and not isinstance(\\n            new, collections.abc.Sequence\\n        ):\\n            raise ValueError(\\n                \\\"'right_inverse' must return a Tensor or a Sequence of tensors (list, tuple...). \\\"\\n                f\\\"Got {type(new).__name__}\\\"\\n            )\\n\\n        # Set the number of original tensors\\n        self.is_tensor = isinstance(new, Tensor)\\n        self.ntensors = 1 if self.is_tensor else len(new)\\n\\n        # Register the tensor(s)\\n        if self.is_tensor:\\n            if original.dtype != new.dtype:\\n                raise ValueError(\\n                    \\\"When `right_inverse` outputs one tensor, it may not change the dtype.\\\\n\\\"\\n                    f\\\"original.dtype: {original.dtype}\\\\n\\\"\\n                    f\\\"right_inverse(original).dtype: {new.dtype}\\\"\\n                )\\n            # Set the original to original so that the user does not need to re-register the parameter\\n            # manually in the optimiser\\n            with torch.no_grad():\\n                _maybe_set(original, new)\\n            _register_parameter_or_buffer(self, \\\"original\\\", original)\\n        else:\\n            for i, originali in enumerate(new):\\n                if not isinstance(originali, Tensor):\\n                    raise ValueError(\\n                        \\\"'right_inverse' must return a Tensor or a Sequence of tensors \\\"\\n                        \\\"(list, tuple...). \\\"\\n                        f\\\"Got element {i} of the sequence with type {type(originali).__name__}.\\\"\\n                    )\\n\\n                # If the original tensor was a Parameter that required grad, we expect the user to\\n                # add the new parameters to the optimizer after registering the parametrization\\n                # (this is documented)\\n                if isinstance(original, Parameter):\\n                    originali = Parameter(originali, original.requires_grad)\\n                originali.requires_grad_(original.requires_grad)\\n                _register_parameter_or_buffer(self, f\\\"original{i}\\\", originali)\\n\\n        if not self.unsafe:\\n            # Consistency checks:\\n            # Since f : A -> B, right_inverse : B -> A, Z and original should live in B\\n            # Z = forward(right_inverse(original))\\n            Z = self()\\n            if not isinstance(Z, Tensor):\\n                raise ValueError(\\n                    f\\\"A parametrization must return a tensor. Got {type(Z).__name__}.\\\"\\n                )\\n            if Z.dtype != original_dtype:\\n                raise ValueError(\\n                    \\\"Registering a parametrization may not change the dtype of the tensor, unless `unsafe` flag is enabled.\\\\n\\\"\\n                    f\\\"unparametrized dtype: {original_dtype}\\\\n\\\"\\n                    f\\\"parametrized dtype: {Z.dtype}\\\"\\n                )\\n            if Z.shape != original_shape:\\n                raise ValueError(\\n                    \\\"Registering a parametrization may not change the shape of the tensor, unless `unsafe` flag is enabled.\\\\n\\\"\\n                    f\\\"unparametrized shape: {original_shape}\\\\n\\\"\\n                    f\\\"parametrized shape: {Z.shape}\\\"\\n                )\\n\\n    def right_inverse(self, value: Tensor) -> None:\\n        r\\\"\\\"\\\"Call the ``right_inverse`` methods of the parametrizations in the inverse registration order.\\n\\n        Then, it stores the result in ``self.original`` if ``right_inverse`` outputs one tensor\\n        or in ``self.original0``, ``self.original1``, ... if it outputs several.\\n\\n        Args:\\n            value (Tensor): Value to which initialize the module\\n        \\\"\\\"\\\"\\n        # All the exceptions in this function should almost never throw.\\n        # They could throw if, for example, right_inverse function returns a different\\n        # dtype when given a different input, which should most likely be caused by a\\n        # bug in the user's code\\n\\n        with torch.no_grad():\\n            # See https://github.com/pytorch/pytorch/issues/53103\\n            for module in reversed(self):  # type: ignore[call-overload]\\n                if hasattr(module, \\\"right_inverse\\\"):\\n                    value = module.right_inverse(value)\\n                else:\\n                    raise RuntimeError(\\n                        f\\\"parametrization {type(module).__name__} does not implement \\\"\\n                        \\\"right_inverse.\\\"\\n                    )\\n            if self.is_tensor:\\n                # These exceptions should only throw when a right_inverse function does not\\n                # return the same dtype for every input, which should most likely be caused by a bug\\n                if not isinstance(value, Tensor):\\n                    raise ValueError(\\n                        f\\\"`right_inverse` should return a tensor. Got {type(value).__name__}\\\"\\n                    )\\n                if value.dtype != self.original.dtype:\\n                    raise ValueError(\\n                        f\\\"The tensor returned by `right_inverse` has dtype {value.dtype} \\\"\\n                        f\\\"while `original` has dtype {self.original.dtype}\\\"\\n                    )\\n                # We know that the result is going to have the same dtype\\n                _maybe_set(self.original, value)\\n            else:\\n                if not isinstance(value, collections.abc.Sequence):\\n                    raise ValueError(\\n                        \\\"'right_inverse' must return a sequence of tensors. \\\"\\n                        f\\\"Got {type(value).__name__}.\\\"\\n                    )\\n                if len(value) != self.ntensors:\\n                    raise ValueError(\\n                        \\\"'right_inverse' must return a sequence of tensors of length \\\"\\n                        f\\\"{self.ntensors}. Got a sequence of length {len(value)}.\\\"\\n                    )\\n                for i, tensor in enumerate(value):\\n                    original_i = getattr(self, f\\\"original{i}\\\")\\n                    if not isinstance(tensor, Tensor):\\n                        raise ValueError(\\n                            f\\\"`right_inverse` must return a sequence of tensors. \\\"\\n                            f\\\"Got element {i} of type {type(tensor).__name__}\\\"\\n                        )\\n                    if original_i.dtype != tensor.dtype:\\n                        raise ValueError(\\n                            f\\\"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} \\\"\\n                            f\\\"while `original{i}` has dtype {original_i.dtype}\\\"\\n                        )\\n                    _maybe_set(original_i, tensor)\\n\\n    def forward(self) -> Tensor:\\n        if torch.jit.is_scripting():\\n            raise RuntimeError(\\\"Parametrization is not working with scripting.\\\")\\n        # Unpack the originals for the first parametrization\\n        if self.is_tensor:\\n            x = self[0](self.original)\\n        else:\\n            originals = (getattr(self, f\\\"original{i}\\\") for i in range(self.ntensors))\\n            x = self[0](*originals)\\n        # It's not possible to call self[1:] here, so we have to be a bit more cryptic\\n        # Also we want to skip all non-integer keys\\n        curr_idx = 1\\n        while hasattr(self, str(curr_idx)):\\n            x = self[curr_idx](x)\\n            curr_idx += 1\\n        return x\\n\\n\\ndef _inject_new_class(module: Module) -> None:\\n    r\\\"\\\"\\\"Set up a module to be parametrized.\\n\\n    This works by substituting the class of the module by a class\\n    that extends it to be able to inject a property\\n\\n    Args:\\n        module (nn.Module): module into which to inject the property\\n    \\\"\\\"\\\"\\n    cls = module.__class__\\n\\n    def default_deepcopy(self, memo):\\n        # Just emulate a standard deepcopy procedure when __deepcopy__ doesn't exist in the current class.\\n        obj = memo.get(id(self), None)\\n        if obj is not None:\\n            return obj\\n        replica = self.__new__(self.__class__)\\n        memo[id(self)] = replica\\n        replica.__dict__ = deepcopy(self.__dict__, memo)\\n        # Also save all slots if they exist.\\n        slots_to_save = copyreg._slotnames(self.__class__)  # type: ignore[attr-defined]\\n        for slot in slots_to_save:\\n            if hasattr(self, slot):\\n                setattr(replica, slot, deepcopy(getattr(self, slot), memo))\\n        return replica\\n\\n    def getstate(self):\\n        raise RuntimeError(\\n            \\\"Serialization of parametrized modules is only \\\"\\n            \\\"supported through state_dict(). See:\\\\n\\\"\\n            \\\"https://pytorch.org/tutorials/beginner/saving_loading_models.html\\\"\\n            \\\"#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training\\\"\\n        )\\n\\n    dct = {\\\"__getstate__\\\": getstate}\\n    # We don't allow serialization of parametrized modules but should still allow deepcopying.\\n    # Default 'deepcopy' function invokes __deepcopy__ method instead of __getstate__ when it exists.\\n    if not hasattr(cls, \\\"__deepcopy__\\\"):\\n        dct[\\\"__deepcopy__\\\"] = default_deepcopy  # type: ignore[assignment]\\n\\n    param_cls = type(\\n        f\\\"Parametrized{cls.__name__}\\\",\\n        (cls,),\\n        dct,\\n    )\\n\\n    module.__class__ = param_cls\\n\\n\\ndef _inject_property(module: Module, tensor_name: str) -> None:\\n    r\\\"\\\"\\\"Injects a property into module[tensor_name].\\n\\n    It assumes that the class in the module has already been modified from its\\n    original one using _inject_new_class and that the tensor under :attr:`tensor_name`\\n    has already been moved out\\n\\n    Args:\\n        module (nn.Module): module into which to inject the property\\n        tensor_name (str): name of the name of the property to create\\n    \\\"\\\"\\\"\\n    # We check the precondition.\\n    # This should never fire if register_parametrization is correctly implemented\\n    assert not hasattr(module, tensor_name)\\n\\n    @torch.jit.unused\\n    def get_cached_parametrization(parametrization) -> Tensor:\\n        global _cache\\n        key = (id(module), tensor_name)\\n        tensor = _cache.get(key)\\n        if tensor is None:\\n            tensor = parametrization()\\n            _cache[key] = tensor\\n        return tensor\\n\\n    def get_parametrized(self) -> Tensor:\\n        if torch.jit.is_scripting():\\n            raise RuntimeError(\\\"Parametrization is not working with scripting.\\\")\\n        parametrization = self.parametrizations[tensor_name]\\n        if _cache_enabled:\\n            if torch.jit.is_scripting():\\n                # Scripting\\n                raise RuntimeError(\\n                    \\\"Caching is not implemented for scripting. \\\"\\n                    \\\"Either disable caching or avoid scripting.\\\"\\n                )\\n            elif torch._C._get_tracing_state() is not None:\\n                # Tracing\\n                raise RuntimeError(\\n                    \\\"Cannot trace a model while caching parametrizations.\\\"\\n                )\\n            else:\\n                return get_cached_parametrization(parametrization)\\n        else:\\n            # If caching is not active, this function just evaluates the parametrization\\n            return parametrization()\\n\\n    def set_original(self, value: Tensor) -> None:\\n        if torch.jit.is_scripting():\\n            raise RuntimeError(\\\"Parametrization is not working with scripting.\\\")\\n        self.parametrizations[tensor_name].right_inverse(value)\\n\\n    setattr(module.__class__, tensor_name, property(get_parametrized, set_original))\\n\\n\\ndef register_parametrization(\\n    module: Module,\\n    tensor_name: str,\\n    parametrization: Module,\\n    *,\\n    unsafe: bool = False,\\n) -> Module:\\n    r\\\"\\\"\\\"Register a parametrization to a tensor in a module.\\n\\n    Assume that ``tensor_name=\\\"weight\\\"`` for simplicity. When accessing ``module.weight``,\\n    the module will return the parametrized version ``parametrization(module.weight)``.\\n    If the original tensor requires a gradient, the backward pass will differentiate\\n    through :attr:`parametrization`, and the optimizer will update the tensor accordingly.\\n\\n    The first time that a module registers a parametrization, this function will add an attribute\\n    ``parametrizations`` to the module of type :class:`~ParametrizationList`.\\n\\n    The list of parametrizations on the tensor ``weight`` will be accessible under\\n    ``module.parametrizations.weight``.\\n\\n    The original tensor will be accessible under\\n    ``module.parametrizations.weight.original``.\\n\\n    Parametrizations may be concatenated by registering several parametrizations\\n    on the same attribute.\\n\\n    The training mode of a registered parametrization is updated on registration\\n    to match the training mode of the host module\\n\\n    Parametrized parameters and buffers have an inbuilt caching system that can be activated\\n    using the context manager :func:`cached`.\\n\\n    A :attr:`parametrization` may optionally implement a method with signature\\n\\n    .. code-block:: python\\n\\n        def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]]\\n\\n    This method is called on the unparametrized tensor when the first parametrization\\n    is registered to compute the initial value of the original tensor.\\n    If this method is not implemented, the original tensor will be just the unparametrized tensor.\\n\\n    If all the parametrizations registered on a tensor implement `right_inverse` it is possible\\n    to initialize a parametrized tensor by assigning to it, as shown in the example below.\\n\\n    It is possible for the first parametrization to depend on several inputs.\\n    This may be implemented returning a tuple of tensors from ``right_inverse``\\n    (see the example implementation of a ``RankOne`` parametrization below).\\n\\n    In this case, the unconstrained tensors are also located under ``module.parametrizations.weight``\\n    with names ``original0``, ``original1``,...\\n\\n    .. note::\\n\\n        If unsafe=False (default) both the forward and right_inverse methods will be called\\n        once to perform a number of consistency checks.\\n        If unsafe=True, then right_inverse will be called if the tensor is not parametrized,\\n        and nothing will be called otherwise.\\n\\n    .. note::\\n\\n        In most situations, ``right_inverse`` will be a function such that\\n        ``forward(right_inverse(X)) == X`` (see\\n        `right inverse <https://en.wikipedia.org/wiki/Inverse_function#Right_inverses>`_).\\n        Sometimes, when the parametrization is not surjective, it may be reasonable\\n        to relax this.\\n\\n    .. warning::\\n\\n        If a parametrization depends on several inputs, :func:`~register_parametrization`\\n        will register a number of new parameters. If such parametrization is registered\\n        after the optimizer is created, these new parameters will need to be added manually\\n        to the optimizer. See :meth:`torch.Optimizer.add_param_group`.\\n\\n    Args:\\n        module (nn.Module): module on which to register the parametrization\\n        tensor_name (str): name of the parameter or buffer on which to register\\n            the parametrization\\n        parametrization (nn.Module): the parametrization to register\\n    Keyword args:\\n        unsafe (bool): a boolean flag that denotes whether the parametrization\\n            may change the dtype and shape of the tensor. Default: `False`\\n            Warning: the parametrization is not checked for consistency upon registration.\\n            Enable this flag at your own risk.\\n\\n    Raises:\\n        ValueError: if the module does not have a parameter or a buffer named :attr:`tensor_name`\\n\\n    Examples:\\n        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)\\n        >>> import torch\\n        >>> import torch.nn as nn\\n        >>> import torch.nn.utils.parametrize as P\\n        >>>\\n        >>> class Symmetric(nn.Module):\\n        >>>     def forward(self, X):\\n        >>>         return X.triu() + X.triu(1).T  # Return a symmetric matrix\\n        >>>\\n        >>>     def right_inverse(self, A):\\n        >>>         return A.triu()\\n        >>>\\n        >>> m = nn.Linear(5, 5)\\n        >>> P.register_parametrization(m, \\\"weight\\\", Symmetric())\\n        >>> print(torch.allclose(m.weight, m.weight.T))  # m.weight is now symmetric\\n        True\\n        >>> A = torch.rand(5, 5)\\n        >>> A = A + A.T   # A is now symmetric\\n        >>> m.weight = A  # Initialize the weight to be the symmetric matrix A\\n        >>> print(torch.allclose(m.weight, A))\\n        True\\n\\n        >>> class RankOne(nn.Module):\\n        >>>     def forward(self, x, y):\\n        >>>         # Form a rank 1 matrix multiplying two vectors\\n        >>>         return x.unsqueeze(-1) @ y.unsqueeze(-2)\\n        >>>\\n        >>>     def right_inverse(self, Z):\\n        >>>         # Project Z onto the rank 1 matrices\\n        >>>         U, S, Vh = torch.linalg.svd(Z, full_matrices=False)\\n        >>>         # Return rescaled singular vectors\\n        >>>         s0_sqrt = S[0].sqrt().unsqueeze(-1)\\n        >>>         return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt\\n        >>>\\n        >>> linear_rank_one = P.register_parametrization(nn.Linear(4, 4), \\\"weight\\\", RankOne())\\n        >>> print(torch.linalg.matrix_rank(linear_rank_one.weight).item())\\n        1\\n\\n    \\\"\\\"\\\"\\n    parametrization.train(module.training)\\n    if is_parametrized(module, tensor_name):\\n        # Correctness checks.\\n        # If A is the space of tensors with shape and dtype equal to module.weight\\n        # we check that parametrization.forward and parametrization.right_inverse are\\n        # functions from A to A\\n        if not unsafe:\\n            Y = getattr(module, tensor_name)\\n            X = parametrization(Y)\\n            if not isinstance(X, Tensor):\\n                raise ValueError(\\n                    f\\\"A parametrization must return a tensor. Got {type(X).__name__}.\\\"\\n                )\\n            if X.dtype != Y.dtype:\\n                raise ValueError(\\n                    \\\"Registering a parametrization may not change the dtype of the tensor, unless the `unsafe` flag is enabled.\\\\n\\\"\\n                    f\\\"module.{tensor_name}.dtype: {Y.dtype}\\\\n\\\"\\n                    f\\\"parametrization(module.{tensor_name}).dtype: {X.dtype}\\\"\\n                )\\n            if X.shape != Y.shape:\\n                raise ValueError(\\n                    \\\"Registering a parametrization may not change the shape of the tensor, unless the `unsafe` flag is enabled.\\\\n\\\"\\n                    f\\\"module.{tensor_name}.shape: {Y.shape}\\\\n\\\"\\n                    f\\\"parametrization(module.{tensor_name}).shape: {X.shape}\\\"\\n                )\\n            if hasattr(parametrization, \\\"right_inverse\\\"):\\n                try:\\n                    Z = parametrization.right_inverse(X)  # type: ignore[operator]\\n                except NotImplementedError:\\n                    pass\\n                else:\\n                    if not isinstance(Z, Tensor):\\n                        raise ValueError(\\n                            f\\\"parametrization.right_inverse must return a tensor. Got: {type(Z).__name__}\\\"\\n                        )\\n                    if Z.dtype != Y.dtype:\\n                        raise ValueError(\\n                            \\\"The tensor returned by parametrization.right_inverse must have the same dtype \\\"\\n                            f\\\"as module.{tensor_name}, unless the `unsafe` flag is enabled.\\\\n\\\"\\n                            f\\\"module.{tensor_name}.dtype: {Y.dtype}\\\\n\\\"\\n                            f\\\"returned dtype: {Z.dtype}\\\"\\n                        )\\n                    if Z.shape != Y.shape:\\n                        raise ValueError(\\n                            \\\"The tensor returned by parametrization.right_inverse must have the same shape \\\"\\n                            f\\\"as module.{tensor_name}, unless the `unsafe` flag is enabled.\\\\n\\\"\\n                            f\\\"module.{tensor_name}.shape: {Y.shape}\\\\n\\\"\\n                            f\\\"returned shape: {Z.shape}\\\"\\n                        )\\n            # else right_inverse is assumed to be the identity\\n\\n        # add the new parametrization to the parametrization list\\n        assert isinstance(module.parametrizations, ModuleDict)  # Make mypy happy\\n        module.parametrizations[tensor_name].append(parametrization)\\n        # If unsafe was True in previous parametrization, keep it enabled\\n        module.parametrizations[tensor_name].unsafe |= unsafe  # type: ignore[index, union-attr]\\n    elif tensor_name in module._buffers or tensor_name in module._parameters:\\n        # Set the parametrization mechanism\\n        # Fetch the original buffer or parameter\\n        original = getattr(module, tensor_name)\\n        # We create this early to check for possible errors\\n        parametrizations = ParametrizationList(\\n            [parametrization], original, unsafe=unsafe\\n        )\\n        # Delete the previous parameter or buffer\\n        delattr(module, tensor_name)\\n        # If this is the first parametrization registered on the module,\\n        # we prepare the module to inject the property\\n        if not is_parametrized(module):\\n            # Change the class\\n            _inject_new_class(module)\\n            # Inject a ``ModuleDict`` into the instance under module.parametrizations\\n            module.parametrizations = ModuleDict()\\n        # Add a property into the class\\n        _inject_property(module, tensor_name)\\n        # Add a ParametrizationList\\n        assert isinstance(module.parametrizations, ModuleDict)  # Make mypy happy\\n        module.parametrizations[tensor_name] = parametrizations\\n    else:\\n        raise ValueError(\\n            f\\\"Module '{module}' does not have a parameter, a buffer, or a \\\"\\n            f\\\"parametrized element with name '{tensor_name}'\\\"\\n        )\\n    return module\\n\\n\\ndef is_parametrized(module: Module, tensor_name: Optional[str] = None) -> bool:\\n    r\\\"\\\"\\\"Determine if a module has a parametrization.\\n\\n    Args:\\n        module (nn.Module): module to query\\n        tensor_name (str, optional): name of the parameter in the module\\n            Default: ``None``\\n    Returns:\\n        ``True`` if :attr:`module` has a parametrization for the parameter named :attr:`tensor_name`,\\n        or if it has any parametrization when :attr:`tensor_name` is ``None``;\\n        otherwise ``False``\\n    \\\"\\\"\\\"\\n    parametrizations = getattr(module, \\\"parametrizations\\\", None)\\n    if parametrizations is None or not isinstance(parametrizations, ModuleDict):\\n        return False\\n    if tensor_name is None:\\n        # Check that there is at least one parametrized buffer or Parameter\\n        return len(parametrizations) > 0\\n    else:\\n        return tensor_name in parametrizations\\n\\n\\ndef remove_parametrizations(\\n    module: Module,\\n    tensor_name: str,\\n    leave_parametrized: bool = True,\\n) -> Module:\\n    r\\\"\\\"\\\"Remove the parametrizations on a tensor in a module.\\n\\n    - If ``leave_parametrized=True``, ``module[tensor_name]`` will be set to\\n      its current output. In this case, the parametrization shall not change the ``dtype``\\n      of the tensor.\\n    - If ``leave_parametrized=False``, ``module[tensor_name]`` will be set to\\n      the unparametrised tensor in ``module.parametrizations[tensor_name].original``.\\n      This is only possible when the parametrization depends on just one tensor.\\n\\n    Args:\\n        module (nn.Module): module from which remove the parametrization\\n        tensor_name (str): name of the parametrization to be removed\\n        leave_parametrized (bool, optional): leave the attribute :attr:`tensor_name` parametrized.\\n            Default: ``True``\\n\\n    Returns:\\n        Module: module\\n\\n    Raises:\\n        ValueError: if ``module[tensor_name]`` is not parametrized\\n        ValueError: if ``leave_parametrized=False`` and the parametrization depends on several tensors\\n    \\\"\\\"\\\"\\n    if not is_parametrized(module, tensor_name):\\n        raise ValueError(\\n            f\\\"Module {module} does not have a parametrization on {tensor_name}\\\"\\n        )\\n\\n    # Fetch the original tensor\\n    assert isinstance(module.parametrizations, ModuleDict)  # Make mypy happy\\n    parametrizations = module.parametrizations[tensor_name]\\n    if parametrizations.is_tensor:\\n        original = parametrizations.original\\n        if leave_parametrized:\\n            with torch.no_grad():\\n                t = getattr(module, tensor_name)\\n            # We know they have the same dtype because we have checked this when registering the\\n            # parametrizations. As such, we can use set_\\n            # We do this so that the parameter does not to change the id()\\n            # This way the user does not need to update the optimizer\\n            with torch.no_grad():\\n                if type(original) is torch.Tensor:\\n                    _maybe_set(original, t)\\n                else:\\n                    try:\\n                        _maybe_set(original, t)\\n                    except RuntimeError as e:\\n                        # TODO: Fix this for tensor subclasses that are parameters:\\n                        # RuntimeError: set_storage is not allowed on a Tensor created from .data or .detach().\\n                        raise RuntimeError(\\n                            \\\"Calling remove_parametrizations() with leave_parametrized=True \\\"\\n                            \\\"for a parameter that is an instance of a tensor subclass requires \\\"\\n                            \\\"set_() to be implemented correctly for the tensor subclass.\\\"\\n                            \\\"Alternatively, one can opt into the swap_tensors path\\\"\\n                            \\\"Either set leave_parametrized=False or provide a working implementation\\\"\\n                            \\\"for set_() in the tensor subclass or set \\\"\\n                            \\\"torch.__future__.set_swap_module_params_on_conversion(True).\\\"\\n                        ) from e\\n    else:\\n        if leave_parametrized:\\n            # We cannot use no_grad because we need to know whether one or more\\n            # original tensors required grad\\n            t = getattr(module, tensor_name)\\n            # We'll have to trust the user to add it to the optimizer\\n            original = Parameter(t) if t.requires_grad else t\\n        else:\\n            raise ValueError(\\n                \\\"Cannot leave unparametrized (`leave_parametrized=False`) a tensor \\\"\\n                \\\"that is parametrized in terms of a sequence of tensors.\\\"\\n            )\\n\\n    # Delete the property that manages the parametrization\\n    delattr(module.__class__, tensor_name)\\n    # Delete the ParametrizationList\\n    del module.parametrizations[tensor_name]\\n\\n    # Restore the parameter / buffer into the main class\\n    _register_parameter_or_buffer(module, tensor_name, original)\\n\\n    # Roll back the parametrized class if no other buffer or parameter\\n    # is currently parametrized in this class\\n    if not is_parametrized(module):\\n        delattr(module, \\\"parametrizations\\\")\\n        # Restore class\\n        orig_cls = module.__class__.__bases__[0]\\n        module.__class__ = orig_cls\\n    return module\\n\\n\\ndef type_before_parametrizations(module: Module) -> type:\\n    r\\\"\\\"\\\"Return the module type before parametrizations were applied and if not, then it returns the module type.\\n\\n    Args:\\n        module (nn.Module): module to get type of\\n    \\\"\\\"\\\"\\n    if is_parametrized(module):\\n        return module.__class__.__bases__[0]\\n    else:\\n        return type(module)\\n\\n\\ndef transfer_parametrizations_and_params(\\n    from_module: Module,\\n    to_module: Module,\\n    tensor_name: Optional[str] = None,\\n) -> Module:\\n    r\\\"\\\"\\\"Transfer parametrizations and the parameters they parametrize from :attr:`from_module` to :attr:`to_module`.\\n\\n    If :attr:`tensor_name` is specified, only transfers the specified parameter, otherwise\\n    transfers all parametrized parameters. If those parameters do not exist in to_module, it will create them.\\n    Does nothing if from_module is not parametrized.\\n\\n    Args:\\n        from_module (nn.Module): module to transfer from\\n        to_module (nn.Module): module to transfer to\\n        tensor_name (str, optional): parameter to transfer\\n\\n    Returns:\\n        Module: to_module\\n    \\\"\\\"\\\"\\n    if is_parametrized(from_module):\\n        assert isinstance(from_module.parametrizations, ModuleDict)  # for mypy\\n\\n        # get list of all params or the single param to transfer\\n        parameters_to_transfer: Union[list, ModuleDict] = (\\n            from_module.parametrizations if tensor_name is None else [tensor_name]\\n        )\\n\\n        assert hasattr(parameters_to_transfer, \\\"__iter__\\\")  # for mypy\\n        for parameter_name in parameters_to_transfer:\\n            # initialize the to-be-transferred param in to_module if it doesn't exist already\\n            if not hasattr(to_module, parameter_name):\\n                setattr(\\n                    to_module,\\n                    parameter_name,\\n                    Parameter(getattr(from_module, parameter_name)),\\n                )\\n\\n            # apply the params's parametrizations to to_module\\n            for param_func in from_module.parametrizations[parameter_name]:\\n                register_parametrization(to_module, parameter_name, param_func)\\n            assert isinstance(to_module.parametrizations, ModuleDict)  # for mypy\\n\\n            # make values match, original values can be stored in either original or\\n            # original0, original1..., need to check both cases\\n            if hasattr(from_module.parametrizations[parameter_name], \\\"original\\\"):\\n                to_module.parametrizations[\\n                    parameter_name\\n                ].original = from_module.parametrizations[parameter_name].original\\n            else:\\n                num = 0\\n                orig_num = \\\"original\\\" + str(num)\\n                # loop through each original# until all values have been set\\n                while hasattr(from_module.parametrizations[parameter_name], orig_num):\\n                    setattr(\\n                        to_module.parametrizations[parameter_name],\\n                        orig_num,\\n                        getattr(from_module.parametrizations[parameter_name], orig_num),\\n                    )\\n                    num = num + 1\\n                    orig_num = \\\"original\\\" + str(num)\\n\\n    return to_module\\n\\n\\n# mypy: allow-untyped-defs\\nfrom enum import auto, Enum\\nfrom typing import Optional\\n\\nimport torch\\nimport torch.nn.functional as F\\nfrom torch import Tensor\\nfrom torch.nn.modules import Module\\nfrom torch.nn.utils import parametrize\\n\\n\\n__all__ = [\\\"orthogonal\\\", \\\"spectral_norm\\\", \\\"weight_norm\\\"]\\n\\n\\ndef _is_orthogonal(Q, eps=None):\\n    n, k = Q.size(-2), Q.size(-1)\\n    Id = torch.eye(k, dtype=Q.dtype, device=Q.device)\\n    # A reasonable eps, but not too large\\n    eps = 10.0 * n * torch.finfo(Q.dtype).eps\\n    return torch.allclose(Q.mH @ Q, Id, atol=eps)\\n\\n\\ndef _make_orthogonal(A):\\n    \\\"\\\"\\\"Assume that A is a tall matrix.\\n\\n    Compute the Q factor s.t. A = QR (A may be complex) and diag(R) is real and non-negative.\\n    \\\"\\\"\\\"\\n    X, tau = torch.geqrf(A)\\n    Q = torch.linalg.householder_product(X, tau)\\n    # The diagonal of X is the diagonal of R (which is always real) so we normalise by its signs\\n    Q *= X.diagonal(dim1=-2, dim2=-1).sgn().unsqueeze(-2)\\n    return Q\\n\\n\\nclass _OrthMaps(Enum):\\n    matrix_exp = auto()\\n    cayley = auto()\\n    householder = auto()\\n\\n\\nclass _Orthogonal(Module):\\n    base: Tensor\\n\\n    def __init__(\\n        self, weight, orthogonal_map: _OrthMaps, *, use_trivialization=True\\n    ) -> None:\\n        super().__init__()\\n\\n        # Note [Householder complex]\\n        # For complex tensors, it is not possible to compute the tensor `tau` necessary for\\n        # linalg.householder_product from the reflectors.\\n        # To see this, note that the reflectors have a shape like:\\n        # 0 0 0\\n        # * 0 0\\n        # * * 0\\n        # which, for complex matrices, give n(n-1) (real) parameters. Now, you need n^2 parameters\\n        # to parametrize the unitary matrices. Saving tau on its own does not work either, because\\n        # not every combination of `(A, tau)` gives a unitary matrix, meaning that if we optimise\\n        # them as independent tensors we would not maintain the constraint\\n        # An equivalent reasoning holds for rectangular matrices\\n        if weight.is_complex() and orthogonal_map == _OrthMaps.householder:\\n            raise ValueError(\\n                \\\"The householder parametrization does not support complex tensors.\\\"\\n            )\\n\\n        self.shape = weight.shape\\n        self.orthogonal_map = orthogonal_map\\n        if use_trivialization:\\n            self.register_buffer(\\\"base\\\", None)\\n\\n    def forward(self, X: torch.Tensor) -> torch.Tensor:\\n        n, k = X.size(-2), X.size(-1)\\n        transposed = n < k\\n        if transposed:\\n            X = X.mT\\n            n, k = k, n\\n        # Here n > k and X is a tall matrix\\n        if (\\n            self.orthogonal_map == _OrthMaps.matrix_exp\\n            or self.orthogonal_map == _OrthMaps.cayley\\n        ):\\n            # We just need n x k - k(k-1)/2 parameters\\n            X = X.tril()\\n            if n != k:\\n                # Embed into a square matrix\\n                X = torch.cat(\\n                    [X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1\\n                )\\n            A = X - X.mH\\n            # A is skew-symmetric (or skew-hermitian)\\n            if self.orthogonal_map == _OrthMaps.matrix_exp:\\n                Q = torch.matrix_exp(A)\\n            elif self.orthogonal_map == _OrthMaps.cayley:\\n                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}\\n                Id = torch.eye(n, dtype=A.dtype, device=A.device)\\n                Q = torch.linalg.solve(\\n                    torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5)\\n                )\\n            # Q is now orthogonal (or unitary) of size (..., n, n)\\n            if n != k:\\n                Q = Q[..., :k]\\n            # Q is now the size of the X (albeit perhaps transposed)\\n        else:\\n            # X is real here, as we do not support householder with complex numbers\\n            A = X.tril(diagonal=-1)\\n            tau = 2.0 / (1.0 + (A * A).sum(dim=-2))\\n            Q = torch.linalg.householder_product(A, tau)\\n            # The diagonal of X is 1's and -1's\\n            # We do not want to differentiate through this or update the diagonal of X hence the casting\\n            Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2)\\n\\n        if hasattr(self, \\\"base\\\"):\\n            Q = self.base @ Q\\n        if transposed:\\n            Q = Q.mT\\n        return Q  # type: ignore[possibly-undefined]\\n\\n    @torch.autograd.no_grad()\\n    def right_inverse(self, Q: torch.Tensor) -> torch.Tensor:\\n        if Q.shape != self.shape:\\n            raise ValueError(\\n                f\\\"Expected a matrix or batch of matrices of shape {self.shape}. \\\"\\n                f\\\"Got a tensor of shape {Q.shape}.\\\"\\n            )\\n\\n        Q_init = Q\\n        n, k = Q.size(-2), Q.size(-1)\\n        transpose = n < k\\n        if transpose:\\n            Q = Q.mT\\n            n, k = k, n\\n\\n        # We always make sure to always copy Q in every path\\n        if not hasattr(self, \\\"base\\\"):\\n            # Note [right_inverse expm cayley]\\n            # If we do not have use_trivialization=True, we just implement the inverse of the forward\\n            # map for the Householder. To see why, think that for the Cayley map,\\n            # we would need to find the matrix X \\\\in R^{n x k} such that:\\n            # Y = torch.cat([X.tril(), X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1)\\n            # A = Y - Y.mH\\n            # cayley(A)[:, :k]\\n            # gives the original tensor. It is not clear how to do this.\\n            # Perhaps via some algebraic manipulation involving the QR like that of\\n            # Corollary 2.2 in Edelman, Arias and Smith?\\n            if (\\n                self.orthogonal_map == _OrthMaps.cayley\\n                or self.orthogonal_map == _OrthMaps.matrix_exp\\n            ):\\n                raise NotImplementedError(\\n                    \\\"It is not possible to assign to the matrix exponential \\\"\\n                    \\\"or the Cayley parametrizations when use_trivialization=False.\\\"\\n                )\\n\\n            # If parametrization == _OrthMaps.householder, make Q orthogonal via the QR decomposition.\\n            # Here Q is always real because we do not support householder and complex matrices.\\n            # See note [Householder complex]\\n            A, tau = torch.geqrf(Q)\\n            # We want to have a decomposition X = QR with diag(R) > 0, as otherwise we could\\n            # decompose an orthogonal matrix Q as Q = (-Q)@(-Id), which is a valid QR decomposition\\n            # The diagonal of Q is the diagonal of R from the qr decomposition\\n            A.diagonal(dim1=-2, dim2=-1).sign_()\\n            # Equality with zero is ok because LAPACK returns exactly zero when it does not want\\n            # to use a particular reflection\\n            A.diagonal(dim1=-2, dim2=-1)[tau == 0.0] *= -1\\n            return A.mT if transpose else A\\n        else:\\n            if n == k:\\n                # We check whether Q is orthogonal\\n                if not _is_orthogonal(Q):\\n                    Q = _make_orthogonal(Q)\\n                else:  # Is orthogonal\\n                    Q = Q.clone()\\n            else:\\n                # Complete Q into a full n x n orthogonal matrix\\n                N = torch.randn(\\n                    *(Q.size()[:-2] + (n, n - k)), dtype=Q.dtype, device=Q.device\\n                )\\n                Q = torch.cat([Q, N], dim=-1)\\n                Q = _make_orthogonal(Q)\\n            self.base = Q\\n\\n            # It is necessary to return the -Id, as we use the diagonal for the\\n            # Householder parametrization. Using -Id makes:\\n            # householder(torch.zeros(m,n)) == torch.eye(m,n)\\n            # Poor man's version of eye_like\\n            neg_Id = torch.zeros_like(Q_init)\\n            neg_Id.diagonal(dim1=-2, dim2=-1).fill_(-1.0)\\n            return neg_Id\\n\\n\\ndef orthogonal(\\n    module: Module,\\n    name: str = \\\"weight\\\",\\n    orthogonal_map: Optional[str] = None,\\n    *,\\n    use_trivialization: bool = True,\\n) -> Module:\\n    r\\\"\\\"\\\"Apply an orthogonal or unitary parametrization to a matrix or a batch of matrices.\\n\\n    Letting :math:`\\\\mathbb{K}` be :math:`\\\\mathbb{R}` or :math:`\\\\mathbb{C}`, the parametrized\\n    matrix :math:`Q \\\\in \\\\mathbb{K}^{m \\\\times n}` is **orthogonal** as\\n\\n    .. math::\\n\\n        \\\\begin{align*}\\n            Q^{\\\\text{H}}Q &= \\\\mathrm{I}_n \\\\mathrlap{\\\\qquad \\\\text{if }m \\\\geq n}\\\\\\\\\\n            QQ^{\\\\text{H}} &= \\\\mathrm{I}_m \\\\mathrlap{\\\\qquad \\\\text{if }m < n}\\n        \\\\end{align*}\\n\\n    where :math:`Q^{\\\\text{H}}` is the conjugate transpose when :math:`Q` is complex\\n    and the transpose when :math:`Q` is real-valued, and\\n    :math:`\\\\mathrm{I}_n` is the `n`-dimensional identity matrix.\\n    In plain words, :math:`Q` will have orthonormal columns whenever :math:`m \\\\geq n`\\n    and orthonormal rows otherwise.\\n\\n    If the tensor has more than two dimensions, we consider it as a batch of matrices of shape `(..., m, n)`.\\n\\n    The matrix :math:`Q` may be parametrized via three different ``orthogonal_map`` in terms of the original tensor:\\n\\n    - ``\\\"matrix_exp\\\"``/``\\\"cayley\\\"``:\\n      the :func:`~torch.matrix_exp` :math:`Q = \\\\exp(A)` and the `Cayley map`_\\n      :math:`Q = (\\\\mathrm{I}_n + A/2)(\\\\mathrm{I}_n - A/2)^{-1}` are applied to a skew-symmetric\\n      :math:`A` to give an orthogonal matrix.\\n    - ``\\\"householder\\\"``: computes a product of Householder reflectors\\n      (:func:`~torch.linalg.householder_product`).\\n\\n    ``\\\"matrix_exp\\\"``/``\\\"cayley\\\"`` often make the parametrized weight converge faster than\\n    ``\\\"householder\\\"``, but they are slower to compute for very thin or very wide matrices.\\n\\n    If ``use_trivialization=True`` (default), the parametrization implements the \\\"Dynamic Trivialization Framework\\\",\\n    where an extra matrix :math:`B \\\\in \\\\mathbb{K}^{n \\\\times n}` is stored under\\n    ``module.parametrizations.weight[0].base``. This helps the\\n    convergence of the parametrized layer at the expense of some extra memory use.\\n    See `Trivializations for Gradient-Based Optimization on Manifolds`_ .\\n\\n    Initial value of :math:`Q`:\\n    If the original tensor is not parametrized and ``use_trivialization=True`` (default), the initial value\\n    of :math:`Q` is that of the original tensor if it is orthogonal (or unitary in the complex case)\\n    and it is orthogonalized via the QR decomposition otherwise (see :func:`torch.linalg.qr`).\\n    Same happens when it is not parametrized and ``orthogonal_map=\\\"householder\\\"`` even when ``use_trivialization=False``.\\n    Otherwise, the initial value is the result of the composition of all the registered\\n    parametrizations applied to the original tensor.\\n\\n    .. note::\\n        This function is implemented using the parametrization functionality\\n        in :func:`~torch.nn.utils.parametrize.register_parametrization`.\\n\\n\\n    .. _`Cayley map`: https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map\\n    .. _`Trivializations for Gradient-Based Optimization on Manifolds`: https://arxiv.org/abs/1909.09501\\n\\n    Args:\\n        module (nn.Module): module on which to register the parametrization.\\n        name (str, optional): name of the tensor to make orthogonal. Default: ``\\\"weight\\\"``.\\n        orthogonal_map (str, optional): One of the following: ``\\\"matrix_exp\\\"``, ``\\\"cayley\\\"``, ``\\\"householder\\\"``.\\n            Default: ``\\\"matrix_exp\\\"`` if the matrix is square or complex, ``\\\"householder\\\"`` otherwise.\\n        use_trivialization (bool, optional): whether to use the dynamic trivialization framework.\\n            Default: ``True``.\\n\\n    Returns:\\n        The original module with an orthogonal parametrization registered to the specified\\n        weight\\n\\n    Example::\\n\\n        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)\\n        >>> orth_linear = orthogonal(nn.Linear(20, 40))\\n        >>> orth_linear\\n        ParametrizedLinear(\\n        in_features=20, out_features=40, bias=True\\n        (parametrizations): ModuleDict(\\n            (weight): ParametrizationList(\\n            (0): _Orthogonal()\\n            )\\n        )\\n        )\\n        >>> # xdoctest: +IGNORE_WANT\\n        >>> Q = orth_linear.weight\\n        >>> torch.dist(Q.T @ Q, torch.eye(20))\\n        tensor(4.9332e-07)\\n    \\\"\\\"\\\"\\n    weight = getattr(module, name, None)\\n    if not isinstance(weight, Tensor):\\n        raise ValueError(\\n            f\\\"Module '{module}' has no parameter or buffer with name '{name}'\\\"\\n        )\\n\\n    # We could implement this for 1-dim tensors as the maps on the sphere\\n    # but I believe it'd bite more people than it'd help\\n    if weight.ndim < 2:\\n        raise ValueError(\\n            \\\"Expected a matrix or batch of matrices. \\\"\\n            f\\\"Got a tensor of {weight.ndim} dimensions.\\\"\\n        )\\n\\n    if orthogonal_map is None:\\n        orthogonal_map = (\\n            \\\"matrix_exp\\\"\\n            if weight.size(-2) == weight.size(-1) or weight.is_complex()\\n            else \\\"householder\\\"\\n        )\\n\\n    orth_enum = getattr(_OrthMaps, orthogonal_map, None)\\n    if orth_enum is None:\\n        raise ValueError(\\n            'orthogonal_map has to be one of \\\"matrix_exp\\\", \\\"cayley\\\", \\\"householder\\\". '\\n            f\\\"Got: {orthogonal_map}\\\"\\n        )\\n    orth = _Orthogonal(weight, orth_enum, use_trivialization=use_trivialization)\\n    parametrize.register_parametrization(module, name, orth, unsafe=True)\\n    return module\\n\\n\\nclass _WeightNorm(Module):\\n    def __init__(\\n        self,\\n        dim: Optional[int] = 0,\\n    ) -> None:\\n        super().__init__()\\n        if dim is None:\\n            dim = -1\\n        self.dim = dim\\n\\n    def forward(self, weight_g, weight_v):\\n        return torch._weight_norm(weight_v, weight_g, self.dim)\\n\\n    def right_inverse(self, weight):\\n        weight_g = torch.norm_except_dim(weight, 2, self.dim)\\n        weight_v = weight\\n\\n        return weight_g, weight_v\\n\\n\\ndef weight_norm(module: Module, name: str = \\\"weight\\\", dim: int = 0):\\n    r\\\"\\\"\\\"Apply weight normalization to a parameter in the given module.\\n\\n    .. math::\\n         \\\\mathbf{w} = g \\\\dfrac{\\\\mathbf{v}}{\\\\|\\\\mathbf{v}\\\\|}\\n\\n    Weight normalization is a reparameterization that decouples the magnitude\\n    of a weight tensor from its direction. This replaces the parameter specified\\n    by :attr:`name` with two parameters: one specifying the magnitude\\n    and one specifying the direction.\\n\\n    By default, with ``dim=0``, the norm is computed independently per output\\n    channel/plane. To compute a norm over the entire weight tensor, use\\n    ``dim=None``.\\n\\n    See https://arxiv.org/abs/1602.07868\\n\\n    Args:\\n        module (Module): containing module\\n        name (str, optional): name of weight parameter\\n        dim (int, optional): dimension over which to compute the norm\\n\\n    Returns:\\n        The original module with the weight norm hook\\n\\n    Example::\\n\\n        >>> m = weight_norm(nn.Linear(20, 40), name='weight')\\n        >>> m\\n        ParametrizedLinear(\\n          in_features=20, out_features=40, bias=True\\n          (parametrizations): ModuleDict(\\n            (weight): ParametrizationList(\\n              (0): _WeightNorm()\\n            )\\n          )\\n        )\\n        >>> m.parametrizations.weight.original0.size()\\n        torch.Size([40, 1])\\n        >>> m.parametrizations.weight.original1.size()\\n        torch.Size([40, 20])\\n\\n    \\\"\\\"\\\"\\n    _weight_norm = _WeightNorm(dim)\\n    parametrize.register_parametrization(module, name, _weight_norm, unsafe=True)\\n\\n    def _weight_norm_compat_hook(\\n        state_dict,\\n        prefix,\\n        local_metadata,\\n        strict,\\n        missing_keys,\\n        unexpected_keys,\\n        error_msgs,\\n    ):\\n        g_key = f\\\"{prefix}{name}_g\\\"\\n        v_key = f\\\"{prefix}{name}_v\\\"\\n        if g_key in state_dict and v_key in state_dict:\\n            original0 = state_dict.pop(g_key)\\n            original1 = state_dict.pop(v_key)\\n            state_dict[f\\\"{prefix}parametrizations.{name}.original0\\\"] = original0\\n            state_dict[f\\\"{prefix}parametrizations.{name}.original1\\\"] = original1\\n\\n    module._register_load_state_dict_pre_hook(_weight_norm_compat_hook)\\n    return module\\n\\n\\nclass _SpectralNorm(Module):\\n    def __init__(\\n        self,\\n        weight: torch.Tensor,\\n        n_power_iterations: int = 1,\\n        dim: int = 0,\\n        eps: float = 1e-12,\\n    ) -> None:\\n        super().__init__()\\n        ndim = weight.ndim\\n        if dim >= ndim or dim < -ndim:\\n            raise IndexError(\\n                \\\"Dimension out of range (expected to be in range of \\\"\\n                f\\\"[-{ndim}, {ndim - 1}] but got {dim})\\\"\\n            )\\n\\n        if n_power_iterations <= 0:\\n            raise ValueError(\\n                \\\"Expected n_power_iterations to be positive, but \\\"\\n                f\\\"got n_power_iterations={n_power_iterations}\\\"\\n            )\\n        self.dim = dim if dim >= 0 else dim + ndim\\n        self.eps = eps\\n        if ndim > 1:\\n            # For ndim == 1 we do not need to approximate anything (see _SpectralNorm.forward)\\n            self.n_power_iterations = n_power_iterations\\n            weight_mat = self._reshape_weight_to_matrix(weight)\\n            h, w = weight_mat.size()\\n\\n            u = weight_mat.new_empty(h).normal_(0, 1)\\n            v = weight_mat.new_empty(w).normal_(0, 1)\\n            self.register_buffer(\\\"_u\\\", F.normalize(u, dim=0, eps=self.eps))\\n            self.register_buffer(\\\"_v\\\", F.normalize(v, dim=0, eps=self.eps))\\n\\n            # Start with u, v initialized to some reasonable values by performing a number\\n            # of iterations of the power method\\n            self._power_method(weight_mat, 15)\\n\\n    def _reshape_weight_to_matrix(self, weight: torch.Tensor) -> torch.Tensor:\\n        # Precondition\\n        assert weight.ndim > 1\\n\\n        if self.dim != 0:\\n            # permute dim to front\\n            weight = weight.permute(\\n                self.dim, *(d for d in range(weight.dim()) if d != self.dim)\\n            )\\n\\n        return weight.flatten(1)\\n\\n    @torch.autograd.no_grad()\\n    def _power_method(self, weight_mat: torch.Tensor, n_power_iterations: int) -> None:\\n        # See original note at torch/nn/utils/spectral_norm.py\\n        # NB: If `do_power_iteration` is set, the `u` and `v` vectors are\\n        #     updated in power iteration **in-place**. This is very important\\n        #     because in `DataParallel` forward, the vectors (being buffers) are\\n        #     broadcast from the parallelized module to each module replica,\\n        #     which is a new module object created on the fly. And each replica\\n        #     runs its own spectral norm power iteration. So simply assigning\\n        #     the updated vectors to the module this function runs on will cause\\n        #     the update to be lost forever. And the next time the parallelized\\n        #     module is replicated, the same randomly initialized vectors are\\n        #     broadcast and used!\\n        #\\n        #     Therefore, to make the change propagate back, we rely on two\\n        #     important behaviors (also enforced via tests):\\n        #       1. `DataParallel` doesn't clone storage if the broadcast tensor\\n        #          is already on correct device; and it makes sure that the\\n        #          parallelized module is already on `device[0]`.\\n        #       2. If the out tensor in `out=` kwarg has correct shape, it will\\n        #          just fill in the values.\\n        #     Therefore, since the same power iteration is performed on all\\n        #     devices, simply updating the tensors in-place will make sure that\\n        #     the module replica on `device[0]` will update the _u vector on the\\n        #     parallelized module (by shared storage).\\n        #\\n        #    However, after we update `u` and `v` in-place, we need to **clone**\\n        #    them before using them to normalize the weight. This is to support\\n        #    backproping through two forward passes, e.g., the common pattern in\\n        #    GAN training: loss = D(real) - D(fake). Otherwise, engine will\\n        #    complain that variables needed to do backward for the first forward\\n        #    (i.e., the `u` and `v` vectors) are changed in the second forward.\\n\\n        # Precondition\\n        assert weight_mat.ndim > 1\\n\\n        for _ in range(n_power_iterations):\\n            # Spectral norm of weight equals to `u^T W v`, where `u` and `v`\\n            # are the first left and right singular vectors.\\n            # This power iteration produces approximations of `u` and `v`.\\n            self._u = F.normalize(\\n                torch.mv(weight_mat, self._v),  # type: ignore[has-type]\\n                dim=0,\\n                eps=self.eps,\\n                out=self._u,  # type: ignore[has-type]\\n            )\\n            self._v = F.normalize(\\n                torch.mv(weight_mat.H, self._u),  # type: ignore[has-type]\\n                dim=0,\\n                eps=self.eps,\\n                out=self._v,  # type: ignore[has-type]\\n            )\\n\\n    def forward(self, weight: torch.Tensor) -> torch.Tensor:\\n        if weight.ndim == 1:\\n            # Faster and more exact path, no need to approximate anything\\n            return F.normalize(weight, dim=0, eps=self.eps)\\n        else:\\n            weight_mat = self._reshape_weight_to_matrix(weight)\\n            if self.training:\\n                self._power_method(weight_mat, self.n_power_iterations)\\n            # See above on why we need to clone\\n            u = self._u.clone(memory_format=torch.contiguous_format)\\n            v = self._v.clone(memory_format=torch.contiguous_format)\\n            # The proper way of computing this should be through F.bilinear, but\\n            # it seems to have some efficiency issues:\\n            # https://github.com/pytorch/pytorch/issues/58093\\n            sigma = torch.vdot(u, torch.mv(weight_mat, v))\\n            return weight / sigma\\n\\n    def right_inverse(self, value: torch.Tensor) -> torch.Tensor:\\n        # we may want to assert here that the passed value already\\n        # satisfies constraints\\n        return value\\n\\n\\ndef spectral_norm(\\n    module: Module,\\n    name: str = \\\"weight\\\",\\n    n_power_iterations: int = 1,\\n    eps: float = 1e-12,\\n    dim: Optional[int] = None,\\n) -> Module:\\n    r\\\"\\\"\\\"Apply spectral normalization to a parameter in the given module.\\n\\n    .. math::\\n        \\\\mathbf{W}_{SN} = \\\\dfrac{\\\\mathbf{W}}{\\\\sigma(\\\\mathbf{W})},\\n        \\\\sigma(\\\\mathbf{W}) = \\\\max_{\\\\mathbf{h}: \\\\mathbf{h} \\\\ne 0} \\\\dfrac{\\\\|\\\\mathbf{W} \\\\mathbf{h}\\\\|_2}{\\\\|\\\\mathbf{h}\\\\|_2}\\n\\n    When applied on a vector, it simplifies to\\n\\n    .. math::\\n        \\\\mathbf{x}_{SN} = \\\\dfrac{\\\\mathbf{x}}{\\\\|\\\\mathbf{x}\\\\|_2}\\n\\n    Spectral normalization stabilizes the training of discriminators (critics)\\n    in Generative Adversarial Networks (GANs) by reducing the Lipschitz constant\\n    of the model. :math:`\\\\sigma` is approximated performing one iteration of the\\n    `power method`_ every time the weight is accessed. If the dimension of the\\n    weight tensor is greater than 2, it is reshaped to 2D in power iteration\\n    method to get spectral norm.\\n\\n\\n    See `Spectral Normalization for Generative Adversarial Networks`_ .\\n\\n    .. _`power method`: https://en.wikipedia.org/wiki/Power_iteration\\n    .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957\\n\\n    .. note::\\n        This function is implemented using the parametrization functionality\\n        in :func:`~torch.nn.utils.parametrize.register_parametrization`. It is a\\n        reimplementation of :func:`torch.nn.utils.spectral_norm`.\\n\\n    .. note::\\n        When this constraint is registered, the singular vectors associated to the largest\\n        singular value are estimated rather than sampled at random. These are then updated\\n        performing :attr:`n_power_iterations` of the `power method`_ whenever the tensor\\n        is accessed with the module on `training` mode.\\n\\n    .. note::\\n        If the `_SpectralNorm` module, i.e., `module.parametrization.weight[idx]`,\\n        is in training mode on removal, it will perform another power iteration.\\n        If you'd like to avoid this iteration, set the module to eval mode\\n        before its removal.\\n\\n    Args:\\n        module (nn.Module): containing module\\n        name (str, optional): name of weight parameter. Default: ``\\\"weight\\\"``.\\n        n_power_iterations (int, optional): number of power iterations to\\n            calculate spectral norm. Default: ``1``.\\n        eps (float, optional): epsilon for numerical stability in\\n            calculating norms. Default: ``1e-12``.\\n        dim (int, optional): dimension corresponding to number of outputs.\\n            Default: ``0``, except for modules that are instances of\\n            ConvTranspose{1,2,3}d, when it is ``1``\\n\\n    Returns:\\n        The original module with a new parametrization registered to the specified\\n        weight\\n\\n    Example::\\n\\n        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> snm = spectral_norm(nn.Linear(20, 40))\\n        >>> snm\\n        ParametrizedLinear(\\n          in_features=20, out_features=40, bias=True\\n          (parametrizations): ModuleDict(\\n            (weight): ParametrizationList(\\n              (0): _SpectralNorm()\\n            )\\n          )\\n        )\\n        >>> torch.linalg.matrix_norm(snm.weight, 2)\\n        tensor(1.0081, grad_fn=<AmaxBackward0>)\\n    \\\"\\\"\\\"\\n    weight = getattr(module, name, None)\\n    if not isinstance(weight, Tensor):\\n        raise ValueError(\\n            f\\\"Module '{module}' has no parameter or buffer with name '{name}'\\\"\\n        )\\n\\n    if dim is None:\\n        if isinstance(\\n            module,\\n            (\\n                torch.nn.ConvTranspose1d,\\n                torch.nn.ConvTranspose2d,\\n                torch.nn.ConvTranspose3d,\\n            ),\\n        ):\\n            dim = 1\\n        else:\\n            dim = 0\\n    parametrize.register_parametrization(\\n        module, name, _SpectralNorm(weight, n_power_iterations, dim, eps)\\n    )\\n    return module\\n\\n\\nfrom . import parametrizations, rnn, stateless\\nfrom .clip_grad import clip_grad_norm, clip_grad_norm_, clip_grad_value_\\nfrom .convert_parameters import parameters_to_vector, vector_to_parameters\\nfrom .fusion import (\\n    fuse_conv_bn_eval,\\n    fuse_conv_bn_weights,\\n    fuse_linear_bn_eval,\\n    fuse_linear_bn_weights,\\n)\\nfrom .init import skip_init\\nfrom .memory_format import (\\n    convert_conv2d_weight_memory_format,\\n    convert_conv3d_weight_memory_format,\\n)\\nfrom .spectral_norm import remove_spectral_norm, spectral_norm\\nfrom .weight_norm import remove_weight_norm, weight_norm\\n\\n\\n__all__ = [\\n    \\\"clip_grad_norm\\\",\\n    \\\"clip_grad_norm_\\\",\\n    \\\"clip_grad_value_\\\",\\n    \\\"convert_conv2d_weight_memory_format\\\",\\n    \\\"convert_conv3d_weight_memory_format\\\",\\n    \\\"fuse_conv_bn_eval\\\",\\n    \\\"fuse_conv_bn_weights\\\",\\n    \\\"fuse_linear_bn_eval\\\",\\n    \\\"fuse_linear_bn_weights\\\",\\n    \\\"parameters_to_vector\\\",\\n    \\\"parametrizations\\\",\\n    \\\"remove_spectral_norm\\\",\\n    \\\"remove_weight_norm\\\",\\n    \\\"rnn\\\",\\n    \\\"skip_init\\\",\\n    \\\"spectral_norm\\\",\\n    \\\"stateless\\\",\\n    \\\"vector_to_parameters\\\",\\n    \\\"weight_norm\\\",\\n]\\n\\n\\n# mypy: allow-untyped-defs\\nimport importlib\\nimport warnings\\nfrom typing import Callable, List\\n\\n\\n_MESSAGE_TEMPLATE = (\\n    r\\\"Usage of '{old_location}' is deprecated; please use '{new_location}' instead.\\\"\\n)\\n\\n\\ndef lazy_deprecated_import(\\n    all: List[str],\\n    old_module: str,\\n    new_module: str,\\n) -> Callable:\\n    r\\\"\\\"\\\"Import utility to lazily import deprecated packages / modules / functional.\\n\\n    The old_module and new_module are also used in the deprecation warning defined\\n    by the `_MESSAGE_TEMPLATE`.\\n\\n    Args:\\n        all: The list of the functions that are imported. Generally, the module's\\n            __all__ list of the module.\\n        old_module: Old module location\\n        new_module: New module location / Migrated location\\n\\n    Returns:\\n        Callable to assign to the `__getattr__`\\n\\n    Usage:\\n\\n        # In the `torch/nn/quantized/functional.py`\\n        from torch.nn.utils._deprecation_utils import lazy_deprecated_import\\n        _MIGRATED_TO = \\\"torch.ao.nn.quantized.functional\\\"\\n        __getattr__ = lazy_deprecated_import(\\n            all=__all__,\\n            old_module=__name__,\\n            new_module=_MIGRATED_TO)\\n    \\\"\\\"\\\"\\n    warning_message = _MESSAGE_TEMPLATE.format(\\n        old_location=old_module, new_location=new_module\\n    )\\n\\n    def getattr_dunder(name):\\n        if name in all:\\n            # We are using the \\\"RuntimeWarning\\\" to make sure it is not\\n            # ignored by default.\\n            warnings.warn(warning_message, RuntimeWarning)\\n            package = importlib.import_module(new_module)\\n            return getattr(package, name)\\n        raise AttributeError(f\\\"Module {new_module!r} has no attribute {name!r}.\\\")\\n\\n    return getattr_dunder\\n\\n\\nfrom __future__ import annotations\\n\\nimport copy\\nfrom typing import Optional, Tuple, TypeVar\\n\\nimport torch\\n\\n\\n__all__ = [\\n    \\\"fuse_conv_bn_eval\\\",\\n    \\\"fuse_conv_bn_weights\\\",\\n    \\\"fuse_linear_bn_eval\\\",\\n    \\\"fuse_linear_bn_weights\\\",\\n]\\n\\nConvT = TypeVar(\\\"ConvT\\\", bound=\\\"torch.nn.modules.conv._ConvNd\\\")\\nLinearT = TypeVar(\\\"LinearT\\\", bound=\\\"torch.nn.Linear\\\")\\n\\n\\ndef fuse_conv_bn_eval(\\n    conv: ConvT,\\n    bn: torch.nn.modules.batchnorm._BatchNorm,\\n    transpose: bool = False,\\n) -> ConvT:\\n    r\\\"\\\"\\\"Fuse a convolutional module and a BatchNorm module into a single, new convolutional module.\\n\\n    Args:\\n        conv (torch.nn.modules.conv._ConvNd): A convolutional module.\\n        bn (torch.nn.modules.batchnorm._BatchNorm): A BatchNorm module.\\n        transpose (bool, optional): If True, transpose the convolutional weight. Defaults to False.\\n\\n    Returns:\\n        torch.nn.modules.conv._ConvNd: The fused convolutional module.\\n\\n    .. note::\\n        Both ``conv`` and ``bn`` must be in eval mode, and ``bn`` must have its running buffers computed.\\n    \\\"\\\"\\\"\\n    assert not (conv.training or bn.training), \\\"Fusion only for eval!\\\"\\n    fused_conv = copy.deepcopy(conv)\\n\\n    assert bn.running_mean is not None and bn.running_var is not None\\n    fused_conv.weight, fused_conv.bias = fuse_conv_bn_weights(\\n        fused_conv.weight,\\n        fused_conv.bias,\\n        bn.running_mean,\\n        bn.running_var,\\n        bn.eps,\\n        bn.weight,\\n        bn.bias,\\n        transpose,\\n    )\\n\\n    return fused_conv\\n\\n\\ndef fuse_conv_bn_weights(\\n    conv_w: torch.Tensor,\\n    conv_b: Optional[torch.Tensor],\\n    bn_rm: torch.Tensor,\\n    bn_rv: torch.Tensor,\\n    bn_eps: float,\\n    bn_w: Optional[torch.Tensor],\\n    bn_b: Optional[torch.Tensor],\\n    transpose: bool = False,\\n) -> Tuple[torch.nn.Parameter, torch.nn.Parameter]:\\n    r\\\"\\\"\\\"Fuse convolutional module parameters and BatchNorm module parameters into new convolutional module parameters.\\n\\n    Args:\\n        conv_w (torch.Tensor): Convolutional weight.\\n        conv_b (Optional[torch.Tensor]): Convolutional bias.\\n        bn_rm (torch.Tensor): BatchNorm running mean.\\n        bn_rv (torch.Tensor): BatchNorm running variance.\\n        bn_eps (float): BatchNorm epsilon.\\n        bn_w (Optional[torch.Tensor]): BatchNorm weight.\\n        bn_b (Optional[torch.Tensor]): BatchNorm bias.\\n        transpose (bool, optional): If True, transpose the conv weight. Defaults to False.\\n\\n    Returns:\\n        Tuple[torch.nn.Parameter, torch.nn.Parameter]: Fused convolutional weight and bias.\\n    \\\"\\\"\\\"\\n    conv_weight_dtype = conv_w.dtype\\n    conv_bias_dtype = conv_b.dtype if conv_b is not None else conv_weight_dtype\\n    if conv_b is None:\\n        conv_b = torch.zeros_like(bn_rm)\\n    if bn_w is None:\\n        bn_w = torch.ones_like(bn_rm)\\n    if bn_b is None:\\n        bn_b = torch.zeros_like(bn_rm)\\n    bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps)\\n\\n    if transpose:\\n        shape = [1, -1] + [1] * (len(conv_w.shape) - 2)\\n    else:\\n        shape = [-1, 1] + [1] * (len(conv_w.shape) - 2)\\n\\n    fused_conv_w = (conv_w * (bn_w * bn_var_rsqrt).reshape(shape)).to(\\n        dtype=conv_weight_dtype\\n    )\\n    fused_conv_b = ((conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b).to(\\n        dtype=conv_bias_dtype\\n    )\\n\\n    return (\\n        torch.nn.Parameter(fused_conv_w, conv_w.requires_grad),\\n        torch.nn.Parameter(fused_conv_b, conv_b.requires_grad),\\n    )\\n\\n\\ndef fuse_linear_bn_eval(\\n    linear: LinearT,\\n    bn: torch.nn.modules.batchnorm._BatchNorm,\\n) -> LinearT:\\n    r\\\"\\\"\\\"Fuse a linear module and a BatchNorm module into a single, new linear module.\\n\\n    Args:\\n        linear (torch.nn.Linear): A Linear module.\\n        bn (torch.nn.modules.batchnorm._BatchNorm): A BatchNorm module.\\n\\n    Returns:\\n        torch.nn.Linear: The fused linear module.\\n\\n    .. note::\\n        Both ``linear`` and ``bn`` must be in eval mode, and ``bn`` must have its running buffers computed.\\n    \\\"\\\"\\\"\\n    assert not (linear.training or bn.training), \\\"Fusion only for eval!\\\"\\n    fused_linear = copy.deepcopy(linear)\\n\\n    \\\"\\\"\\\"\\n    Linear-BN needs to be fused while preserving the shapes of linear weight/bias.\\n    To preserve the shapes of linear weight/bias, the channel dim of bn needs to be broadcastable with the last dim of linear,\\n    because bn operates over the channel dim, (N, C_in, H, W) while linear operates over the last dim, (*, H_in).\\n    To be broadcastable, the number of features in bn and\\n    the number of output features from linear must satisfy the following condition:\\n    1. they are equal, or\\n    2. the number of features in bn is 1\\n    Otherwise, skip the folding path\\n    \\\"\\\"\\\"\\n    assert (\\n        linear.out_features == bn.num_features or bn.num_features == 1\\n    ), \\\"To fuse, linear.out_features == bn.num_features or bn.num_features == 1\\\"\\n\\n    assert bn.running_mean is not None and bn.running_var is not None\\n    fused_linear.weight, fused_linear.bias = fuse_linear_bn_weights(\\n        fused_linear.weight,\\n        fused_linear.bias,\\n        bn.running_mean,\\n        bn.running_var,\\n        bn.eps,\\n        bn.weight,\\n        bn.bias,\\n    )\\n\\n    return fused_linear\\n\\n\\ndef fuse_linear_bn_weights(\\n    linear_w: torch.Tensor,\\n    linear_b: Optional[torch.Tensor],\\n    bn_rm: torch.Tensor,\\n    bn_rv: torch.Tensor,\\n    bn_eps: float,\\n    bn_w: torch.Tensor,\\n    bn_b: torch.Tensor,\\n) -> Tuple[torch.nn.Parameter, torch.nn.Parameter]:\\n    r\\\"\\\"\\\"Fuse linear module parameters and BatchNorm module parameters into new linear module parameters.\\n\\n    Args:\\n        linear_w (torch.Tensor): Linear weight.\\n        linear_b (Optional[torch.Tensor]): Linear bias.\\n        bn_rm (torch.Tensor): BatchNorm running mean.\\n        bn_rv (torch.Tensor): BatchNorm running variance.\\n        bn_eps (float): BatchNorm epsilon.\\n        bn_w (torch.Tensor): BatchNorm weight.\\n        bn_b (torch.Tensor): BatchNorm bias.\\n\\n    Returns:\\n        Tuple[torch.nn.Parameter, torch.nn.Parameter]: Fused linear weight and bias.\\n    \\\"\\\"\\\"\\n    linear_weight_dtype = linear_w.dtype\\n    linear_bias_dtype = linear_b.dtype if linear_b is not None else linear_weight_dtype\\n    if linear_b is None:\\n        linear_b = torch.zeros_like(bn_rm)\\n    bn_scale = bn_w * torch.rsqrt(bn_rv + bn_eps)\\n\\n    fused_w = linear_w * bn_scale.unsqueeze(-1).to(dtype=linear_weight_dtype)\\n    fused_b = ((linear_b - bn_rm) * bn_scale + bn_b).to(dtype=linear_bias_dtype)\\n\\n    return torch.nn.Parameter(fused_w, linear_w.requires_grad), torch.nn.Parameter(\\n        fused_b, linear_b.requires_grad\\n    )\\n\\n\\n# mypy: allow-untyped-defs\\nimport inspect\\n\\nimport torch\\n\\n\\ndef skip_init(module_cls, *args, **kwargs):\\n    r\\\"\\\"\\\"\\n    Given a module class object and args / kwargs, instantiate the module without initializing parameters / buffers.\\n\\n    This can be useful if initialization is slow or if custom initialization will\\n    be performed, making the default initialization unnecessary. There are some caveats to this, due to\\n    the way this function is implemented:\\n\\n    1. The module must accept a `device` arg in its constructor that is passed to any parameters\\n    or buffers created during construction.\\n\\n    2. The module must not perform any computation on parameters in its constructor except\\n    initialization (i.e. functions from :mod:`torch.nn.init`).\\n\\n    If these conditions are satisfied, the module can be instantiated with parameter / buffer values\\n    uninitialized, as if having been created using :func:`torch.empty`.\\n\\n    Args:\\n        module_cls: Class object; should be a subclass of :class:`torch.nn.Module`\\n        args: args to pass to the module's constructor\\n        kwargs: kwargs to pass to the module's constructor\\n\\n    Returns:\\n        Instantiated module with uninitialized parameters / buffers\\n\\n    Example::\\n\\n        >>> # xdoctest: +IGNORE_WANT(\\\"non-deterministic\\\")\\n        >>> import torch\\n        >>> m = torch.nn.utils.skip_init(torch.nn.Linear, 5, 1)\\n        >>> m.weight\\n        Parameter containing:\\n        tensor([[0.0000e+00, 1.5846e+29, 7.8307e+00, 2.5250e-29, 1.1210e-44]],\\n               requires_grad=True)\\n        >>> m2 = torch.nn.utils.skip_init(torch.nn.Linear, in_features=6, out_features=1)\\n        >>> m2.weight\\n        Parameter containing:\\n        tensor([[-1.4677e+24,  4.5915e-41,  1.4013e-45,  0.0000e+00, -1.4677e+24,\\n                  4.5915e-41]], requires_grad=True)\\n\\n    \\\"\\\"\\\"\\n    if not issubclass(module_cls, torch.nn.Module):\\n        raise RuntimeError(f\\\"Expected a Module; got {module_cls}\\\")\\n    if \\\"device\\\" not in inspect.signature(module_cls).parameters:\\n        raise RuntimeError(\\\"Module must support a 'device' arg to skip initialization\\\")\\n\\n    final_device = kwargs.pop(\\\"device\\\", \\\"cpu\\\")\\n    kwargs[\\\"device\\\"] = \\\"meta\\\"\\n    return module_cls(*args, **kwargs).to_empty(device=final_device)\\n\\n\\n# mypy: allow-untyped-defs\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .conv_utils import (\\n    conv_args_and_kwargs,\\n    conv_backward,\\n    conv_input_for_string_padding,\\n    conv_picker,\\n)\\nfrom .expanded_weights_impl import ExpandedWeight, implements_per_sample_grads\\nfrom .expanded_weights_utils import forward_helper\\n\\n\\n@implements_per_sample_grads(F.conv1d)\\n@implements_per_sample_grads(F.conv2d)\\n@implements_per_sample_grads(F.conv3d)\\nclass ConvPerSampleGrad(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, kwarg_names, conv_fn, *expanded_args_and_kwargs):\\n        expanded_args, expanded_kwargs = conv_args_and_kwargs(\\n            kwarg_names, expanded_args_and_kwargs\\n        )\\n        orig_input = expanded_args[0]\\n        was_same_padding = expanded_kwargs[\\\"padding\\\"] == \\\"same\\\"\\n\\n        if isinstance(expanded_kwargs[\\\"padding\\\"], str):\\n            # if padding is a string, we'll do the necessary padding (slowly) using F.pad\\n            kernel_size = expanded_args[1].shape[2:]\\n            padding, dilation = expanded_kwargs[\\\"padding\\\"], expanded_kwargs[\\\"dilation\\\"]\\n            input = conv_input_for_string_padding(\\n                conv_fn, padding, expanded_args[0], dilation, kernel_size\\n            )\\n            expanded_args = (input, expanded_args[1])\\n            # since we've already done the padding, don't need any more\\n            expanded_kwargs[\\\"padding\\\"] = 0\\n\\n        output = forward_helper(conv_fn, expanded_args, expanded_kwargs)\\n        input, weight = expanded_args\\n        batched_dim_size = conv_picker(conv_fn, 3, 4, 5)\\n        if input.dim() != batched_dim_size:\\n            raise RuntimeError(\\n                f\\\"Expanded Weights only support convolution with batched input, got {conv_fn} with an\\\"\\n                f\\\"unbatched input of dim {input.dim()}, expected input of dim {batched_dim_size}\\\"\\n            )\\n\\n        ctx.conv_fn = conv_fn\\n\\n        ctx.batch_size = orig_input.shape[0]\\n        ctx.input_required_grad = orig_input.requires_grad\\n        ctx.orig_input_shape = orig_input.shape\\n        ctx.was_same_padding = was_same_padding\\n        ctx.stride, ctx.padding = expanded_kwargs[\\\"stride\\\"], expanded_kwargs[\\\"padding\\\"]\\n        ctx.dilation, ctx.groups = (\\n            expanded_kwargs[\\\"dilation\\\"],\\n            expanded_kwargs[\\\"groups\\\"],\\n        )\\n\\n        if isinstance(weight, ExpandedWeight):\\n            ctx.input = input\\n        ctx.weight = weight\\n        ctx.bias = expanded_kwargs[\\\"bias\\\"]\\n\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        return conv_backward(ctx.conv_fn, ctx, grad_output)\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import List, Optional\\n\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .expanded_weights_impl import implements_per_sample_grads\\nfrom .expanded_weights_utils import (\\n    forward_helper,\\n    is_batch_first,\\n    set_grad_sample_if_exists,\\n    unpack_expanded_weight_or_tensor,\\n)\\n\\n\\n@implements_per_sample_grads(F.linear)\\nclass LinearPerSampleGrad(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, _, __, *expanded_args_and_kwargs):\\n        if len(expanded_args_and_kwargs[0].shape) <= 1:\\n            raise RuntimeError(\\n                \\\"Input does not have a batch dimension. Expanded Weights expected input \\\"\\n                f\\\"of at least rank 2, got of rank {len(expanded_args_and_kwargs[0].shape)}\\\"\\n            )\\n        expanded_kwargs = {\\n            \\\"bias\\\": expanded_args_and_kwargs[2]\\n            if len(expanded_args_and_kwargs) == 3\\n            else None\\n        }\\n        expanded_args = expanded_args_and_kwargs[:2]\\n        ctx.batch_first = is_batch_first(expanded_args_and_kwargs)\\n        output = forward_helper(F.linear, expanded_args, expanded_kwargs)\\n        ctx.args = expanded_args\\n        ctx.kwargs = expanded_kwargs\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        input, weight = ctx.args\\n        bias = ctx.kwargs[\\\"bias\\\"]\\n        results: List[Optional[torch.Tensor]] = []\\n        results.append(None)  # for kwarg_names\\n        results.append(None)  # for op reference\\n\\n        if input.requires_grad:\\n            results.append(grad_output.matmul(unpack_expanded_weight_or_tensor(weight)))\\n        else:\\n            results.append(None)\\n        results.extend([None] * 2)  # weight and bias don't compute batched gradients\\n\\n        if not ctx.batch_first:\\n            grad_output = grad_output.transpose(0, 1)\\n            input = input.transpose(0, 1)\\n\\n        # weight and bias get their grad_sample fields set directly if they exist\\n        set_grad_sample_if_exists(\\n            weight, lambda _: torch.einsum(\\\"n...i,n...j->nij\\\", grad_output, input)\\n        )\\n        set_grad_sample_if_exists(\\n            bias, lambda _: torch.einsum(\\\"n...k->nk\\\", grad_output)\\n        )\\n        return tuple(results)\\n\\n\\n# mypy: allow-untyped-defs\\nfrom functools import partial\\nfrom typing import List, Optional\\n\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .expanded_weights_impl import implements_per_sample_grads\\nfrom .expanded_weights_utils import (\\n    forward_helper,\\n    set_grad_sample_if_exists,\\n    standard_kwargs,\\n    unpack_expanded_weight_or_tensor,\\n)\\n\\n\\n@implements_per_sample_grads(F.instance_norm)\\nclass InstanceNormPerSampleGrad(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs):\\n        instance_norm = partial(torch.instance_norm, cudnn_enabled=True)\\n        expanded_args, expanded_kwargs = standard_kwargs(\\n            kwarg_names, expanded_args_and_kwargs\\n        )\\n        output = forward_helper(instance_norm, expanded_args, expanded_kwargs)\\n        ctx.input = expanded_args[0]\\n        ctx.running_mean, ctx.running_var = (\\n            expanded_kwargs[\\\"running_mean\\\"],\\n            expanded_kwargs[\\\"running_var\\\"],\\n        )\\n        ctx.weight, ctx.bias, ctx.eps = (\\n            expanded_kwargs[\\\"weight\\\"],\\n            expanded_kwargs[\\\"bias\\\"],\\n            expanded_kwargs[\\\"eps\\\"],\\n        )\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        input, running_mean, running_var = ctx.input, ctx.running_mean, ctx.running_var\\n        weight, bias, eps = ctx.weight, ctx.bias, ctx.eps\\n\\n        results: List[Optional[torch.Tensor]] = []\\n        results.append(None)  # for kwarg names\\n        results.append(None)  # for op reference\\n        if input.requires_grad:\\n            b = input.shape[0]\\n            c = input.shape[1]\\n            new_shape = (1, b * c, *input.shape[2:])\\n\\n            weight_ = unpack_expanded_weight_or_tensor(\\n                weight, lambda orig_weight: orig_weight.repeat(b)\\n            )\\n            running_mean_ = running_mean.repeat(b) if running_mean is not None else None\\n            running_var_ = running_var.repeat(b) if running_var is not None else None\\n            input_reshaped = input.contiguous().view(new_shape)\\n            grad_output_reshaped = grad_output.contiguous().view(new_shape)\\n            mean = torch.mean(\\n                input_reshaped, (0,) + tuple(range(2, input.dim())), False\\n            )\\n            var = torch.var(\\n                input_reshaped,\\n                (0,) + tuple(range(2, input.dim())),\\n                keepdim=False,\\n                unbiased=False,\\n            )\\n            rstd = 1 / torch.sqrt(var + eps)\\n\\n            # must use native batch norm since it supports all inputs. This may have used cuda or openmi during the forward but\\n            # it didn't save the metadata, so we don't know during the backward\\n            res = torch.ops.aten.native_batch_norm_backward(\\n                grad_output_reshaped,\\n                input_reshaped,\\n                weight_,\\n                running_mean_,\\n                running_var_,\\n                mean,\\n                rstd,\\n                True,\\n                eps,\\n                (True, False, False),\\n            )\\n            results.append(res[0].reshape(input.shape))\\n        else:\\n            results.append(None)\\n\\n        # weight and bias don't compute batched gradients; no other arguments are differentiable (2 are not saved from the forward)\\n        results = results + [None] * 7\\n\\n        # set grad_sample field for weight and bias with per sample gradients\\n        set_grad_sample_if_exists(\\n            weight,\\n            lambda _: torch.einsum(\\n                \\\"ni...->ni\\\", F.instance_norm(input, eps=eps) * grad_output\\n            ),\\n        )\\n        set_grad_sample_if_exists(\\n            bias, lambda _: torch.einsum(\\\"ni...->ni\\\", grad_output)\\n        )\\n        return tuple(results)\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import List, Optional\\n\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .expanded_weights_impl import ExpandedWeight, implements_per_sample_grads\\nfrom .expanded_weights_utils import (\\n    forward_helper,\\n    set_grad_sample_if_exists,\\n    standard_kwargs,\\n    sum_over_all_but_batch_and_last_n,\\n    unpack_expanded_weight_or_tensor,\\n)\\n\\n\\n@implements_per_sample_grads(F.layer_norm)\\nclass LayerNormPerSampleGrad(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs):\\n        expanded_args, expanded_kwargs = standard_kwargs(\\n            kwarg_names, expanded_args_and_kwargs\\n        )\\n        input = expanded_args[0]\\n        normalized_shape = expanded_args[1]\\n        if len(input.shape) <= len(normalized_shape):\\n            raise RuntimeError(\\n                \\\"Expanded Weights: Layer norm should not normalize over batch dimension for per sample gradient\\\"\\n                f\\\"computations but got that normalized shape, {normalized_shape}, matched input shape.\\\"\\n            )\\n        output, mean, rstd = forward_helper(\\n            torch.native_layer_norm, expanded_args, expanded_kwargs\\n        )\\n        ctx.args = expanded_args\\n\\n        if input.requires_grad or isinstance(expanded_kwargs[\\\"weight\\\"], ExpandedWeight):\\n            ctx.weight = expanded_kwargs[\\\"weight\\\"]\\n        if input.requires_grad or isinstance(expanded_kwargs[\\\"bias\\\"], ExpandedWeight):\\n            ctx.bias = expanded_kwargs[\\\"bias\\\"]\\n        ctx.eps = expanded_kwargs[\\\"eps\\\"]\\n        ctx.mean, ctx.rstd = mean, rstd\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        def weight_per_sample_grad(weight):\\n            return sum_over_all_but_batch_and_last_n(\\n                F.layer_norm(input, normalized_shape, eps=ctx.eps) * grad_output,\\n                weight.dim(),\\n            )\\n\\n        input, normalized_shape = ctx.args\\n        mean, rstd = ctx.mean, ctx.rstd\\n\\n        results: List[Optional[torch.Tensor]] = []\\n        results.append(None)  # for kwarg names\\n        results.append(None)  # for op reference\\n        if input.requires_grad:\\n            weight_ = unpack_expanded_weight_or_tensor(ctx.weight)\\n            bias_ = unpack_expanded_weight_or_tensor(ctx.bias)\\n            results.append(\\n                torch.ops.aten.native_layer_norm_backward(\\n                    grad_output,\\n                    input,\\n                    normalized_shape,\\n                    mean,\\n                    rstd,\\n                    weight_,\\n                    bias_,\\n                    (True, False, False),\\n                )[0]\\n            )\\n        else:\\n            results.append(None)\\n\\n        # weight and bias don't compute batched gradients; no other arguments are differentiable\\n        results = results + [None] * 4\\n\\n        # set grad_sample field for weight and bias with per sample gradients\\n        if hasattr(ctx, \\\"weight\\\"):\\n            set_grad_sample_if_exists(ctx.weight, weight_per_sample_grad)\\n        if hasattr(ctx, \\\"bias\\\"):\\n            set_grad_sample_if_exists(\\n                ctx.bias,\\n                lambda bias: sum_over_all_but_batch_and_last_n(grad_output, bias.dim()),\\n            )\\n        return tuple(results)\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import List, Optional\\n\\nimport numpy as np\\n\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .expanded_weights_utils import (\\n    set_grad_sample_if_exists,\\n    unpack_expanded_weight_or_tensor,\\n)\\n\\n\\nTHRESHOLD = 32\\n\\n\\ndef conv_picker(func, conv1dOpt, conv2dOpt, conv3dOpt):\\n    if func == F.conv1d:\\n        return conv1dOpt\\n    if func == F.conv2d:\\n        return conv2dOpt\\n    else:\\n        assert func == F.conv3d\\n        return conv3dOpt\\n\\n\\ndef conv_args_and_kwargs(kwarg_names, expanded_args_and_kwargs):\\n    args = expanded_args_and_kwargs[: len(expanded_args_and_kwargs) - len(kwarg_names)]\\n    kwargs = expanded_args_and_kwargs[\\n        len(expanded_args_and_kwargs) - len(kwarg_names) :\\n    ]\\n    kwargs = dict(zip(kwarg_names, kwargs))\\n\\n    return conv_normalizer(*args, **kwargs)\\n\\n\\ndef conv_normalizer(\\n    input,\\n    weight,\\n    bias=None,\\n    stride=1,\\n    padding=0,\\n    dilation=1,\\n    groups=1,\\n):\\n    return (input, weight), {\\n        \\\"bias\\\": bias,\\n        \\\"stride\\\": stride,\\n        \\\"padding\\\": padding,\\n        \\\"dilation\\\": dilation,\\n        \\\"groups\\\": groups,\\n    }\\n\\n\\ndef conv_input_for_string_padding(func, padding_style, input, dilation, kernel_size):\\n    if padding_style == \\\"valid\\\":\\n        return input\\n    else:\\n        padding = int_padding_for_string_padding(\\n            func, padding_style, dilation, kernel_size\\n        )\\n        return F.pad(input, padding)\\n\\n\\ndef int_padding_for_string_padding(func, padding_style, dilation, kernel_size):\\n    def get_dilation(i):\\n        return dilation[i] if isinstance(dilation, tuple) else dilation\\n\\n    if padding_style == \\\"same\\\":\\n        padding: List[int] = []\\n        # F.pad needs the padding in reverse order from what conv expects\\n        for i in range(conv_picker(func, 0, 1, 2), -1, -1):\\n            padding += conv_padding_for_same(get_dilation(i), kernel_size[i])\\n        return padding\\n    elif padding_style == \\\"valid\\\":\\n        return conv_picker(func, 2, 4, 6) * (0,)\\n    else:\\n        raise RuntimeError(\\n            f\\\"got padding type of {padding_style}, only accept 'same' or 'valid'\\\"\\n        )\\n\\n\\ndef conv_padding_for_same(dilation, kernel_size):\\n    total_pad = dilation * (kernel_size - 1)\\n    left_pad = total_pad // 2\\n    right_pad = total_pad - left_pad\\n    return left_pad, right_pad\\n\\n\\ndef conv_backward(func, ctx, grad_output):\\n    def weight_grad_sample(weight):\\n        if batch_size < THRESHOLD and groups == 1:\\n            return conv_group_weight_grad_sample(\\n                ctx.input,\\n                grad_output,\\n                weight_shape,\\n                stride,\\n                padding,\\n                dilation,\\n                batch_size,\\n                func,\\n            )\\n        else:\\n            return conv_unfold_weight_grad_sample(\\n                ctx.input,\\n                grad_output,\\n                weight_shape,\\n                kernel_size,\\n                stride,\\n                padding,\\n                dilation,\\n                groups,\\n                func,\\n            )\\n\\n    def expand(param):\\n        if isinstance(param, int):\\n            return conv_picker(func, (param,), (param, param), (param, param, param))\\n        else:\\n            return param\\n\\n    def calc_total_padding(func, was_same, padding, dilation, kernel_size):\\n        if was_same:\\n            all_padding = int_padding_for_string_padding(\\n                func, \\\"same\\\", dilation, kernel_size\\n            )\\n            # F.pad needs the padding in reverse order from what conv expects\\n            total_padding = tuple(\\n                all_padding[i] + all_padding[i - 1]\\n                for i in range(len(all_padding) - 1, -1, -2)\\n            )\\n            return total_padding\\n        else:\\n            return tuple(2 * pad for pad in padding)\\n\\n    weight_shape = ctx.weight.shape\\n    stride, padding, dilation, groups = (\\n        expand(ctx.stride),\\n        expand(ctx.padding),\\n        expand(ctx.dilation),\\n        ctx.groups,\\n    )\\n\\n    kernel_size = []\\n    for i in range(2, conv_picker(func, 3, 4, 5)):\\n        kernel_size.append(weight_shape[i])\\n\\n    batch_size = ctx.batch_size\\n    results: List[Optional[torch.Tensor]] = []\\n    results.append(None)  # for kwarg names\\n    results.append(None)  # for op reference\\n\\n    # \\\"same\\\" padding may give uneven padding on either side so we need to separate the \\\"padding\\\" attr and total padding\\n    total_padding = calc_total_padding(\\n        func, ctx.was_same_padding, padding, dilation, kernel_size\\n    )\\n\\n    if ctx.input_required_grad:\\n        output_padding = []\\n        input_dims = conv_picker(func, 1, 2, 3)\\n        for i in range(input_dims):\\n            input_dim = ctx.orig_input_shape[2 + i]\\n            output_padding.append(\\n                (\\n                    total_padding[i]\\n                    + input_dim\\n                    - (kernel_size[i] * dilation[i] - dilation[i] + 1)\\n                )\\n                % stride[i]\\n            )\\n        weight_ = unpack_expanded_weight_or_tensor(ctx.weight)\\n        transpose_func = conv_picker(\\n            func, F.conv_transpose1d, F.conv_transpose2d, F.conv_transpose3d\\n        )\\n        out = transpose_func(\\n            grad_output,\\n            weight_,\\n            None,\\n            stride,\\n            padding,\\n            tuple(output_padding),\\n            groups,\\n            dilation,\\n        )\\n\\n        if ctx.was_same_padding:\\n            for i in range(len(total_padding)):\\n                out = torch.narrow(\\n                    out, 2 + i, total_padding[i] // 2, ctx.orig_input_shape[2 + i]\\n                )\\n\\n        results.append(out)\\n    else:\\n        results.append(None)\\n    # weight and bias don't compute batched gradients; no other arguments are differentiable\\n    results = results + [None] * 6\\n\\n    # set grad_sample field for weight and bias with per sample gradients\\n    set_grad_sample_if_exists(ctx.weight, weight_grad_sample)\\n    set_grad_sample_if_exists(\\n        ctx.bias, lambda _: grad_output.reshape(*grad_output.shape[:2], -1).sum(dim=2)\\n    )\\n    return tuple(results)\\n\\n\\ndef conv_unfold_weight_grad_sample(\\n    input,\\n    grad_output,\\n    weight_shape,\\n    kernel_size,\\n    stride,\\n    padding,\\n    dilation,\\n    groups,\\n    func,\\n):\\n    n = input.shape[0]\\n    in_channels = input.shape[1]\\n\\n    unfold_func = conv_picker(\\n        func,\\n        lambda: F.unfold(\\n            input.unsqueeze(-2),\\n            kernel_size=(1, kernel_size[0]),\\n            dilation=(1, dilation[0]),\\n            padding=(0, padding[0]),\\n            stride=(1, stride[0]),\\n        ),\\n        lambda: F.unfold(\\n            input, kernel_size, dilation=dilation, padding=padding, stride=stride\\n        ),\\n        lambda: unfold3d(input, kernel_size, padding, stride, dilation),\\n    )\\n\\n    input = unfold_func()\\n    grad_output = grad_output.reshape(n, -1, input.shape[-1])\\n\\n    # n=batch_sz; o=num_out_channels; p=(num_in_channels/groups)*kernel_sz\\n    weight_grad_sample = torch.einsum(\\\"noq,npq->nop\\\", grad_output, input)\\n    # rearrange the above tensor and extract diagonals.\\n    weight_grad_sample = weight_grad_sample.view(\\n        n,\\n        groups,\\n        -1,\\n        groups,\\n        int(in_channels / groups),\\n        np.prod(kernel_size),\\n    )\\n    weight_grad_sample = torch.einsum(\\n        \\\"ngrg...->ngr...\\\", weight_grad_sample\\n    ).contiguous()\\n    shape = [n] + list(weight_shape)\\n    weight_grad_sample = weight_grad_sample.view(shape)\\n    return weight_grad_sample\\n\\n\\ndef conv_group_weight_grad_sample(\\n    input,\\n    grad_output,\\n    weight_shape,\\n    stride,\\n    padding,\\n    dilation,\\n    batch_size,\\n    func,\\n):\\n    I = input.shape[1]\\n    O = grad_output.shape[1]\\n\\n    input_ = input.transpose(0, 1)\\n    grad_output_ = grad_output.view(\\n        grad_output.shape[0] * grad_output.shape[1], 1, *grad_output.shape[2:]\\n    )\\n\\n    weight_grad_sample = func(\\n        input_,\\n        grad_output_,\\n        None,\\n        stride=dilation,\\n        padding=padding,\\n        dilation=stride,\\n        groups=batch_size,\\n    )\\n    input_dims = conv_picker(func, 3, 4, 5)\\n    for i in range(2, input_dims):\\n        weight_grad_sample = weight_grad_sample.narrow(i, 0, weight_shape[i])\\n    weight_grad_sample = weight_grad_sample.view(\\n        I, batch_size, O, *weight_grad_sample.shape[2:]\\n    )\\n    weight_grad_sample = weight_grad_sample.movedim(0, 2)\\n    return weight_grad_sample\\n\\n\\ndef unfold3d(\\n    tensor,\\n    kernel_size,\\n    padding,\\n    stride,\\n    dilation,\\n):\\n    r\\\"\\\"\\\"\\n    Extract sliding local blocks from an batched input tensor.\\n\\n    :class:`torch.nn.Unfold` only supports 4D inputs (batched image-like tensors).\\n    This method implements the same action for 5D inputs\\n    Args:\\n        tensor: An input tensor of shape ``(B, C, D, H, W)``.\\n        kernel_size: the size of the sliding blocks\\n        padding: implicit zero padding to be added on both sides of input\\n        stride: the stride of the sliding blocks in the input spatial dimensions\\n        dilation: the spacing between the kernel points.\\n    Returns:\\n        A tensor of shape ``(B, C * np.prod(kernel_size), L)``, where L - output spatial dimensions.\\n        See :class:`torch.nn.Unfold` for more details\\n    Example:\\n        >>> # xdoctest: +SKIP\\n        >>> B, C, D, H, W = 3, 4, 5, 6, 7\\n        >>> tensor = torch.arange(1, B * C * D * H * W + 1.).view(B, C, D, H, W)\\n        >>> unfold3d(tensor, kernel_size=2, padding=0, stride=1).shape\\n        torch.Size([3, 32, 120])\\n    \\\"\\\"\\\"\\n    if len(tensor.shape) != 5:\\n        raise ValueError(\\n            f\\\"Input tensor must be of the shape [B, C, D, H, W]. Got{tensor.shape}\\\"\\n        )\\n\\n    if dilation != (1, 1, 1):\\n        raise NotImplementedError(f\\\"dilation={dilation} not supported.\\\")\\n\\n    batch_size, channels, _, _, _ = tensor.shape\\n\\n    # Input shape: (B, C, D, H, W)\\n    tensor = F.pad(\\n        tensor, (padding[2], padding[2], padding[1], padding[1], padding[0], padding[0])\\n    )\\n    # Output shape: (B, C, D+2*padding[2], H+2*padding[1], W+2*padding[0])\\n\\n    tensor = tensor.unfold(dimension=2, size=kernel_size[0], step=stride[0])\\n    tensor = tensor.unfold(dimension=3, size=kernel_size[1], step=stride[1])\\n    tensor = tensor.unfold(dimension=4, size=kernel_size[2], step=stride[2])\\n    # Output shape: (B, C, D_out, H_out, W_out, kernel_size[0], kernel_size[1], kernel_size[2])\\n    # For D_out, H_out, W_out definitions see :class:`torch.nn.Unfold`\\n\\n    tensor = tensor.permute(0, 2, 3, 4, 1, 5, 6, 7)\\n    # Output shape: (B, D_out, H_out, W_out, C, kernel_size[0], kernel_size[1], kernel_size[2])\\n\\n    tensor = tensor.reshape(batch_size, -1, channels * np.prod(kernel_size)).transpose(\\n        1, 2\\n    )\\n    # Output shape: (B, D_out * H_out * W_out, C * kernel_size[0] * kernel_size[1] * kernel_size[2]\\n\\n    return tensor\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import List, Optional\\n\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .expanded_weights_impl import implements_per_sample_grads\\nfrom .expanded_weights_utils import (\\n    forward_helper,\\n    set_grad_sample_if_exists,\\n    standard_kwargs,\\n)\\n\\n\\n@implements_per_sample_grads(F.embedding)\\nclass EmbeddingPerSampleGrad(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs):\\n        expanded_args, expanded_kwargs = standard_kwargs(\\n            kwarg_names, expanded_args_and_kwargs\\n        )\\n        if len(expanded_args[0].shape) == 1:\\n            raise RuntimeError(\\n                f\\\"Expanded Weights needs an input with a batch size, got a 1D tensor, {expanded_args[0]}\\\"\\n            )\\n        output = forward_helper(F.embedding, expanded_args, expanded_kwargs)\\n        ctx.input, ctx.weight = expanded_args\\n        ctx.padding_idx, ctx.scale_grad_by_freq = (\\n            expanded_kwargs[\\\"padding_idx\\\"],\\n            expanded_kwargs[\\\"scale_grad_by_freq\\\"],\\n        )\\n        ctx.sparse = expanded_kwargs[\\\"sparse\\\"]\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        input, weight = ctx.input, ctx.weight\\n        padding_idx, scale_grad_by_freq, sparse = (\\n            ctx.padding_idx,\\n            ctx.scale_grad_by_freq,\\n            ctx.sparse,\\n        )\\n\\n        def weight_per_sample_grad(weight):\\n            batch_size = input.shape[0]\\n            embedding_dim = weight.shape[1]\\n            index = (\\n                input.unsqueeze(-1)\\n                .expand(*input.shape, embedding_dim)\\n                .reshape(batch_size, -1, embedding_dim)\\n            )\\n            grad_sample = torch.zeros(\\n                batch_size, *weight.shape, device=weight.device, dtype=grad_output.dtype\\n            )\\n            return grad_sample.scatter_add_(\\n                1, index, grad_output.reshape(batch_size, -1, embedding_dim)\\n            )\\n\\n        results: List[Optional[torch.Tensor]] = []\\n        results.append(None)  # for kwarg names\\n        results.append(None)  # for op reference\\n\\n        if input.requires_grad:\\n            bw_fn = torch.ops.aten.embedding_backward\\n            results.append(\\n                bw_fn(\\n                    grad_output,\\n                    input,\\n                    weight.shape[0],\\n                    padding_idx,\\n                    scale_grad_by_freq,\\n                    sparse,\\n                )\\n            )\\n        else:\\n            results.append(None)\\n\\n        # weight doesn't compute batched gradients; no other arguments are differentiable (2 not saved from forward)\\n        results = results + [None] * 6\\n\\n        # set grad_sample field for weight with per sample gradients\\n        set_grad_sample_if_exists(weight, weight_per_sample_grad)\\n        return tuple(results)\\n\\n\\n# mypy: allow-untyped-defs\\nimport functools\\nfrom contextlib import contextmanager\\nfrom typing import Callable, Dict\\n\\nimport torch\\nfrom torch._decomp import decomposition_table\\nfrom torch.utils._pytree import tree_map_only\\n\\n\\nHANDLED_FUNCTIONS: Dict[Callable, torch.autograd.Function] = {}\\n\\naten = torch._ops.ops.aten\\n# __torch_function__ runs before the pydispatcher so we need to manually use the same\\n# decompositions indexed by their torch equivalent\\nexpanded_weights_rnn_decomps = {\\n    # func: (input_decomp, data_decomp)\\n    torch.rnn_relu: (\\n        decomposition_table[aten.rnn_relu.input],\\n        decomposition_table[aten.rnn_relu.data],\\n    ),\\n    torch.rnn_tanh: (\\n        decomposition_table[aten.rnn_tanh.input],\\n        decomposition_table[aten.rnn_tanh.data],\\n    ),\\n    torch.lstm: (\\n        decomposition_table[aten.lstm.input],\\n        decomposition_table[aten.lstm.data],\\n    ),\\n    torch.gru: (\\n        decomposition_table[aten.gru.input],\\n        decomposition_table[aten.gru.data],\\n    ),\\n}\\n\\n\\n# all of the RNN decomps run linear with the batch dimension second, even if batch_first was set\\n@contextmanager\\ndef batch_second(args, kwargs):\\n    def set_batch_second(ew):\\n        ew.set_batch_first(False)\\n\\n    def reset_batch_first(ew):\\n        ew.set_batch_first(True)\\n\\n    tree_map_only(ExpandedWeight, set_batch_second, args)\\n    tree_map_only(ExpandedWeight, set_batch_second, kwargs)\\n    try:\\n        yield\\n    finally:\\n        tree_map_only(ExpandedWeight, reset_batch_first, args)\\n        tree_map_only(ExpandedWeight, reset_batch_first, kwargs)\\n\\n\\n# to support packed sequences, we need to allow for smaller batches. Expanded weights represents the largest batch\\n@contextmanager\\ndef allow_smaller_batches(args, kwargs):\\n    def allow(ew):\\n        ew.set_allow_smaller_batches(True)\\n\\n    def reset(ew):\\n        ew.set_allow_smaller_batches(False)\\n\\n    tree_map_only(ExpandedWeight, allow, args)\\n    tree_map_only(ExpandedWeight, allow, kwargs)\\n    try:\\n        yield\\n    finally:\\n        tree_map_only(ExpandedWeight, reset, args)\\n        tree_map_only(ExpandedWeight, reset, kwargs)\\n\\n\\n@contextmanager\\ndef setup_rnn(use_input_variant, args, kwargs):\\n    with batch_second(args, kwargs) if use_input_variant else allow_smaller_batches(\\n        args, kwargs\\n    ):\\n        yield\\n\\n\\ndef implements_per_sample_grads(torch_function):\\n    @functools.wraps(torch_function)\\n    def decorator(autograd_func):\\n        HANDLED_FUNCTIONS[torch_function] = autograd_func\\n        return autograd_func\\n\\n    return decorator\\n\\n\\n# ExpandedWeight represents a weight (parameter) Tensor that has an expanded\\n# batch dimension. Operations on the ExpandedWeight Tensor act exactly like\\n# those without an expanded batch dimension but a call to .backward() populates\\n# the original (unexpanded) tensor with per-sample-gradients for in the grad_sample field\\n#\\n# ExpandedWeight has a fallback that always fails since we cannot know what the batch\\n# dimension of the input tensor is and therefore cannot know if this is a valid call\\n#\\n# This is a __torch_function__ object but it could have also been a Tensor Extension\\n# with a dispatch key.\\n#\\n# Needs to be a tensor subclass to allow reparamaterization\\nclass ExpandedWeight(torch.Tensor):\\n    def __init__(self, orig_weight, batch_size, loss_reduction):\\n        self.batch_size = batch_size\\n        self.batch_first = True\\n        self.allow_smaller_batches = False\\n        self.orig_weight = orig_weight\\n        self.loss_reduction = loss_reduction\\n\\n    handled_functions = HANDLED_FUNCTIONS\\n\\n    def __new__(cls, orig_weight, batch_size, loss_reduction):\\n        if not isinstance(orig_weight, torch.Tensor):\\n            raise RuntimeError(\\n                f\\\"Can only make Expanded Weights of Tensors, got {type(orig_weight).__name__}\\\"\\n            )\\n        if not orig_weight.requires_grad:\\n            raise RuntimeError(\\n                \\\"Can only build ExpandedWeights objects of tensors that require_grad\\\"\\n            )\\n        ret = torch.Tensor._make_subclass(cls, orig_weight, True)\\n        return ret\\n\\n    @classmethod\\n    def __torch_function__(cls, func, _, args=(), kwargs=None):\\n        if kwargs is None:\\n            kwargs = {}\\n        if func in expanded_weights_rnn_decomps:\\n            # in aten, choosing the input or data variants is done by parsing logic. This mimics some of that\\n            decomp_opts = expanded_weights_rnn_decomps[func]\\n            use_input_variant = isinstance(\\n                args[2], list\\n            )  # data variant uses a list here\\n            decomp = decomp_opts[0] if use_input_variant else decomp_opts[1]\\n\\n            if decomp is not None:\\n                with setup_rnn(use_input_variant, args, kwargs):\\n                    return decomp(*args, **kwargs)\\n        if func == torch._cudnn_rnn_flatten_weight:\\n            # since we aren't using the fused cuda kernels for RNNs, don't do this\\n            return\\n        if func in cls.handled_functions:\\n            return cls.handled_functions[func].apply(\\n                tuple(kwargs.keys()), func, *(args + tuple(kwargs.values()))\\n            )\\n        # We cannot use a fallback here because we do not know the batch dimension for any regular tensor inputs,\\n        # i.e. torch.add(torch.Tensor, ExpandedWeight)\\n        raise RuntimeError(\\n            f\\\"Expanded Weights encountered but cannot handle function {func.__name__}\\\"\\n        )\\n\\n    @property\\n    def dtype(self):\\n        return self.orig_weight.dtype\\n\\n    @property\\n    def data(self):\\n        return self.orig_weight.data\\n\\n    @property\\n    def shape(self):\\n        return self.orig_weight.shape\\n\\n    @property\\n    def device(self):\\n        return self.orig_weight.device\\n\\n    @property\\n    def is_cuda(self):\\n        return self.orig_weight.is_cuda\\n\\n    def data_ptr(self):\\n        return self.orig_weight.data_ptr()\\n\\n    def get_device(self):\\n        return self.orig_weight.get_device()\\n\\n    def set_allow_smaller_batches(self, is_allow_smaller_batches):\\n        self.allow_smaller_batches = is_allow_smaller_batches\\n\\n    def set_batch_first(self, is_batch_first=True):\\n        self.batch_first = is_batch_first\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Optional\\n\\nimport torch\\n\\nfrom .expanded_weights_impl import ExpandedWeight\\n\\n\\ndef is_batch_first(expanded_args_and_kwargs):\\n    batch_first = None\\n    for arg in expanded_args_and_kwargs:\\n        if not isinstance(arg, ExpandedWeight):\\n            continue\\n\\n        if not batch_first:\\n            batch_first = arg.batch_first\\n        elif arg.batch_first != batch_first:\\n            raise RuntimeError(\\n                \\\"Got conflicting batch_first arguments in the same layer\\\"\\n            )\\n    return batch_first\\n\\n\\ndef standard_kwargs(kwarg_names, expanded_args):\\n    r\\\"\\\"\\\"Separate args and kwargs from `__torch_function__`s that standardize kwargs.\\n\\n    Most `__torch_function__`s standardize the kwargs that they give, so this will separate\\n    the args and kwargs they pass. Functions that don't are linear and convND.\\n    \\\"\\\"\\\"\\n    kwarg_values = expanded_args[len(expanded_args) - len(kwarg_names) :]\\n    expanded_args_without_kwargs = expanded_args[\\n        : len(expanded_args) - len(kwarg_names)\\n    ]\\n    expanded_kwargs = dict(zip(kwarg_names, kwarg_values))\\n    return expanded_args_without_kwargs, expanded_kwargs\\n\\n\\ndef forward_helper(func, expanded_args, expanded_kwargs):\\n    r\\\"\\\"\\\"Compute the forward pass for a function that has expanded weight(s) passed to it.\\n\\n    It will run the forward pass where all ExpandedWeights are their original\\n    weight. It runs checks on the given arguments and detaches the outputs.\\n\\n    .. note:: First argument in :attr:`expanded_args` must be the input with the batch\\n    dimension as the first element of the shape\\n\\n    .. note:: :attr:`func` must return a Tensor or tuple of Tensors\\n\\n    Args:\\n        func: The function to be called\\n        expanded_args: Arguments to be passed to :attr:`func`. Will include arguments\\n          that need to be unpacked because they are ExpandedWeights\\n        expanded_kwargs: Keyword arguments to be passed to :attr:`func`.\\n          Similar to :attr:`expanded_args`.\\n    \\\"\\\"\\\"\\n    unexpanded_args, unexpanded_kwargs = _check_and_unexpand_args(\\n        func, expanded_args, expanded_kwargs\\n    )\\n    return func(*unexpanded_args, **unexpanded_kwargs)\\n\\n\\ndef _check_and_unexpand_args(func, expanded_args, expanded_kwargs):\\n    # input must be the first argument passed\\n    input = expanded_args[0]\\n    if isinstance(input, ExpandedWeight):\\n        raise RuntimeError(\\n            \\\"Expanded Weights do not support inputs that are also ExpandedWeights. \\\"\\n            f\\\"Input must be a Tensor, got {type(input).__name__} in function {func.__name__}\\\"\\n        )\\n    if not isinstance(input, torch.Tensor):\\n        raise RuntimeError(\\n            \\\"Expanded Weights requires a Tensor as the first input to get the batch dimension, \\\"\\n            f\\\"got {type(input).__name__} in function {func.__name__}\\\"\\n        )\\n    if len(input.shape) == 0:\\n        raise RuntimeError(\\n            f\\\"Expanded Weights requires a batch dimension but got an input of size 0 in function {func.__name__}\\\"\\n        )\\n    if input.shape[0] == 0:\\n        raise RuntimeError(\\n            \\\"0 is not a valid batch size for Expanded Weights but got input tensor of \\\"\\n            f\\\"{input} in function {func.__name__}\\\"\\n        )\\n    for arg in expanded_args + tuple(expanded_kwargs.values()):\\n        if not isinstance(arg, ExpandedWeight):\\n            continue\\n        batch_size = input.shape[0] if arg.batch_first else input.shape[1]\\n        if (arg.allow_smaller_batches and batch_size > arg.batch_size) or (\\n            not arg.allow_smaller_batches and arg.batch_size != batch_size\\n        ):\\n            raise RuntimeError(\\n                \\\"Expected ExpandedWeights to have batch size matching input but got \\\"\\n                f\\\"input batch size of {batch_size} with ExpandedWeight of batch size {arg.batch_size}\\\"\\n            )\\n\\n    loss_reduction: Optional[str] = None\\n    for arg in expanded_args + tuple(expanded_kwargs.values()):\\n        if isinstance(arg, ExpandedWeight):\\n            if loss_reduction is None:\\n                loss_reduction = arg.loss_reduction\\n            elif loss_reduction != arg.loss_reduction:\\n                raise RuntimeError(\\n                    \\\"Expected ExpandedWeights to all have the same loss_reduction argument but got one\\\"\\n                    f\\\"with {loss_reduction} and one with {arg.loss_reduction}\\\"\\n                )\\n\\n    unexpanded_args = tuple(\\n        arg.orig_weight if isinstance(arg, ExpandedWeight) else arg\\n        for arg in expanded_args\\n    )\\n    unexpanded_kwargs = {\\n        name: arg.orig_weight if isinstance(arg, ExpandedWeight) else arg\\n        for (name, arg) in expanded_kwargs.items()\\n    }\\n    return unexpanded_args, unexpanded_kwargs\\n\\n\\ndef maybe_scale_by_batch_size(grad_sample, expanded_weight):\\n    if expanded_weight.loss_reduction == \\\"mean\\\":\\n        return grad_sample * expanded_weight.batch_size\\n    else:\\n        return grad_sample\\n\\n\\ndef set_grad_sample_if_exists(maybe_expanded_weight, per_sample_grad_fn):\\n    unpacked = unpack_expanded_weight_or_tensor(maybe_expanded_weight)\\n    if isinstance(maybe_expanded_weight, ExpandedWeight):\\n        grad_sample_contribution = maybe_scale_by_batch_size(\\n            per_sample_grad_fn(unpacked), maybe_expanded_weight\\n        )\\n\\n        if maybe_expanded_weight.batch_size > grad_sample_contribution.shape[0]:\\n            # this only passes the other checks if the arg allows smaller batch sizes\\n            intermediate = torch.zeros(\\n                maybe_expanded_weight.batch_size,\\n                *grad_sample_contribution.shape[1:],\\n                dtype=grad_sample_contribution.dtype,\\n                device=grad_sample_contribution.device,\\n            )\\n            intermediate[: grad_sample_contribution.shape[0]] = grad_sample_contribution\\n            grad_sample_contribution = intermediate\\n\\n        if hasattr(unpacked, \\\"grad_sample\\\") and unpacked.grad_sample is not None:\\n            unpacked.grad_sample = unpacked.grad_sample + grad_sample_contribution\\n        else:\\n            unpacked.grad_sample = grad_sample_contribution\\n\\n\\ndef unpack_expanded_weight_or_tensor(maybe_expanded_weight, func=lambda x: x):\\n    if isinstance(maybe_expanded_weight, ExpandedWeight):\\n        orig_weight = maybe_expanded_weight.orig_weight\\n        return func(orig_weight)\\n    elif (\\n        isinstance(maybe_expanded_weight, torch.Tensor)\\n        and not maybe_expanded_weight.requires_grad\\n    ):\\n        return func(maybe_expanded_weight)\\n    elif isinstance(maybe_expanded_weight, torch.Tensor):\\n        raise RuntimeError(\\n            \\\"ExpandedWeights currently does not support a mixture of ExpandedWeight parameters \\\"\\n            \\\"and normal Parameters. Please file and issue with pytorch/pytorch\\\"\\n        )\\n\\n\\ndef sum_over_all_but_batch_and_last_n(\\n    tensor: torch.Tensor,\\n    n_dims: int,\\n) -> torch.Tensor:\\n    r\\\"\\\"\\\"\\n    Calculate the sum over all dimensions, except the first (batch dimension), and excluding the last n_dims.\\n\\n    This function will ignore the first dimension and it will\\n    not aggregate over the last n_dims dimensions.\\n    Args:\\n        tensor: An input tensor of shape ``(B, ..., X[n_dims-1])``.\\n        n_dims: Number of dimensions to keep.\\n    Example:\\n        >>> tensor = torch.ones(1, 2, 3, 4, 5)\\n        >>> sum_over_all_but_batch_and_last_n(tensor, n_dims=2).shape\\n        torch.Size([1, 4, 5])\\n    Returns:\\n        A tensor of shape ``(B, ..., X[n_dims-1])``\\n    \\\"\\\"\\\"\\n    if tensor.dim() == n_dims + 1:\\n        return tensor\\n    else:\\n        dims = list(range(1, tensor.dim() - n_dims))\\n        return tensor.sum(dim=dims)\\n\\n\\n# mypy: allow-untyped-defs\\nimport operator\\nfrom functools import reduce\\nfrom typing import List, Optional\\n\\nimport torch\\nimport torch.nn.functional as F\\n\\nfrom .expanded_weights_impl import ExpandedWeight, implements_per_sample_grads\\nfrom .expanded_weights_utils import (\\n    forward_helper,\\n    set_grad_sample_if_exists,\\n    standard_kwargs,\\n    unpack_expanded_weight_or_tensor,\\n)\\n\\n\\n@implements_per_sample_grads(F.group_norm)\\nclass GroupNormPerSampleGrad(torch.autograd.Function):\\n    @staticmethod\\n    def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs):\\n        expanded_args, expanded_kwargs = standard_kwargs(\\n            kwarg_names, expanded_args_and_kwargs\\n        )\\n        input, num_groups = expanded_args\\n        N = input.shape[0]\\n        C = input.shape[1]\\n        HxW = reduce(operator.mul, input.shape[2:], 1)\\n        weight, bias, eps = (\\n            expanded_kwargs[\\\"weight\\\"],\\n            expanded_kwargs[\\\"bias\\\"],\\n            expanded_kwargs[\\\"eps\\\"],\\n        )\\n        output, mean, rstd = forward_helper(\\n            torch.native_group_norm,\\n            (input, weight, bias, N, C, HxW, num_groups, eps),\\n            {},\\n        )\\n        ctx.input, ctx.num_groups = input, num_groups\\n        ctx.weight, ctx.eps = weight, eps\\n        ctx.mean, ctx.rstd = mean, rstd\\n        if isinstance(bias, ExpandedWeight):\\n            ctx.bias = bias\\n        if input.requires_grad and isinstance(weight, ExpandedWeight):\\n            ctx.weight = weight\\n        return output\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        input, num_groups = ctx.input, ctx.num_groups\\n        weight, bias, eps = ctx.weight, ctx.bias, ctx.eps\\n        mean, rstd = ctx.mean, ctx.rstd\\n\\n        results: List[Optional[torch.Tensor]] = []\\n        results.append(None)  # for kwarg names\\n        results.append(None)  # for op reference\\n\\n        if input.requires_grad:\\n            weight_c = unpack_expanded_weight_or_tensor(\\n                weight, lambda t: t.contiguous()\\n            )\\n            input_c = input.contiguous()\\n            grad_output_c = (\\n                grad_output.contiguous() if grad_output is not None else None\\n            )\\n            N = input.shape[0]\\n            C = input.shape[1]\\n            HxW = 1\\n            for s in input.shape[2:]:\\n                HxW *= s\\n            bw_fn = torch.ops.aten.native_group_norm_backward\\n            results.append(\\n                bw_fn(\\n                    grad_output_c,\\n                    input_c,\\n                    mean,\\n                    rstd,\\n                    weight_c,\\n                    N,\\n                    C,\\n                    HxW,\\n                    num_groups,\\n                    (True, False, False),\\n                )[0]\\n            )\\n        else:\\n            results.append(None)\\n\\n        # weight and bias don't compute batched gradients; no other arguments are differentiable\\n        results = results + [None] * 4\\n\\n        # set grad_sample field for weight and bias with per sample gradients\\n        if hasattr(ctx, \\\"weight\\\"):\\n            set_grad_sample_if_exists(\\n                weight,\\n                lambda _: torch.einsum(\\n                    \\\"ni...->ni\\\", F.group_norm(input, num_groups, eps=eps) * grad_output\\n                ),\\n            )\\n        if hasattr(ctx, \\\"bias\\\"):\\n            set_grad_sample_if_exists(\\n                bias, lambda _: torch.einsum(\\\"ni...->ni\\\", grad_output)\\n            )\\n        return tuple(results)\\n\\n\\nfrom .conv_expanded_weights import ConvPerSampleGrad\\nfrom .embedding_expanded_weights import EmbeddingPerSampleGrad\\nfrom .expanded_weights_impl import ExpandedWeight\\nfrom .group_norm_expanded_weights import GroupNormPerSampleGrad\\nfrom .instance_norm_expanded_weights import InstanceNormPerSampleGrad\\nfrom .layer_norm_expanded_weights import LayerNormPerSampleGrad\\nfrom .linear_expanded_weights import LinearPerSampleGrad\\n\\n\\n__all__ = [\\\"ExpandedWeight\\\"]\\n\\n\\n# mypy: allow-untyped-defs\\n# this is for historical pickle deserialization, it is not used otherwise\\n\\n\\ndef _get_thnn_function_backend():\\n    pass\\n\\n\\n\\n\\nfrom torch.nn.quantizable.modules import *  # noqa: F403\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantizable Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantizable`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantizable/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.quantizable.modules.rnn import LSTM, LSTMCell\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Quantizable Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/quantizable`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/quantizable/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\nfrom torch.ao.nn.quantizable.modules.activation import MultiheadAttention\\n\\n\\nfrom torch.ao.nn.quantizable.modules.activation import MultiheadAttention\\nfrom torch.ao.nn.quantizable.modules.rnn import LSTM, LSTMCell\\n\\n\\n__all__ = [\\n    \\\"LSTM\\\",\\n    \\\"LSTMCell\\\",\\n    \\\"MultiheadAttention\\\",\\n]\\n\\n\\n# mypy: allow-untyped-defs\\nimport warnings\\nfrom typing import List\\n\\nimport torch\\nfrom torch._utils import (\\n    _flatten_dense_tensors,\\n    _get_device_index,\\n    _handle_complex,\\n    _reorder_tensors_as,\\n    _take_tensors,\\n    _unflatten_dense_tensors,\\n)\\nfrom torch.cuda import nccl\\n\\n\\ndef broadcast(tensor, devices=None, *, out=None):\\n    r\\\"\\\"\\\"Broadcasts a tensor to specified GPU devices.\\n\\n    Args:\\n        tensor (Tensor): tensor to broadcast. Can be on CPU or GPU.\\n        devices (Iterable[torch.device, str or int], optional): an iterable of\\n          GPU devices, among which to broadcast.\\n        out (Sequence[Tensor], optional, keyword-only): the GPU tensors to\\n          store output results.\\n\\n    .. note::\\n        Exactly one of :attr:`devices` and :attr:`out` must be specified.\\n\\n    Returns:\\n        - If :attr:`devices` is specified,\\n            a tuple containing copies of :attr:`tensor`, placed on\\n            :attr:`devices`.\\n        - If :attr:`out` is specified,\\n            a tuple containing :attr:`out` tensors, each containing a copy of\\n            :attr:`tensor`.\\n    \\\"\\\"\\\"\\n    tensor = _handle_complex(tensor)\\n    if not ((devices is None) ^ (out is None)):\\n        raise RuntimeError(\\n            f\\\"Exactly one of 'devices' and 'out' must be specified, but got devices={devices} and out={out}\\\"\\n        )\\n    if devices is not None:\\n        devices = [_get_device_index(d) for d in devices]\\n        return torch._C._broadcast(tensor, devices)\\n    else:\\n        return torch._C._broadcast_out(tensor, out)\\n\\n\\ndef broadcast_coalesced(tensors, devices, buffer_size=10485760):\\n    \\\"\\\"\\\"Broadcast a sequence of tensors to the specified GPUs.\\n\\n    Small tensors are first coalesced into a buffer to reduce the number of synchronizations.\\n\\n    Args:\\n        tensors (sequence): tensors to broadcast. Must be on the same device,\\n          either CPU or GPU.\\n        devices (Iterable[torch.device, str or int]): an iterable of GPU\\n          devices, among which to broadcast.\\n        buffer_size (int): maximum size of the buffer used for coalescing\\n\\n    Returns:\\n        A tuple containing copies of :attr:`tensor`, placed on :attr:`devices`.\\n    \\\"\\\"\\\"\\n    devices = [_get_device_index(d) for d in devices]\\n    tensors = [_handle_complex(t) for t in tensors]\\n    return torch._C._broadcast_coalesced(tensors, devices, buffer_size)\\n\\n\\ndef reduce_add(inputs, destination=None):\\n    \\\"\\\"\\\"Sum tensors from multiple GPUs.\\n\\n    All inputs should have matching shapes, dtype, and layout. The output tensor\\n    will be of the same shape, dtype, and layout.\\n\\n    Args:\\n        inputs (Iterable[Tensor]): an iterable of tensors to add.\\n        destination (int, optional): a device on which the output will be\\n            placed (default: current device).\\n\\n    Returns:\\n        A tensor containing an elementwise sum of all inputs, placed on the\\n        :attr:`destination` device.\\n    \\\"\\\"\\\"\\n    destination = _get_device_index(destination, optional=True)\\n    input_size = inputs[0].size()\\n    root_index = None  # index of input tensor that already is on the correct device\\n    for i, inp in enumerate(inputs):\\n        assert inp.device.type != \\\"cpu\\\", \\\"reduce_add expects all inputs to be on GPUs\\\"\\n        if inp.get_device() == destination:\\n            root_index = i\\n        if inp.size() != input_size:\\n            got = \\\"x\\\".join(str(x) for x in inp.size())\\n            expected = \\\"x\\\".join(str(x) for x in input_size)\\n            raise ValueError(\\n                f\\\"input {i} has invalid size: got {got}, but expected {expected}\\\"\\n            )\\n    if root_index is None:\\n        raise RuntimeError(\\n            \\\"reduce_add expects destination to be on the same GPU with one of the tensors\\\"\\n        )\\n\\n    if len(inputs) == 1:\\n        return inputs[0]\\n\\n    if nccl.is_available(inputs):\\n        result = torch.empty_like(inputs[root_index])\\n        nccl.reduce(inputs, output=result, root=root_index)\\n    else:\\n        destination_device = torch.device(inputs[root_index].device.type, destination)\\n        nonroot = [t for i, t in enumerate(inputs) if i != root_index]\\n        # make a new tensor w/o clone\\n        result = inputs[root_index] + nonroot[0].to(\\n            device=destination_device, non_blocking=True\\n        )\\n        for other in nonroot[1:]:\\n            result.add_(other.to(device=destination_device, non_blocking=True))\\n    return result\\n\\n\\ndef reduce_add_coalesced(inputs, destination=None, buffer_size=10485760):\\n    \\\"\\\"\\\"Sum tensors from multiple GPUs.\\n\\n    Small tensors are first coalesced into a buffer to reduce the number\\n    of synchronizations.\\n\\n    Args:\\n        inputs (Iterable[Iterable[Tensor]]): iterable of iterables that\\n            contain tensors from a single device.\\n        destination (int, optional): a device on which the output will be\\n            placed (default: current device).\\n        buffer_size (int): maximum size of the buffer used for coalescing\\n\\n    Returns:\\n        A tuple of tensors containing an elementwise sum of each group of\\n        inputs, placed on the ``destination`` device.\\n    \\\"\\\"\\\"\\n    # TODO: When `len(inputs) == 1` and all inputs are on `destination`, just\\n    #       return `inputs`.\\n    dense_tensors: List[List] = [[] for _ in inputs]  # shape (num_gpus, num_tensors)\\n    output = []\\n    ref_order = []\\n    # process sparse ones first since they may have different sizes on different gpus\\n    for tensor_at_gpus in zip(*inputs):\\n        if all(t.is_sparse for t in tensor_at_gpus):\\n            result = reduce_add(tensor_at_gpus, destination)  # this will be sparse too\\n            output.append(result)\\n            ref_order.append(tensor_at_gpus[0])\\n        else:\\n            for coll, t in zip(dense_tensors, tensor_at_gpus):\\n                coll.append(t.to_dense() if t.is_sparse else t)\\n            ref_order.append(dense_tensors[0][-1])\\n    itrs = [_take_tensors(tensors, buffer_size) for tensors in dense_tensors]\\n    # now the dense ones, which have consistent sizes\\n    for chunks in zip(*itrs):\\n        flat_tensors = [\\n            _flatten_dense_tensors(chunk) for chunk in chunks\\n        ]  # (num_gpus,)\\n        flat_result = reduce_add(flat_tensors, destination)\\n        for t in _unflatten_dense_tensors(flat_result, chunks[0]):\\n            # The unflattened tensors do not share storage, and we don't expose\\n            # base flat tensor anyways, so give them different version counters.\\n            # See NOTE [ Version Counter in comm.*_coalesced ]\\n            output.append(t.data)\\n    return tuple(_reorder_tensors_as(output, ref_order))\\n\\n\\ndef scatter(tensor, devices=None, chunk_sizes=None, dim=0, streams=None, *, out=None):\\n    \\\"\\\"\\\"Scatters tensor across multiple GPUs.\\n\\n    Args:\\n        tensor (Tensor): tensor to scatter. Can be on CPU or GPU.\\n        devices (Iterable[torch.device, str or int], optional): an iterable of\\n          GPU devices, among which to scatter.\\n        chunk_sizes (Iterable[int], optional): sizes of chunks to be placed on\\n          each device. It should match :attr:`devices` in length and sums to\\n          ``tensor.size(dim)``. If not specified, :attr:`tensor` will be divided\\n          into equal chunks.\\n        dim (int, optional): A dimension along which to chunk :attr:`tensor`.\\n          Default: ``0``.\\n        streams (Iterable[torch.cuda.Stream], optional): an iterable of Streams, among\\n          which to execute the scatter. If not specified, the default stream will\\n          be utilized.\\n        out (Sequence[Tensor], optional, keyword-only): the GPU tensors to\\n          store output results. Sizes of these tensors must match that of\\n          :attr:`tensor`, except for :attr:`dim`, where the total size must\\n          sum to ``tensor.size(dim)``.\\n\\n    .. note::\\n        Exactly one of :attr:`devices` and :attr:`out` must be specified. When\\n        :attr:`out` is specified, :attr:`chunk_sizes` must not be specified and\\n        will be inferred from sizes of :attr:`out`.\\n\\n    Returns:\\n        - If :attr:`devices` is specified,\\n            a tuple containing chunks of :attr:`tensor`, placed on\\n            :attr:`devices`.\\n        - If :attr:`out` is specified,\\n            a tuple containing :attr:`out` tensors, each containing a chunk of\\n            :attr:`tensor`.\\n    \\\"\\\"\\\"\\n    tensor = _handle_complex(tensor)\\n    if out is None:\\n        devices = [_get_device_index(d) for d in devices]\\n        return tuple(torch._C._scatter(tensor, devices, chunk_sizes, dim, streams))\\n    else:\\n        if devices is not None:\\n            raise RuntimeError(\\n                f\\\"'devices' must not be specified when 'out' is specified, but got devices={devices}\\\"\\n            )\\n        if chunk_sizes is not None:\\n            raise RuntimeError(\\n                f\\\"'chunk_sizes' must not be specified when 'out' is specified, but got chunk_sizes={chunk_sizes}\\\"\\n            )\\n        return tuple(torch._C._scatter_out(tensor, out, dim, streams))\\n\\n\\ndef gather(tensors, dim=0, destination=None, *, out=None):\\n    r\\\"\\\"\\\"Gathers tensors from multiple GPU devices.\\n\\n    Args:\\n        tensors (Iterable[Tensor]): an iterable of tensors to gather.\\n          Tensor sizes in all dimensions other than :attr:`dim` have to match.\\n        dim (int, optional): a dimension along which the tensors will be\\n          concatenated. Default: ``0``.\\n        destination (torch.device, str, or int, optional): the output device.\\n          Can be CPU or CUDA. Default: the current CUDA device.\\n        out (Tensor, optional, keyword-only): the tensor to store gather result.\\n          Its sizes must match those of :attr:`tensors`, except for :attr:`dim`,\\n          where the size must equal ``sum(tensor.size(dim) for tensor in tensors)``.\\n          Can be on CPU or CUDA.\\n\\n    .. note::\\n        :attr:`destination` must not be specified when :attr:`out` is specified.\\n\\n    Returns:\\n        - If :attr:`destination` is specified,\\n            a tensor located on :attr:`destination` device, that is a result of\\n            concatenating :attr:`tensors` along :attr:`dim`.\\n        - If :attr:`out` is specified,\\n            the :attr:`out` tensor, now containing results of concatenating\\n            :attr:`tensors` along :attr:`dim`.\\n    \\\"\\\"\\\"\\n    tensors = [_handle_complex(t) for t in tensors]\\n    if out is None:\\n        if destination == -1:\\n            warnings.warn(\\n                \\\"Using -1 to represent CPU tensor is deprecated. Please use a \\\"\\n                'device object or string instead, e.g., \\\"cpu\\\".',\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n        destination = _get_device_index(destination, allow_cpu=True, optional=True)\\n        return torch._C._gather(tensors, dim, destination)\\n    else:\\n        if destination is not None:\\n            raise RuntimeError(\\n                f\\\"'destination' must not be specified when 'out' is specified, but got destination={destination}\\\"\\n            )\\n        return torch._C._gather_out(tensors, out, dim)\\n\\n\\nimport warnings\\nfrom typing import List, Optional\\n\\nimport torch\\nfrom torch._utils import _get_device_index\\nfrom torch.autograd import Function\\nfrom torch.nn.parallel import comm\\n\\n\\nclass Broadcast(Function):\\n    @staticmethod\\n    def forward(ctx, target_gpus, *inputs):\\n        assert all(\\n            i.device.type != \\\"cpu\\\" for i in inputs\\n        ), \\\"Broadcast function not implemented for CPU tensors\\\"\\n        target_gpus = [_get_device_index(x, True) for x in target_gpus]\\n        ctx.target_gpus = target_gpus\\n        if len(inputs) == 0:\\n            return ()\\n        ctx.num_inputs = len(inputs)\\n        ctx.input_device = inputs[0].get_device()\\n        outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)\\n        non_differentiables = []\\n        for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]):\\n            if not input_requires_grad:\\n                for output in outputs:\\n                    non_differentiables.append(output[idx])\\n        ctx.mark_non_differentiable(*non_differentiables)\\n        return tuple([t for tensors in outputs for t in tensors])\\n\\n    @staticmethod\\n    def backward(ctx, *grad_outputs):\\n        return (None,) + ReduceAddCoalesced.apply(\\n            ctx.input_device, ctx.num_inputs, *grad_outputs\\n        )\\n\\n\\nclass ReduceAddCoalesced(Function):\\n    @staticmethod\\n    def forward(ctx, destination, num_inputs, *grads):\\n        ctx.target_gpus = [\\n            grads[i].get_device() for i in range(0, len(grads), num_inputs)\\n        ]\\n\\n        grads_ = [grads[i : i + num_inputs] for i in range(0, len(grads), num_inputs)]\\n        return comm.reduce_add_coalesced(grads_, destination)\\n\\n    @staticmethod\\n    def backward(ctx, *grad_outputs):\\n        return (\\n            None,\\n            None,\\n        ) + Broadcast.apply(ctx.target_gpus, *grad_outputs)\\n\\n\\nclass Gather(Function):\\n    @staticmethod\\n    def forward(ctx, target_device, dim, *inputs):\\n        assert all(\\n            i.device.type != \\\"cpu\\\" for i in inputs\\n        ), \\\"Gather function not implemented for CPU tensors\\\"\\n        if target_device == \\\"cpu\\\":\\n            ctx.target_device = \\\"cpu\\\"\\n        else:\\n            target_device = _get_device_index(target_device, True)\\n            ctx.target_device = target_device\\n        ctx.dim = dim\\n        ctx.input_gpus = tuple(i.get_device() for i in inputs)\\n        if all(t.dim() == 0 for t in inputs) and dim == 0:\\n            inputs = tuple(t.view(1) for t in inputs)\\n            warnings.warn(\\n                \\\"Was asked to gather along dimension 0, but all \\\"\\n                \\\"input tensors were scalars; will instead unsqueeze \\\"\\n                \\\"and return a vector.\\\"\\n            )\\n            ctx.unsqueezed_scalar = True\\n        else:\\n            ctx.unsqueezed_scalar = False\\n        ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)\\n        return comm.gather(inputs, ctx.dim, ctx.target_device)\\n\\n    @staticmethod\\n    def backward(ctx, grad_output):\\n        scattered_grads = Scatter.apply(\\n            ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output\\n        )\\n        if ctx.unsqueezed_scalar:\\n            scattered_grads = tuple(g[0] for g in scattered_grads)\\n        return (None, None) + scattered_grads\\n\\n\\nclass Scatter(Function):\\n    @staticmethod\\n    def forward(ctx, target_gpus, chunk_sizes, dim, input):\\n        target_gpus = [_get_device_index(x, True) for x in target_gpus]\\n        ctx.dim = dim\\n        ctx.input_device = input.get_device() if input.device.type != \\\"cpu\\\" else -1\\n        streams = None\\n        if torch.cuda.is_available() and ctx.input_device == -1:\\n            # Perform CPU to GPU copies in a background stream\\n            streams = [\\n                _get_stream(torch.device(\\\"cuda\\\", device)) for device in target_gpus\\n            ]\\n        outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)\\n        # Synchronize with the copy stream\\n        if streams is not None:\\n            for i, output in enumerate(outputs):\\n                with torch.cuda.device(target_gpus[i]):\\n                    main_stream = torch.cuda.current_stream()\\n                    main_stream.wait_stream(streams[i])\\n                    output.record_stream(main_stream)\\n        return outputs\\n\\n    @staticmethod\\n    def backward(ctx, *grad_output):\\n        return None, None, None, Gather.apply(ctx.input_device, ctx.dim, *grad_output)\\n\\n\\n# background streams used for copying\\n_streams: Optional[List[Optional[torch.Stream]]] = None\\n\\n\\ndef _get_stream(device: torch.device):\\n    \\\"\\\"\\\"Get a background stream for copying between CPU and target device.\\\"\\\"\\\"\\n    global _streams\\n    if device.type == \\\"cpu\\\":\\n        return None\\n    device_mod = getattr(torch, device.type, None)\\n    if device_mod is None:\\n        return None\\n    if _streams is None:\\n        _streams = [None] * device_mod.device_count()\\n    if _streams[device.index] is None:\\n        _streams[device.index] = device_mod.Stream(device.index)\\n    return _streams[device.index]\\n\\n\\n# mypy: allow-untyped-defs\\nimport copy\\nimport functools\\nimport inspect\\nimport itertools\\nimport logging\\nimport os\\nimport sys\\nimport warnings\\nimport weakref\\nfrom collections import defaultdict, deque\\nfrom contextlib import contextmanager\\nfrom dataclasses import dataclass, fields, is_dataclass\\nfrom enum import auto, Enum\\nfrom typing import Any, Callable, List, Optional, Tuple, Type, TYPE_CHECKING\\n\\nimport torch\\nimport torch.distributed as dist\\nfrom torch._utils import _get_device_index\\nfrom torch.autograd import Function, Variable\\nfrom torch.distributed.algorithms.join import Join, Joinable, JoinHook\\nfrom torch.nn.modules import Module\\nfrom torch.nn.parallel.scatter_gather import gather, scatter_kwargs\\nfrom torch.utils._pytree import tree_flatten, tree_unflatten\\n\\n\\nRPC_AVAILABLE = False\\nif dist.is_available():\\n    from torch.distributed.distributed_c10d import (\\n        _get_default_group,\\n        _rank_not_in_group,\\n        ReduceOp,\\n    )\\n    from torch.distributed.utils import (\\n        _alloc_storage,\\n        _cast_forward_inputs,\\n        _free_storage,\\n        _sync_module_states,\\n        _to_kwargs,\\n        _verify_param_shape_across_processes,\\n    )\\nif dist.rpc.is_available():\\n    RPC_AVAILABLE = True\\n    from torch.distributed.rpc import RRef\\n\\nif TYPE_CHECKING:\\n    from torch.utils.hooks import RemovableHandle\\n\\n\\n__all__ = [\\\"DistributedDataParallel\\\"]\\n\\nlogger = logging.getLogger(__name__)\\n\\n\\n@dataclass\\nclass _MixedPrecision:\\n    \\\"\\\"\\\"\\n    This configures DDP-native mixed precision training.\\n\\n    Attributes:\\n        param_dtype (torch.dtype): This specifies the dtype for model\\n            parameters, inputs (when ``cast_forward_inputs`` is set to\\n            ``True``), and therefore the dtype for computation.\\n            However, outside the forward and backward passes, parameters are in\\n            full precision. Model checkpointing always happens in full\\n            precision.\\n        reduce_dtype (torch.dtype): This specifies the dtype for gradient\\n            reduction, which is permitted to differ from ``param_dtype``.\\n        buffer_dtype (torch.dtype): This specifies the dtype for buffers.\\n\\n    .. note:: This API is experimental and subject to change.\\n\\n    .. note:: Only floating point tensors are cast to their specified dtypes.\\n\\n    .. note:: ``state_dict`` checkpoints parameters and buffers in full\\n        precision.\\n\\n    .. note:: Each low precision dtype must be specified explicitly. For\\n        example, ``_MixedPrecision(reduce_dtype=torch.float16)`` only specifies\\n        the reduction dtype to be low precision, and DDP will not cast\\n        parameters or buffers.\\n\\n    .. note:: If a ``reduce_dtype`` is not specified, then gradient reduction\\n        happens in ``param_dtype`` if specified or the original parameter dtype\\n        otherwise. For example, ``_MixedPrecision(param_dtype=torch.float16)``\\n        would result in communication occurring in fp16.\\n    \\\"\\\"\\\"\\n\\n    param_dtype: Optional[torch.dtype] = None\\n    reduce_dtype: Optional[torch.dtype] = None\\n    buffer_dtype: Optional[torch.dtype] = None\\n    # TODO (rohan-varma): keep_low_precision_grads: bool = False\\n    # TODO (rohan-varma): APIs to allow users to run batchnorm and layernorm\\n    # in full precision. For DDP, this can be implemented by not performing the\\n    # parameter cast for BN and LN units.\\n\\n\\ndef _cast_buffers(mixed_precision_config, root_module):\\n    \\\"\\\"\\\"Casts buffers to the given ``buffer_dtype``.\\\"\\\"\\\"\\n    for buf in root_module.buffers():\\n        if hasattr(buf, \\\"_ddp_ignored\\\") and buf._ddp_ignored:\\n            continue\\n\\n        buf.data = buf.to(dtype=mixed_precision_config.buffer_dtype)\\n\\n\\ndef _setup_mixed_precision_params(mixed_precision_config, root_module):\\n    \\\"\\\"\\\"Create and free storage for the mixed precision parameters.\\\"\\\"\\\"\\n    for param in root_module.parameters():\\n        # Do not setup mixed precision for DDP ignored parameters.\\n        if hasattr(param, \\\"_ddp_ignored\\\") and param._ddp_ignored:\\n            continue\\n\\n        if not hasattr(param, \\\"_mp_param\\\"):\\n            param._mp_param = torch.zeros_like(\\n                param,\\n                device=param.device,\\n                dtype=mixed_precision_config.param_dtype,\\n                requires_grad=param.requires_grad,\\n            )\\n            _free_storage(param._mp_param)\\n            # _fp_param will point to the full precision param so it can be switched\\n            # back to at the end of forward / backward.\\n            param._fp_param = param.data\\n\\n\\ndef _tree_flatten_with_rref(output):\\n    output_is_rref = RPC_AVAILABLE and isinstance(output, RRef)\\n    if output_is_rref:\\n        output_tensor_list, treespec = tree_flatten(output.local_value())\\n    else:\\n        output_tensor_list, treespec = tree_flatten(output)\\n    # Need to return flattened tensors, spec to re-pack them, as well\\n    # as if the return type was actually an RRef to reconstruct.\\n    return output_tensor_list, treespec, output_is_rref\\n\\n\\ndef _tree_unflatten_with_rref(output, treespec, output_is_rref):\\n    output = tree_unflatten(output, treespec)\\n    if output_is_rref:\\n        output = RRef(output)\\n    return output\\n\\n\\ndef _find_tensors(obj):\\n    r\\\"\\\"\\\"Recursively find all tensors contained in the specified object.\\\"\\\"\\\"\\n    if RPC_AVAILABLE and isinstance(obj, RRef):\\n        # If the current node is the owner of the RRef, unwrap it and try to\\n        # find Tensors.\\n        # TODO: Expand to remote RRefs.\\n        if obj.is_owner():\\n            return _find_tensors(obj.local_value())\\n    if isinstance(obj, torch.Tensor):\\n        return [obj]\\n    if isinstance(obj, (list, tuple)):\\n        return itertools.chain.from_iterable(map(_find_tensors, obj))\\n    if isinstance(obj, dict):\\n        return itertools.chain.from_iterable(map(_find_tensors, obj.values()))\\n    if is_dataclass(obj):\\n        return itertools.chain.from_iterable(\\n            map(_find_tensors, (getattr(obj, f.name) for f in fields(obj)))\\n        )\\n\\n    return []\\n\\n\\ndef _dump_DDP_relevant_env_vars():\\n    relevant_env_vars = [\\n        \\\"RANK\\\",\\n        \\\"LOCAL_RANK\\\",\\n        \\\"WORLD_SIZE\\\",\\n        \\\"MASTER_PORT\\\",\\n        \\\"MASTER_ADDR\\\",\\n        \\\"CUDA_VISIBLE_DEVICES\\\",\\n        \\\"GLOO_SOCKET_IFNAME\\\",\\n        \\\"GLOO_DEVICE_TRANSPORT\\\",\\n        \\\"NCCL_SOCKET_IFNAME\\\",\\n        \\\"TORCH_NCCL_BLOCKING_WAIT\\\",\\n        \\\"NCCL_DEBUG\\\",\\n        \\\"NCCL_DEBUG_SUBSYS\\\",\\n        \\\"NCCL_IB_DISABLE\\\",\\n        # More NCCL env vars:\\n        \\\"NCCL_P2P_DISABLE\\\",\\n        \\\"NCCL_P2P_LEVEL\\\",\\n        \\\"NCCL_SHM_DISABLE\\\",\\n        \\\"NCCL_SOCKET_NTHREADS\\\",\\n        \\\"NCCL_NSOCKS_PERTHREAD\\\",\\n        \\\"NCCL_BUFFSIZE\\\",\\n        \\\"NCCL_NTHREADS\\\",\\n        \\\"NCCL_RINGS\\\",\\n        \\\"NCCL_MAX_NCHANNELS\\\",\\n        \\\"NCCL_MIN_NCHANNELS\\\",\\n        \\\"NCCL_CHECKS_DISABLE\\\",\\n        \\\"NCCL_CHECK_POINTERS\\\",\\n        \\\"NCCL_LAUNCH_MODE\\\",\\n        \\\"NCCL_IB_HCA\\\",\\n        \\\"NCCL_IB_TIMEOUT\\\",\\n        \\\"NCCL_IB_RETRY_CNT\\\",\\n        \\\"NCCL_IB_GID_INDEX\\\",\\n        \\\"NCCL_IB_SL\\\",\\n        \\\"NCCL_IB_TC\\\",\\n        \\\"NCCL_IB_AR_THRESHOLD\\\",\\n        \\\"NCCL_IB_CUDA_SUPPORT\\\",\\n        \\\"NCCL_NET_GDR_LEVEL\\\",\\n        \\\"NCCL_NET_GDR_READ\\\",\\n        \\\"NCCL_SINGLE_RING_THRESHOLD\\\",\\n        \\\"NCCL_LL_THRESHOLD\\\",\\n        \\\"NCCL_TREE_THRESHOLD\\\",\\n        \\\"NCCL_ALGO\\\",\\n        \\\"NCCL_PROTO\\\",\\n        \\\"NCCL_IGNORE_CPU_AFFINITY\\\",\\n        \\\"NCCL_DEBUG_FILE\\\",\\n        \\\"NCCL_COLLNET_ENABLE\\\",\\n        \\\"NCCL_TOPO_FILE\\\",\\n        \\\"NCCL_TOPO_DUMP_FILE\\\",\\n        \\\"TORCH_NCCL_ASYNC_ERROR_HANDLING\\\",\\n    ]\\n    formatted_output = \\\"\\\"\\n    for var in relevant_env_vars:\\n        value = os.environ[var] if var in os.environ else \\\"N/A\\\"\\n        formatted_output += f\\\"env:{var}={value}\\\\n\\\"\\n    print(formatted_output)\\n\\n\\nclass _BufferCommHookLocation(Enum):\\n    PRE_FORWARD = auto()\\n    POST_FORWARD = auto()\\n\\n\\n@dataclass\\nclass _BufferCommHook:\\n    buffer_comm_hook: Callable\\n    buffer_comm_hook_state: Any\\n    buffer_comm_hook_location: _BufferCommHookLocation\\n\\n\\n# Add a DDPSink to run various functions when backwards starts, such as\\n# queueing call back of out-most backward/graph task,\\n# this helps call back is fired after all gradients' calculation\\n# is completed.\\nclass _DDPSink(Function):\\n    @staticmethod\\n    def forward(ctx, ddp_weakref, *inputs):\\n        # set_materialize_grads(False) will ensure that None gradients stay as\\n        # None and are not filled with zeros.\\n        ctx.set_materialize_grads(False)\\n        ctx.ddp_weakref = ddp_weakref\\n        ret = inputs\\n        if ddp_weakref()._ddp_sink_clone:\\n            ret = tuple(\\n                inp.clone() if isinstance(inp, torch.Tensor) else inp for inp in inputs\\n            )\\n        return ret\\n\\n    @staticmethod\\n    def backward(ctx, *grad_outputs):\\n        # Enqueue delay allreduce for static graph training on the first\\n        # iteration.\\n        ddp_weakref = ctx.ddp_weakref()\\n        reducer = ddp_weakref.reducer\\n        static_graph = ddp_weakref.static_graph\\n        delay_ar_enqueued = (\\n            static_graph and ddp_weakref._static_graph_delay_allreduce_enqueued\\n        )\\n        if static_graph and not delay_ar_enqueued:\\n            Variable._execution_engine.queue_callback(  # type: ignore[call-arg,misc]\\n                reducer._delay_all_reduce\\n            )\\n            ddp_weakref._static_graph_delay_allreduce_enqueued = True\\n\\n        return (None, *grad_outputs)\\n\\n\\nclass _DDPJoinHook(JoinHook):\\n    def __init__(self, ddp, divide_by_initial_world_size):\\n        \\\"\\\"\\\"Set config variables for internal usage.\\\"\\\"\\\"\\n        assert isinstance(ddp, DistributedDataParallel), (\\n            \\\"DDP join hook requires passing in a DistributedDataParallel \\\"\\n            \\\"instance as the state\\\"\\n        )\\n        assert ddp.logger is not None\\n        ddp.logger._set_uneven_input_join()\\n        self.ddp = ddp\\n        self.ddp._divide_by_initial_world_size = divide_by_initial_world_size\\n        super().__init__()\\n\\n    def main_hook(self):\\n        \\\"\\\"\\\"Shadow the DDP collective communication operations in the forward and backward passes.\\\"\\\"\\\"\\n        ddp = self.ddp\\n        # Buckets are rebuilt only once during a training period\\n        ddp.reducer._rebuild_buckets()\\n\\n        # Schedule a broadcast if we are syncing module buffers in the\\n        # forward pass\\n        # TODO: make DDP uneven inputs context manager support buffer\\n        # comm hook (https://github.com/pytorch/pytorch/issues/65436)\\n        ddp._check_and_sync_module_buffers()\\n\\n        # Check if need to sync in the backward pass\\n        should_sync_backwards = ddp._check_global_requires_backward_grad_sync(\\n            is_joined_rank=True\\n        )\\n        # Forward parameter sync is disabled in the next iteration if we\\n        # are skipping gradient sync this iteration, so set\\n        # `require_forward_param_sync` accordingly\\n        ddp.require_forward_param_sync = should_sync_backwards\\n        if not should_sync_backwards:\\n            return\\n\\n        # Schedule one allreduce per gradient bucket to match the backward\\n        # pass allreduce\\n        ddp._match_all_reduce_for_bwd_pass()\\n\\n        # Check if we need to allreduce locally unused parameters\\n        if ddp.find_unused_parameters:\\n            ddp._match_unused_params_allreduce()\\n\\n        # Rebuilt parameters are pushed only once during a training period\\n        ddp.reducer._push_all_rebuilt_params()\\n\\n    def post_hook(self, is_last_joiner: bool):\\n        \\\"\\\"\\\"Sync the final model to ensure that the model is the same across all processes.\\\"\\\"\\\"\\n        self.ddp._sync_final_model(is_last_joiner)\\n\\n\\nclass DistributedDataParallel(Module, Joinable):\\n    r\\\"\\\"\\\"Implement distributed data parallelism based on ``torch.distributed`` at module level.\\n\\n    This container provides data parallelism by synchronizing gradients\\n    across each model replica. The devices to synchronize across are\\n    specified by the input ``process_group``, which is the entire world\\n    by default. Note that ``DistributedDataParallel`` does not chunk or\\n    otherwise shard the input across participating GPUs; the user is\\n    responsible for defining how to do so, for example through the use\\n    of a :class:`DistributedSampler`.\\n\\n    See also: :ref:`distributed-basics` and :ref:`cuda-nn-ddp-instead`.\\n    The same constraints on input as in :class:`torch.nn.DataParallel` apply.\\n\\n    Creation of this class requires that ``torch.distributed`` to be already\\n    initialized, by calling :func:`torch.distributed.init_process_group`.\\n\\n    ``DistributedDataParallel`` is proven to be significantly faster than\\n    :class:`torch.nn.DataParallel` for single-node multi-GPU data\\n    parallel training.\\n\\n    To use ``DistributedDataParallel`` on a host with N GPUs, you should spawn\\n    up ``N`` processes, ensuring that each process exclusively works on a single\\n    GPU from 0 to N-1. This can be done by either setting\\n    ``CUDA_VISIBLE_DEVICES`` for every process or by calling:\\n\\n        >>> # xdoctest: +SKIP(\\\"undefined variables\\\")\\n        >>> torch.cuda.set_device(i)\\n\\n    where i is from 0 to N-1. In each process, you should refer the following\\n    to construct this module:\\n\\n        >>> # xdoctest: +SKIP(\\\"undefined variables\\\")\\n        >>> torch.distributed.init_process_group(\\n        >>>     backend='nccl', world_size=N, init_method='...'\\n        >>> )\\n        >>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)\\n\\n    In order to spawn up multiple processes per node, you can use either\\n    ``torch.distributed.launch`` or ``torch.multiprocessing.spawn``.\\n\\n    .. note::\\n        Please refer to `PyTorch Distributed Overview <https://pytorch.org/tutorials/beginner/dist_overview.html>`__\\n        for a brief introduction to all features related to distributed training.\\n\\n    .. note::\\n        ``DistributedDataParallel`` can be used in conjunction with\\n        :class:`torch.distributed.optim.ZeroRedundancyOptimizer` to reduce\\n        per-rank optimizer states memory footprint. Please refer to\\n        `ZeroRedundancyOptimizer recipe <https://pytorch.org/tutorials/recipes/zero_redundancy_optimizer.html>`__\\n        for more details.\\n\\n    .. note:: ``nccl`` backend is currently the fastest and highly recommended\\n        backend when using GPUs. This applies to both single-node and\\n        multi-node distributed training.\\n\\n    .. note:: This module also supports mixed-precision distributed training.\\n        This means that your model can have different types of parameters such\\n        as mixed types of ``fp16`` and ``fp32``, the gradient reduction on these\\n        mixed types of parameters will just work fine.\\n\\n    .. note:: If you use ``torch.save`` on one process to checkpoint the module,\\n        and ``torch.load`` on some other processes to recover it, make sure that\\n        ``map_location`` is configured properly for every process. Without\\n        ``map_location``, ``torch.load`` would recover the module to devices\\n        where the module was saved from.\\n\\n    .. note:: When a model is trained on ``M`` nodes with ``batch=N``, the\\n        gradient will be ``M`` times smaller when compared to the same model\\n        trained on a single node with ``batch=M*N`` if the loss is summed (NOT\\n        averaged as usual) across instances in a batch (because the gradients\\n        between different nodes are averaged). You should take this into\\n        consideration when you want to obtain a mathematically equivalent\\n        training process compared to the local training counterpart. But in most\\n        cases, you can just treat a DistributedDataParallel wrapped model, a\\n        DataParallel wrapped model and an ordinary model on a single GPU as the\\n        same (E.g. using the same learning rate for equivalent batch size).\\n\\n    .. note::\\n        Parameters are never broadcast between processes. The module performs\\n        an all-reduce step on gradients and assumes that they will be modified\\n        by the optimizer in all processes in the same way. Buffers\\n        (e.g. BatchNorm stats) are broadcast from the module in process of rank\\n        0, to all other replicas in the system in every iteration.\\n\\n    .. note::\\n        If you are using DistributedDataParallel in conjunction with the\\n        :ref:`distributed-rpc-framework`, you should always use\\n        :meth:`torch.distributed.autograd.backward` to compute gradients and\\n        :class:`torch.distributed.optim.DistributedOptimizer` for optimizing\\n        parameters.\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined variables\\\")\\n            >>> import torch.distributed.autograd as dist_autograd\\n            >>> from torch.nn.parallel import DistributedDataParallel as DDP\\n            >>> import torch\\n            >>> from torch import optim\\n            >>> from torch.distributed.optim import DistributedOptimizer\\n            >>> import torch.distributed.rpc as rpc\\n            >>> from torch.distributed.rpc import RRef\\n            >>>\\n            >>> t1 = torch.rand((3, 3), requires_grad=True)\\n            >>> t2 = torch.rand((3, 3), requires_grad=True)\\n            >>> rref = rpc.remote(\\\"worker1\\\", torch.add, args=(t1, t2))\\n            >>> ddp_model = DDP(my_model)\\n            >>>\\n            >>> # Setup optimizer\\n            >>> optimizer_params = [rref]\\n            >>> for param in ddp_model.parameters():\\n            >>>     optimizer_params.append(RRef(param))\\n            >>>\\n            >>> dist_optim = DistributedOptimizer(\\n            >>>     optim.SGD,\\n            >>>     optimizer_params,\\n            >>>     lr=0.05,\\n            >>> )\\n            >>>\\n            >>> with dist_autograd.context() as context_id:\\n            >>>     pred = ddp_model(rref.to_here())\\n            >>>     loss = loss_func(pred, target)\\n            >>>     dist_autograd.backward(context_id, [loss])\\n            >>>     dist_optim.step(context_id)\\n\\n    .. note::\\n        DistributedDataParallel currently offers limited support for gradient\\n        checkpointing with :meth:`torch.utils.checkpoint`.\\n        If the checkpoint is done with use_reentrant=False (recommended), DDP\\n        will work as expected without any limitations.\\n        If, however, the checkpoint is done with use_reentrant=True (the default),\\n        DDP will work as expected when there are no unused parameters in the model\\n        and each layer is checkpointed at most once (make sure you are not passing\\n        `find_unused_parameters=True` to DDP). We currently do not support the\\n        case where a layer is checkpointed multiple times, or when there unused\\n        parameters in the checkpointed model.\\n\\n    .. note::\\n        To let a non-DDP model load a state dict from a DDP model,\\n        :meth:`~torch.nn.modules.utils.consume_prefix_in_state_dict_if_present`\\n        needs to be applied to strip the prefix \\\"module.\\\" in the DDP state dict before loading.\\n\\n    .. warning::\\n        Constructor, forward method, and differentiation of the output (or a\\n        function of the output of this module) are distributed synchronization\\n        points. Take that into account in case different processes might be\\n        executing different code.\\n\\n    .. warning::\\n        This module assumes all parameters are registered in the model by the\\n        time it is created. No parameters should be added nor removed later.\\n        Same applies to buffers.\\n\\n    .. warning::\\n        This module assumes all parameters are registered in the model of each\\n        distributed processes are in the same order. The module itself will\\n        conduct gradient ``allreduce`` following the reverse order of the\\n        registered parameters of the model. In other words, it is users'\\n        responsibility to ensure that each distributed process has the exact\\n        same model and thus the exact same parameter registration order.\\n\\n    .. warning::\\n        This module allows parameters with non-rowmajor-contiguous strides.\\n        For example, your model may contain some parameters whose\\n        :class:`torch.memory_format` is ``torch.contiguous_format``\\n        and others whose format is ``torch.channels_last``.  However,\\n        corresponding parameters in different processes must have the\\n        same strides.\\n\\n    .. warning::\\n        This module doesn't work with :func:`torch.autograd.grad` (i.e. it will\\n        only work if gradients are to be accumulated in ``.grad`` attributes of\\n        parameters).\\n\\n    .. warning::\\n        If you plan on using this module with a ``nccl`` backend or a ``gloo``\\n        backend (that uses Infiniband), together with a DataLoader that uses\\n        multiple workers, please change the multiprocessing start method to\\n        ``forkserver`` (Python 3 only) or ``spawn``. Unfortunately\\n        Gloo (that uses Infiniband) and NCCL2 are not fork safe, and you will\\n        likely experience deadlocks if you don't change this setting.\\n\\n    .. warning::\\n        You should never try to change your model's parameters after wrapping\\n        up your model with ``DistributedDataParallel``. Because, when\\n        wrapping up your model with ``DistributedDataParallel``, the constructor\\n        of ``DistributedDataParallel`` will register the additional gradient\\n        reduction functions on all the parameters of the model itself at the\\n        time of construction. If you change the model's parameters afterwards,\\n        gradient reduction functions no longer match the correct set of\\n        parameters.\\n\\n    .. warning::\\n        Using ``DistributedDataParallel`` in conjunction with the\\n        :ref:`distributed-rpc-framework` is experimental and subject to change.\\n\\n    Args:\\n        module (Module): module to be parallelized\\n        device_ids (list of int or torch.device): CUDA devices.\\n                   1) For single-device modules, ``device_ids`` can\\n                   contain exactly one device id, which represents the only\\n                   CUDA device where the input module corresponding to this process resides.\\n                   Alternatively, ``device_ids`` can also be ``None``.\\n                   2) For multi-device modules and CPU modules,\\n                   ``device_ids`` must be ``None``.\\n\\n                   When ``device_ids`` is ``None`` for both cases,\\n                   both the input data for the forward pass and the actual module\\n                   must be placed on the correct device.\\n                   (default: ``None``)\\n        output_device (int or torch.device): Device location of output for\\n                      single-device CUDA modules. For multi-device modules and\\n                      CPU modules, it must be ``None``, and the module itself\\n                      dictates the output location. (default: ``device_ids[0]``\\n                      for single-device modules)\\n        broadcast_buffers (bool): Flag that enables syncing (broadcasting)\\n                          buffers of the module at beginning of the ``forward``\\n                          function. (default: ``True``)\\n        process_group: The process group to be used for distributed data\\n                       all-reduction. If ``None``, the default process group, which\\n                       is created by :func:`torch.distributed.init_process_group`,\\n                       will be used. (default: ``None``)\\n        bucket_cap_mb: ``DistributedDataParallel`` will bucket parameters into\\n                       multiple buckets so that gradient reduction of each\\n                       bucket can potentially overlap with backward computation.\\n                       :attr:`bucket_cap_mb` controls the bucket size in\\n                       MebiBytes (MiB). If ``None``, a default size of 25 MiB\\n                       will be used. (default: ``None``)\\n        find_unused_parameters (bool): Traverse the autograd graph from all\\n                               tensors contained in the return value of the\\n                               wrapped module's ``forward`` function. Parameters\\n                               that don't receive gradients as part of this\\n                               graph are preemptively marked as being ready to\\n                               be reduced. In addition, parameters that may have\\n                               been used in the wrapped module's ``forward``\\n                               function but were not part of loss computation and\\n                               thus would also not receive gradients are\\n                               preemptively marked as ready to be reduced.\\n                               (default: ``False``)\\n        check_reduction: This argument is deprecated.\\n        gradient_as_bucket_view (bool): When set to ``True``, gradients will be views\\n                      pointing to different offsets of ``allreduce`` communication\\n                      buckets. This can reduce peak memory usage, where the\\n                      saved memory size will be equal to the total gradients\\n                      size. Moreover, it avoids the overhead of copying between\\n                      gradients and ``allreduce`` communication buckets. When\\n                      gradients are views, ``detach_()`` cannot be called on the\\n                      gradients. If hitting such errors, please fix it by\\n                      referring to the :meth:`~torch.optim.Optimizer.zero_grad`\\n                      function in ``torch/optim/optimizer.py`` as a solution.\\n                      Note that gradients will be views after first iteration, so\\n                      the peak memory saving should be checked after first iteration.\\n        static_graph (bool): When set to ``True``, DDP knows the trained graph is\\n                     static. Static graph means 1) The set of used and unused\\n                     parameters will not change during the whole training loop; in\\n                     this case, it does not matter whether users set\\n                     ``find_unused_parameters = True`` or not. 2) How the graph is trained\\n                     will not change during the whole training loop (meaning there is\\n                     no control flow depending on iterations).\\n                     When static_graph is set to be ``True``, DDP will support cases that\\n                     can not be supported in the past:\\n                     1) Reentrant backwards.\\n                     2) Activation checkpointing multiple times.\\n                     3) Activation checkpointing when model has unused parameters.\\n                     4) There are model parameters that are outside of forward function.\\n                     5) Potentially improve performance when there are unused parameters,\\n                     as DDP will not search graph in each iteration to detect unused\\n                     parameters when static_graph is set to be ``True``.\\n                     To check whether you can set static_graph to be ``True``, one way is to\\n                     check ddp logging data at the end of your previous model training,\\n                     if ``ddp_logging_data.get(\\\"can_set_static_graph\\\") == True``, mostly you\\n                     can set ``static_graph = True`` as well.\\n\\n                     Example::\\n                         >>> # xdoctest: +SKIP(\\\"undefined variables\\\")\\n                         >>> model_DDP = torch.nn.parallel.DistributedDataParallel(model)\\n                         >>> # Training loop\\n                         >>> ...\\n                         >>> ddp_logging_data = model_DDP._get_ddp_logging_data()\\n                         >>> static_graph = ddp_logging_data.get(\\\"can_set_static_graph\\\")\\n        delay_all_reduce_named_params (list of tuple of str and torch.nn.Parameter): a list\\n                    of named parameters whose all reduce will be delayed when the gradient of\\n                    the parameter specified in ``param_to_hook_all_reduce`` is ready. Other\\n                    arguments of DDP do not apply to named params specified in this argument\\n                    as these named params will be ignored by DDP reducer.\\n        param_to_hook_all_reduce (torch.nn.Parameter): a parameter to hook delayed all reduce\\n                    of parameters specified in ``delay_all_reduce_named_params``.\\n\\n\\n    Attributes:\\n        module (Module): the module to be parallelized.\\n\\n    Example::\\n\\n        >>> # xdoctest: +SKIP(\\\"undefined variables\\\")\\n        >>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')\\n        >>> net = torch.nn.parallel.DistributedDataParallel(model)\\n    \\\"\\\"\\\"\\n\\n    # used to track whether the given thread is inside ddp forward for torchdynamo purposes\\n    _active_ddp_module: Optional[\\\"DistributedDataParallel\\\"] = None\\n\\n    def __init__(\\n        self,\\n        module,\\n        device_ids=None,\\n        output_device=None,\\n        dim=0,\\n        broadcast_buffers=True,\\n        process_group=None,\\n        bucket_cap_mb=None,\\n        find_unused_parameters=False,\\n        check_reduction=False,\\n        gradient_as_bucket_view=False,\\n        static_graph=False,\\n        delay_all_reduce_named_params=None,\\n        param_to_hook_all_reduce=None,\\n        mixed_precision: Optional[_MixedPrecision] = None,\\n        device_mesh=None,\\n    ):\\n        super().__init__()\\n        Joinable.__init__(self)\\n        self.logger: Optional[dist.Logger] = None\\n        if bool(delay_all_reduce_named_params is not None) != bool(\\n            param_to_hook_all_reduce is not None\\n        ):\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"delay_all_reduce_named_params and param_to_hook_all_reduce \\\"\\n                \\\"need to be set at the same time.\\\",\\n            )\\n\\n        if process_group and device_mesh is not None:\\n            raise RuntimeError(\\n                \\\"Cannot specify both process_group and device_mesh arguments.\\\"\\n            )\\n        elif process_group is None and device_mesh is None:\\n            self.process_group = _get_default_group()\\n        elif device_mesh is None:\\n            self.process_group = process_group\\n        else:\\n            if device_mesh.ndim != 1:\\n                raise RuntimeError(\\n                    f\\\"Only 1D device mesh is supported, but got {device_mesh}.\\\"\\n                )\\n            self.device_mesh = device_mesh\\n            self.process_group = device_mesh.get_group(mesh_dim=0)\\n            from torch.distributed.device_mesh import _mesh_resources\\n\\n            root_mesh = _mesh_resources.get_root_mesh(device_mesh)\\n            # if a root mesh is not the same as device_mesh,\\n            # meaning the device_mesh is sliced out from the root mesh.\\n            if root_mesh != device_mesh:\\n                # TODO: This is a temporary work around to enable DDP + TP.\\n                # We should do the logic in DDP so that the 2D implementation is\\n                # sound and the state_dict works out of the box.\\n                # This has to be done before check UninitializedParameter.\\n                from torch.distributed.tensor.parallel.ddp import (\\n                    _pre_dp_module_transform,\\n                )\\n\\n                _pre_dp_module_transform(module)\\n\\n        self._delay_all_reduce_params = []\\n        if hasattr(module, \\\"_ddp_params_and_buffers_to_ignore\\\"):\\n            self.parameters_to_ignore = set(module._ddp_params_and_buffers_to_ignore)\\n        else:\\n            self.parameters_to_ignore = set()\\n        if delay_all_reduce_named_params is not None:\\n            for name, param in delay_all_reduce_named_params:\\n                self.parameters_to_ignore.add(name)\\n                self._delay_all_reduce_params.append(param)\\n\\n        self._module_parameters = [\\n            p\\n            for n, p in module.named_parameters()\\n            if n not in self.parameters_to_ignore\\n        ]\\n        if not any(p.requires_grad for p in self._module_parameters):\\n            if len(self._delay_all_reduce_params):\\n                logger.info(\\\"Delay the AllReduce of all parameters.\\\")\\n            else:\\n                self._log_and_throw(\\n                    RuntimeError,\\n                    \\\"DistributedDataParallel is not needed when a module \\\"\\n                    \\\"doesn't have any parameter that requires a gradient.\\\",\\n                )\\n\\n        if device_ids is not None and len(device_ids) > 1:\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"device_ids can only be None or contain a single element.\\\",\\n            )\\n\\n        self.is_multi_device_module = (\\n            len({p.device for p in self._module_parameters}) > 1\\n        )\\n        distinct_device_types = {\\n            p.device.type for p in self._module_parameters if p.device is not None\\n        }\\n        if len(distinct_device_types) != 1:\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"DistributedDataParallel's input module must be on \\\"\\n                f\\\"the same type of devices, but input module parameters locate in {distinct_device_types}.\\\",\\n            )\\n\\n        self.device_type = next(iter(distinct_device_types))\\n\\n        if (\\n            device_ids is None\\n            or len(device_ids) == 0  # For backward compatibility.\\n            or self.device_type == \\\"cpu\\\"\\n            or self.is_multi_device_module\\n        ):\\n            if device_ids or output_device:\\n                self._log_and_throw(\\n                    ValueError,\\n                    \\\"DistributedDataParallel device_ids and output_device arguments \\\"\\n                    \\\"only work with single-device/multiple-device GPU modules or CPU modules, \\\"\\n                    f\\\"but got device_ids {device_ids}, output_device {output_device}, \\\"\\n                    f\\\"and module parameters {({p.device for p in self._module_parameters})}.\\\",\\n                )\\n\\n            self.device_ids = None\\n            self.output_device = None\\n        else:\\n            self.device_ids = [_get_device_index(x, True) for x in device_ids]\\n\\n            if output_device is None:\\n                output_device = device_ids[0]\\n\\n            self.output_device = _get_device_index(output_device, True)\\n\\n        self.static_graph = False\\n        self.dim = dim\\n        self.module = module\\n        self.device = next(iter(self._module_parameters)).device\\n        self.broadcast_buffers = broadcast_buffers\\n        self.find_unused_parameters = find_unused_parameters\\n        self.require_backward_grad_sync = True\\n        self.require_forward_param_sync = True\\n        self.gradient_as_bucket_view = gradient_as_bucket_view\\n        self.mixed_precision = mixed_precision\\n        if self.mixed_precision is not None:\\n            logger.warning(\\\"Received mixed precision config %s\\\", self.mixed_precision)\\n\\n        if check_reduction:\\n            # This argument is no longer used since the reducer\\n            # will ensure reduction completes even if some parameters\\n            # do not receive gradients.\\n            warnings.warn(\\n                \\\"The `check_reduction` argument in `DistributedDataParallel` \\\"\\n                \\\"module is deprecated. Please avoid using it.\\\",\\n                FutureWarning,\\n                stacklevel=2,\\n            )\\n\\n        # Check that a module does not have Uninitialized parameters\\n        for param in self._module_parameters:\\n            if isinstance(param, torch.nn.parameter.UninitializedParameter):\\n                self._log_and_throw(\\n                    RuntimeError,\\n                    \\\"Modules with uninitialized parameters can't be used with `DistributedDataParallel`. \\\"\\n                    \\\"Run a dummy forward pass to correctly initialize the modules\\\",\\n                )\\n        # used for intra-node param sync and inter-node sync as well\\n        self.broadcast_bucket_size = int(250 * 1024 * 1024)\\n\\n        # reduction bucket size\\n        if bucket_cap_mb is None:\\n            # default case (bucket cap is 25 MiB)\\n            bucket_cap_mb = 25\\n            self.bucket_bytes_cap_default = True\\n        else:\\n            self.bucket_bytes_cap_default = False\\n        self.bucket_bytes_cap = int(bucket_cap_mb * 1024 * 1024)\\n\\n        # Whether to perform input tensor CPU to GPU copies on a side-stream\\n        self.use_side_stream_for_tensor_copies = (\\n            os.environ.get(\\\"PYTORCH_DDP_USE_SIDE_STREAM\\\", \\\"1\\\") == \\\"1\\\"\\n        )\\n\\n        # Initialize gradient buffers and register all reduce hook\\n        self._delay_grad_buffer: Optional[torch.Tensor] = None\\n        self._delay_grad_views: List[torch.Tensor] = []\\n        self._delay_all_reduce_all_params = False\\n        if len(self._delay_all_reduce_params) != 0:\\n            self._register_delay_all_reduce_hook(\\n                bucket_cap_mb=bucket_cap_mb,\\n                param_to_hook_all_reduce=param_to_hook_all_reduce,\\n                device_ids=device_ids,\\n            )\\n            if self._delay_all_reduce_all_params:\\n                return\\n\\n        # Build parameters for reducer.\\n        parameters, expect_sparse_gradient = self._build_params_for_reducer()\\n        # Verify model equivalence.\\n        _verify_param_shape_across_processes(self.process_group, parameters)\\n        # Sync params and buffers. Ensures all DDP models start off at the same value.\\n        _sync_module_states(\\n            module=self.module,\\n            process_group=self.process_group,\\n            broadcast_bucket_size=self.broadcast_bucket_size,\\n            src=0,\\n            params_and_buffers_to_ignore=self.parameters_to_ignore,\\n            broadcast_buffers=self.broadcast_buffers,\\n        )\\n        # In debug mode, build a mapping of parameter index -> parameter.\\n        param_to_name_mapping = self._build_debug_param_to_name_mapping(parameters)\\n\\n        # Builds reducer.\\n        self._ddp_init_helper(\\n            parameters,\\n            expect_sparse_gradient,\\n            param_to_name_mapping,\\n            static_graph,\\n        )\\n        self._comm_hooks: List[Tuple[Callable, object]] = []\\n\\n        if self.mixed_precision is not None:\\n            _setup_mixed_precision_params(self.mixed_precision, self.module)\\n            _cast_buffers(self.mixed_precision, self.module)\\n            # Stream used for async low precision copies.\\n            self._mp_stream = torch.cuda.Stream()\\n            self._submodule_to_event = defaultdict(deque)  # type: ignore[var-annotated]\\n            # Add forward pre-hook to root module to kick off copies to lower\\n            # precision.\\n            self.module.register_forward_pre_hook(\\n                self._root_copy_hook, prepend=False, with_kwargs=True\\n            )\\n            # Add forward pre hook to all submodules to wait for copy events\\n            # before running computation.\\n            for module in self.module.modules():\\n                module.register_forward_pre_hook(\\n                    self._module_wait_for_copy_hook,\\n                    prepend=False,\\n                    with_kwargs=True,\\n                )\\n            # Set up callbacks in backward to upcast and use full precision\\n            # params. TODO (rohan-varma): Make this compose with general\\n            # comm hooks and apply_optimizer_in_backward. Importing inline to\\n            # avoid circular import issue.\\n            from torch.distributed.algorithms.ddp_comm_hooks.mixed_precision_hooks import (\\n                _AllreduceUpcastHookState,\\n                _reducer_allreduce_and_upcast_hook,\\n            )\\n\\n            upcast_hook_state = _AllreduceUpcastHookState(\\n                ddp_weakref=weakref.ref(self),\\n                upcast_stream=torch.cuda.Stream(),\\n            )\\n            self.register_comm_hook(\\n                upcast_hook_state,\\n                _reducer_allreduce_and_upcast_hook,\\n            )\\n            # Inform reducer of reduced precision param dtype for correctness\\n            # of type checks between gradient and bucket.\\n            self.reducer._set_mixed_precision_param_dtype(  # type: ignore[attr-defined]\\n                self.mixed_precision.param_dtype\\n            )\\n\\n        self._has_rebuilt_buckets = False\\n\\n        if static_graph:\\n            self._set_static_graph()\\n\\n        self._lazy_init_ran = False\\n\\n        # Register the AccumulateGrad post hooks if optimize_ddp is\\n        # True. The hooks will be deregistered if compiled_autograd is not\\n        # enabled.\\n        self._accum_grad_hooks: List[RemovableHandle] = []\\n        optimize_ddp = torch._dynamo.config._get_optimize_ddp_mode()\\n        self._use_python_reducer = optimize_ddp in (\\n            \\\"python_reducer\\\",\\n            \\\"python_reducer_without_compiled_forward\\\",\\n        )\\n        if self._use_python_reducer:\\n            torch._inductor.config._fuse_ddp_communication = True\\n            torch._inductor.config._fuse_ddp_bucket_size = bucket_cap_mb\\n            # Directly adding this to the trace rule will disturb the users\\n            # who are using DDPOptimizer.\\n            torch._dynamo.trace_rules.LEGACY_MOD_INLINELIST.add(\\n                \\\"torch.nn.parallel.distributed\\\"\\n            )\\n            torch._dynamo.trace_rules.get_legacy_mod_inlinelist.cache_clear()\\n        self._force_to_disable_cpp_reducer = (\\n            optimize_ddp == \\\"python_reducer_without_compiled_forward\\\"\\n        )\\n        if self._use_python_reducer:\\n            self._register_accum_grad_hook()\\n\\n        # Whether or not DDPSink performs a clone.\\n        self._ddp_sink_clone = True\\n\\n    def _register_accum_grad_hook(self):\\n        import torch.distributed._functional_collectives as fcol\\n\\n        def compiled_accum_grad_hook(\\n            param,\\n            *,\\n            param_index: int,\\n        ):\\n            if not self.require_backward_grad_sync:\\n                return\\n\\n            if param.grad is None:\\n                return\\n\\n            if self._comm_hooks:\\n                for hook, state in self._comm_hooks:\\n                    hook(state, (param.grad, param))\\n            else:\\n                gradient = param.grad / self.process_group.size()\\n                gradient = fcol.all_reduce(gradient, \\\"sum\\\", self.process_group)\\n                param.grad.copy_(gradient)\\n\\n        for index, param in enumerate(self._module_parameters):\\n            if not param.requires_grad:\\n                continue\\n            self._accum_grad_hooks.append(\\n                param.register_post_accumulate_grad_hook(\\n                    functools.partial(\\n                        compiled_accum_grad_hook,\\n                        param_index=index,\\n                    )\\n                )\\n            )\\n\\n    def _delayed_all_reduce_hook(self, grad):\\n        world_size = dist.get_world_size(self.process_group)\\n\\n        self._delay_grad_buffer.div_(world_size)  # type: ignore[union-attr]\\n        _ = dist.all_reduce(\\n            self._delay_grad_buffer, group=self.process_group, async_op=True\\n        )\\n        return grad\\n\\n    def _register_delay_all_reduce_hook(\\n        self,\\n        bucket_cap_mb,\\n        param_to_hook_all_reduce,\\n        device_ids,\\n    ):\\n        # 1. Create gradient buffer\\n        device = torch.device(\\\"cpu\\\") if device_ids is None else device_ids[0]\\n        self._delay_grad_buffer = torch.zeros(\\n            sum(p.numel() for p in self._delay_all_reduce_params),\\n            device=device,\\n        )\\n\\n        # 2. Broadcast the parameters\\n        detached_params = [p.detach() for p in self._delay_all_reduce_params]\\n        dist._broadcast_coalesced(self.process_group, detached_params, bucket_cap_mb, 0)\\n\\n        # 3. Hook all reduce to the specified parameter\\n        param_to_hook_all_reduce.register_hook(self._delayed_all_reduce_hook)\\n\\n        # 4. Build tensor views for gradients\\n        offset = 0\\n        for param in self._delay_all_reduce_params:\\n            grad_view = self._delay_grad_buffer[offset : (offset + param.numel())].view(\\n                param.shape\\n            )\\n            self._delay_grad_views.append(grad_view)\\n            offset = offset + param.numel()\\n\\n        # 5. Check whether the all reduce of all params requiring grad is delayed.\\n        for module_name, module in self.module.named_modules():\\n            for param_name, param in module.named_parameters(recurse=False):\\n                if param.requires_grad:\\n                    full_name = f\\\"{module_name}.{param_name}\\\"\\n                    if full_name not in self.parameters_to_ignore:\\n                        # There is at least a param whose all reduce will not be delayed.\\n                        # In this case, we should not set self._delay_all_reduce_all_params\\n                        # to True.\\n                        return\\n        self._delay_all_reduce_all_params = True\\n\\n    def _setup_in_backward_optimizers(self):\\n        # Check if user has used apply_optim_in_backward to overlap optimizer\\n        # step + DDP backward. Current constraints:\\n        # 1. Only allreduce is supported at the moment, no custom communication.\\n        # 2. For DDP-managed parameters that have their optimizer run in\\n        # backward, their gradients are set to ``None``. If your use case\\n        # requires DDP parameters grad not to be set to ``None`` after their\\n        # in-backward optimizer runs, please ping\\n        # https://github.com/pytorch/pytorch/issues/90052.\\n        # NOTE: we use self._module_parameters instead of .parameters() since\\n        # the former excludes ignored (non-DDP managed) parameters.\\n        if any(hasattr(p, \\\"_in_backward_optimizers\\\") for p in self._module_parameters):\\n            torch._C._log_api_usage_once(\\\"ddp.optimizer_in_backward\\\")\\n            # Remove hooks that apply_optim_in_backward had registered because\\n            # DDP customizes how optimizer is overlapped with backward due to\\n            # the allreduce.\\n            param_to_handle_map = (\\n                dist.optim.apply_optimizer_in_backward.param_to_optim_hook_handle_map\\n            )\\n            for p in self._module_parameters:\\n                for handle in param_to_handle_map.get(p, []):\\n                    handle.remove()\\n\\n            # Need a weakref to DDP instance to run all_reduce (from reducer)\\n            # and get managed DDP parameters.\\n            ddp_weakref = weakref.ref(self)\\n            # Note: importing in function, otherwise this will cause a circular\\n            # import.\\n            from torch.distributed.algorithms.ddp_comm_hooks.optimizer_overlap_hooks import (\\n                _apply_optim_in_backward_hook,\\n            )\\n\\n            self.register_comm_hook(\\n                ddp_weakref,\\n                _apply_optim_in_backward_hook(\\n                    gradient_is_bucket_view=self.gradient_as_bucket_view\\n                ),\\n            )\\n\\n            self.reducer._set_optimizer_in_backward()  # type: ignore[attr-defined]\\n\\n    def _fire_reducer_autograd_hook(self, idx, *unused):\\n        \\\"\\\"\\\"\\n        Fire the reducer's autograd hook to allreduce params in a Reducer bucket.\\n\\n        Note that this is only used during mixed precision training as the\\n        Reducer's hooks installed during construction time would not be called\\n        as we're working in the low precision parameter setting.\\n        \\\"\\\"\\\"\\n        self.reducer._autograd_hook(idx)  # type: ignore[attr-defined]\\n\\n    def _root_copy_hook(self, *args: Any, **kwargs: Any) -> None:\\n        \\\"\\\"\\\"\\n        For DDP mixed precision, put low precision copies on separate stream and create events to wait for them.\\n\\n        When training with DDP mixed precision, this root pre-forward hook kicks\\n        off low precision copies on a separate stream and creates respective\\n        events to wait for them.\\n        \\\"\\\"\\\"\\n        # Clear out previous iteration submodule to event. This is because we\\n        # may have populated some events for modules that didn't end up being\\n        # used.\\n        self._submodule_to_event = defaultdict(deque)  # type: ignore[var-annotated]\\n        with torch.cuda.stream(self._mp_stream):\\n            for submodule in self.module.modules():\\n                for param in submodule.parameters(recurse=False):\\n                    # Do not cast DDP ignored parameters.\\n                    if hasattr(param, \\\"_ddp_ignored\\\") and param._ddp_ignored:\\n                        continue\\n                    _alloc_storage(param._mp_param, param.size())\\n                    # copy() implicitly casts to low precision\\n                    with torch.no_grad():\\n                        param._mp_param.copy_(param.data)\\n                        # TODO: when zero_grad(set_to_none=False) or in grad\\n                        # accumulation case, accumulated grads can be in fp32\\n                        # which can cause errors when running DDP backwards due\\n                        # to mismatched incoming and accumulated gradient types.\\n                        # So we manually cast the accumulated grad down for now,\\n                        # in the future we may shift to FSDP style gradient\\n                        # accumulation management where the accumulated gradient\\n                        # is saved and .grad field is set to None, bypassing\\n                        # this issue.\\n                        if param.grad is not None:\\n                            param.grad.data = param.grad.to(\\n                                self.mixed_precision.param_dtype  # type: ignore[union-attr]\\n                            )\\n                    param.data = param._mp_param\\n                copy_event = torch.cuda.Event()\\n                copy_event.record()\\n                self._submodule_to_event[submodule].append(copy_event)\\n\\n    def _module_wait_for_copy_hook(\\n        self,\\n        module,\\n        *args: Any,\\n        **kwargs: Any,\\n    ) -> None:\\n        \\\"\\\"\\\"Before carrying out computation, wait on the appropriate event to ensure low precision copies have finished.\\\"\\\"\\\"\\n        try:\\n            event = self._submodule_to_event[module].popleft()\\n        except IndexError:\\n            # copy event has already been waited on\\n            return\\n\\n        event.wait(stream=torch.cuda.current_stream())\\n        for p in module.parameters(recurse=False):\\n            # Don't register hooks if param does not require grad\\n            if not p.requires_grad or (hasattr(p, \\\"_ddp_ignored\\\") and p._ddp_ignored):\\n                continue\\n            # We need to register autograd hook here instead of DDP's ctor\\n            # since we're working with the low precision param. Register them\\n            # via obtaining the gradient accumulator.\\n            tmp = p.expand_as(p)\\n            grad_acc = tmp.grad_fn.next_functions[0][0]\\n\\n            hook = grad_acc.register_hook(\\n                functools.partial(self._fire_reducer_autograd_hook, p._idx)\\n            )\\n            p._ddp_mp_hook_state = (grad_acc, hook)\\n\\n    def _log_and_throw(self, err_type, err_msg):\\n        if self.logger is not None:\\n            self.logger.set_error_and_log(f\\\"{str(err_type)}: {err_msg}\\\")\\n        raise err_type(err_msg)\\n\\n    def _ddp_init_helper(\\n        self,\\n        parameters,\\n        expect_sparse_gradient,\\n        param_to_name_mapping,\\n        static_graph,\\n    ):\\n        \\\"\\\"\\\"\\n        DDP init helper function to manage parameters, grad hooks, logging, and SyncBatchNorm.\\n\\n        Initialization helper function that does the following:\\n        (1) bucketing the parameters for reductions\\n        (2) resetting the bucketing states\\n        (3) registering the grad hooks\\n        (4) Logging construction-time DDP logging data\\n        (5) passing a handle of DDP to SyncBatchNorm Layer\\n        \\\"\\\"\\\"\\n        # Notice, the parameters order is not in the order in which they are used,\\n        # especially in models with control flow.\\n        #\\n        # Alongside parameters are not presented in the real execution order,\\n        # if a certain model happens to also\\n        #   1) have other collectives comm ops in its backward graph.\\n        #   2) have unused parameter in subset ranks of the whole world.\\n        # bucketing could insert ALL-REDUCE comm op too early on the rank with unused parameter,\\n        # matching up with other collectives comm ops on other ranks unexpectedly.\\n        #\\n        # In order to handle this corner case, when the parameters are not in the real execution order,\\n        # we don't do bucketing, thus only one ALL-REDUCE is inserted after all the gradients\\n        # of the whole graph are computed.\\n        #\\n        # Notice, here we only disable bucketing for the first iteration.\\n        # After the first iteration, it's OK to rebuild buckets,\\n        # because \\\"bucket rebuild\\\" bucketizes parameters based on its real execution order in backward graph.\\n\\n        # Can remove this branching once #73732 is landed.\\n        if static_graph is True or self.find_unused_parameters is False:\\n            bucket_size_limits = [sys.maxsize]\\n        else:\\n            if self.bucket_bytes_cap_default:\\n                bucket_size_limits = [\\n                    dist._DEFAULT_FIRST_BUCKET_BYTES,\\n                    self.bucket_bytes_cap,\\n                ]\\n            else:\\n                bucket_size_limits = [self.bucket_bytes_cap]\\n        (\\n            bucket_indices,\\n            per_bucket_size_limits,\\n        ) = dist._compute_bucket_assignment_by_size(\\n            parameters,\\n            bucket_size_limits,\\n            expect_sparse_gradient,\\n        )\\n\\n        # Remember index for parameters if we are in mixed precision, as we\\n        # need to pass in index to Reducer's autograd hook via python.\\n        if self.mixed_precision is not None:\\n            for i, p in enumerate(parameters):\\n                p._idx = i\\n\\n        # Note: reverse list of buckets because we want to approximate the\\n        # order in which their gradients are produced, and assume they\\n        # are used in the forward pass in the order they are defined.\\n        self.reducer = dist.Reducer(\\n            parameters,\\n            list(reversed(bucket_indices)),\\n            list(reversed(per_bucket_size_limits)),\\n            self.process_group,\\n            expect_sparse_gradient,\\n            # The bucket size limit is specified in the constructor.\\n            # Additionally, we allow for a single small bucket for parameters\\n            # that are defined first, such that their gradients don't spill into\\n            # a much larger bucket, adding unnecessary latency after gradient\\n            # computation finishes. Experiments showed 1MB is a reasonable value.\\n            self.bucket_bytes_cap,\\n            self.find_unused_parameters,\\n            self.gradient_as_bucket_view,\\n            param_to_name_mapping,\\n            # User can set dist._DEFAULT_FIRST_BUCKET_BYTES to tune DDP first\\n            # bucket.\\n            (\\n                dist._DEFAULT_FIRST_BUCKET_BYTES\\n                if self.bucket_bytes_cap_default\\n                else self.bucket_bytes_cap\\n            ),\\n        )\\n\\n        self.logger = dist.Logger(self.reducer)\\n        # Set as a weak reference to avoid reference cycle between\\n        # logger and reducer.\\n        self.reducer.set_logger(self.logger)\\n\\n        has_sync_bn = False\\n        for submodule in self.module.modules():\\n            if isinstance(submodule, torch.nn.SyncBatchNorm):\\n                has_sync_bn = True\\n                break\\n\\n        # Set logging data that can be got during construction time.\\n        self.logger.set_construction_data_and_log(\\n            self.module.__class__.__name__,\\n            [] if self.device_ids is None else self.device_ids,\\n            -1 if self.output_device is None else self.output_device,\\n            self.broadcast_buffers,\\n            has_sync_bn,\\n            static_graph,\\n        )\\n\\n        # passing a handle to torch.nn.SyncBatchNorm layer\\n        self._passing_sync_batchnorm_handle(self.module)\\n\\n    def __getstate__(self):\\n        self._check_default_group()\\n        attrs = copy.copy(self.__dict__)\\n        del attrs[\\\"process_group\\\"]\\n        del attrs[\\\"reducer\\\"]\\n        del attrs[\\\"logger\\\"]\\n        return attrs\\n\\n    def __setstate__(self, state):\\n        # If serializable, then the process group should be the default one\\n        self.process_group = _get_default_group()\\n        super().__setstate__(state)\\n        self.__dict__.setdefault(\\\"require_forward_param_sync\\\", True)\\n        self.__dict__.setdefault(\\\"require_backward_grad_sync\\\", True)\\n        parameters, expect_sparse_gradient = self._build_params_for_reducer()\\n        # In debug mode, build a mapping of parameter index -> parameter.\\n        param_to_name_mapping = self._build_debug_param_to_name_mapping(parameters)\\n        # Builds reducer.\\n        self._ddp_init_helper(\\n            parameters,\\n            expect_sparse_gradient,\\n            param_to_name_mapping,\\n            self.static_graph,\\n        )\\n        if self.static_graph:\\n            self.reducer._set_static_graph()\\n            assert self.logger is not None\\n            self.logger._set_static_graph()\\n\\n    def _build_params_for_reducer(self):\\n        # Build tuple of (module, parameter) for all parameters that require grads.\\n        modules_and_parameters = [\\n            (module, parameter)\\n            for module_name, module in self.module.named_modules()\\n            for parameter in [\\n                param\\n                # Note that we access module.named_parameters instead of\\n                # parameters(module). parameters(module) is only needed in the\\n                # single-process multi device case, where it accesses replicated\\n                # parameters through _former_parameters.\\n                for param_name, param in module.named_parameters(recurse=False)\\n                if param.requires_grad\\n                and f\\\"{module_name}.{param_name}\\\" not in self.parameters_to_ignore\\n            ]\\n        ]\\n\\n        # Deduplicate any parameters that might be shared across child modules.\\n        memo = set()\\n        modules_and_parameters = [\\n            # \\\"p not in memo\\\" is the deduplication check.\\n            # \\\"not memo.add(p)\\\" is always True, and it's only there to cause \\\"add(p)\\\" if needed.\\n            (m, p)\\n            for m, p in modules_and_parameters\\n            if p not in memo and not memo.add(p)  # type: ignore[func-returns-value]\\n        ]\\n\\n        # Build list of parameters.\\n        parameters = [parameter for _, parameter in modules_and_parameters]\\n\\n        # Checks if a module will produce a sparse gradient.\\n        def produces_sparse_gradient(module):\\n            if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)):\\n                return module.sparse\\n            return False\\n\\n        # Build list of booleans indicating whether or not to expect sparse\\n        # gradients for the corresponding parameters.\\n        expect_sparse_gradient = [\\n            produces_sparse_gradient(module) for module, _ in modules_and_parameters\\n        ]\\n\\n        self._assign_modules_buffers()\\n\\n        return parameters, expect_sparse_gradient\\n\\n    def _assign_modules_buffers(self):\\n        \\\"\\\"\\\"\\n        Assign self.module.named_buffers to self.modules_buffers.\\n\\n        Assigns module buffers to self.modules_buffers which are then used to\\n        broadcast across ranks when broadcast_buffers=True. Note that this\\n        must be called every time buffers need to be synced because buffers can\\n        be reassigned by user module,\\n        see https://github.com/pytorch/pytorch/issues/63916.\\n        \\\"\\\"\\\"\\n        # Collect buffers for modules, filtering out buffers that should be ignored.\\n        named_module_buffers = [\\n            (buffer, buffer_name)\\n            for buffer_name, buffer in self.module.named_buffers()\\n            if buffer_name not in self.parameters_to_ignore\\n        ]\\n        self.modules_buffers = [\\n            buffer for (buffer, buffer_name) in named_module_buffers\\n        ]\\n        # Dict[str, tensor] representing module buffers not ignored by DDP.\\n        self.named_module_buffers = {\\n            buffer_name: buffer for (buffer, buffer_name) in named_module_buffers\\n        }\\n\\n    def _build_debug_param_to_name_mapping(self, parameters):\\n        param_to_param_index = {parameters[i]: i for i in range(len(parameters))}\\n        param_set = set(parameters)\\n        param_index_to_param_fqn = {}\\n        for module_name, module in self.module.named_modules():\\n            for param_name, param in module.named_parameters(recurse=False):\\n                fqn = f\\\"{module_name}.{param_name}\\\"\\n                # Bypass ignored parameters since those are not reduced by DDP\\n                # to begin with.\\n                if fqn not in self.parameters_to_ignore and param.requires_grad:\\n                    if param not in param_set:\\n                        self._log_and_throw(\\n                            ValueError,\\n                            f\\\"Param with name {fqn} found in module parameters, but not DDP parameters.\\\"\\n                            \\\" This indicates a bug in DDP, please report an issue to PyTorch.\\\",\\n                        )\\n                    param_index = param_to_param_index[param]\\n                    param_index_to_param_fqn[param_index] = fqn\\n\\n        # Ensure we covered all parameters\\n        if len(param_set) != len(param_index_to_param_fqn):\\n            self._log_and_throw(\\n                ValueError,\\n                (\\n                    \\\"Expected param to name mapping to cover all parameters, but\\\"\\n                    f\\\" got conflicting lengths: {len(param_set)} vs \\\"\\n                    f\\\"{len(param_index_to_param_fqn)}. This indicates a bug in DDP\\\"\\n                    \\\", please report an issue to PyTorch.\\\"\\n                ),\\n            )\\n\\n        return param_index_to_param_fqn\\n\\n    def _get_parameters(self, m, recurse=True):\\n        \\\"\\\"\\\"Return a generator of module parameters.\\\"\\\"\\\"\\n\\n        def model_parameters(m):\\n            ps = (\\n                m._former_parameters.values()\\n                if hasattr(m, \\\"_former_parameters\\\")\\n                else m.parameters(recurse=False)\\n            )\\n            yield from ps\\n\\n        for mod in m.modules() if recurse else [m]:\\n            yield from model_parameters(mod)\\n\\n    def _check_default_group(self):\\n        pickle_not_supported = False\\n        try:\\n            if self.process_group != _get_default_group():\\n                pickle_not_supported = True\\n        except RuntimeError:\\n            pickle_not_supported = True\\n\\n        if pickle_not_supported:\\n            self._log_and_throw(\\n                RuntimeError,\\n                \\\"DDP Pickling/Unpickling are only supported \\\"\\n                \\\"when using DDP with the default process \\\"\\n                \\\"group. That is, when you have called \\\"\\n                \\\"init_process_group and have not passed \\\"\\n                \\\"process_group argument to DDP constructor\\\",\\n            )\\n\\n    @contextmanager\\n    def no_sync(self):\\n        r\\\"\\\"\\\"\\n        Context manager to disable gradient synchronizations across DDP processes.\\n\\n        Within this context, gradients will be accumulated on module\\n        variables, which will later be synchronized in the first\\n        forward-backward pass exiting the context.\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"undefined variables\\\")\\n            >>> ddp = torch.nn.parallel.DistributedDataParallel(model, pg)\\n            >>> with ddp.no_sync():\\n            >>>     for input in inputs:\\n            >>>         ddp(input).backward()  # no synchronization, accumulate grads\\n            >>> ddp(another_input).backward()  # synchronize grads\\n\\n        .. warning::\\n            The forward pass should be included inside the context manager, or\\n            else gradients will still be synchronized.\\n        \\\"\\\"\\\"\\n        old_require_backward_grad_sync = self.require_backward_grad_sync\\n        self.require_backward_grad_sync = False\\n        try:\\n            yield\\n        finally:\\n            self.require_backward_grad_sync = old_require_backward_grad_sync\\n\\n    @classmethod\\n    def _get_active_ddp_module(cls):\\n        \\\"\\\"\\\"`TorchDynamo` requires DDP's status and module for cooperative optimization.\\\"\\\"\\\"\\n        return cls._active_ddp_module\\n\\n    # note, this ctxmgr function is marked 'skip' in torchdynamo, so dynamo only kicks in\\n    # for the 'module_to_run' underneath\\n    # see torch._dynamo/eval_frame.py TorchPatcher.patch for more details\\n    @contextmanager\\n    @torch._disable_dynamo(recursive=False)\\n    def _inside_ddp_forward(self):\\n        DistributedDataParallel._active_ddp_module = self\\n        try:\\n            yield\\n        finally:\\n            DistributedDataParallel._active_ddp_module = None\\n\\n    def _run_ddp_forward(self, *inputs, **kwargs):\\n        if self._use_python_reducer:\\n            return self.module(*inputs, **kwargs)  # type: ignore[index]\\n        else:\\n            with self._inside_ddp_forward():\\n                return self.module(*inputs, **kwargs)  # type: ignore[index]\\n\\n    def _clear_grad_buffer(self):\\n        # Making param.grad points to the grad buffers before backward is based on the\\n        # assumption that the grad accumulation is done in place in autograd engine,\\n        # for some edge cases, if the grad accumulation in autograd engine is not in\\n        # place, then the param.grad and grad buffers are detached.\\n        if self._delay_grad_buffer is not None:\\n            # We batch zero_grad for all params by resetting the whole grad\\n            # buffer when the grad of all params is set to None.\\n            all_param_grad_none = all(\\n                param.grad is None for param in self._delay_all_reduce_params\\n            )\\n\\n            for index, param in enumerate(self._delay_all_reduce_params):\\n                if param.grad is None:\\n                    param.grad = self._delay_grad_views[index]\\n                    if not all_param_grad_none:\\n                        param.grad.zero_()\\n\\n            if all_param_grad_none:\\n                self._delay_grad_buffer.zero_()\\n\\n    def _lazy_init(self):\\n        # Initialization for DDP that occurs after construction, but lazily\\n        # before the first forward pass.\\n        self._setup_in_backward_optimizers()\\n        self._lazy_init_ran = True\\n\\n    def _should_disable_cpp_reducer(self) -> bool:\\n        return self._use_python_reducer and (\\n            torch._utils.is_compiling() or self._force_to_disable_cpp_reducer\\n        )\\n\\n    def _pre_forward(self, *inputs, **kwargs):\\n        if self._should_disable_cpp_reducer():\\n            return inputs, kwargs\\n\\n        # Disable the python reducer if compiled_autograd is not enabled.\\n        if self._accum_grad_hooks:\\n            for index, h in enumerate(self._accum_grad_hooks):\\n                h.remove()\\n            self._accum_grad_hooks.clear()\\n\\n        if not self._lazy_init_ran and not torch._utils.is_compiling():\\n            self._lazy_init()\\n\\n        if self._delay_all_reduce_all_params:\\n            return inputs, kwargs\\n\\n        if torch.is_grad_enabled() and self.require_backward_grad_sync:\\n            assert self.logger is not None\\n            self.logger.set_runtime_stats_and_log()\\n            self.reducer.prepare_for_forward()\\n\\n        # Notify the join context that this process has not joined, if\\n        # needed\\n        work = Join.notify_join_context(self)\\n        if work:\\n            self.reducer._set_forward_pass_work_handle(\\n                work, self._divide_by_initial_world_size  # type: ignore[arg-type]\\n            )\\n\\n        # Calling _rebuild_buckets before forward computation,\\n        # It may allocate new buckets before deallocating old buckets\\n        # inside _rebuild_buckets. To save peak memory usage,\\n        # call _rebuild_buckets before the peak memory usage increases\\n        # during forward computation.\\n        # This should be called only once during whole training period.\\n        if torch.is_grad_enabled() and self.reducer._rebuild_buckets():\\n            logger.info(\\\"Reducer buckets have been rebuilt in this iteration.\\\")\\n            self._has_rebuilt_buckets = True\\n\\n        # sync params according to location (before/after forward) user\\n        # specified as part of hook, if hook was specified.\\n        if self._check_sync_bufs_pre_fwd():\\n            self._sync_buffers()\\n\\n        if self._join_config.enable:\\n            # Notify joined ranks whether they should sync in backwards pass or not.\\n            self._check_global_requires_backward_grad_sync(is_joined_rank=False)\\n\\n        if self.device_ids:\\n            moved_inputs, moved_kwargs = _to_kwargs(\\n                inputs,\\n                kwargs,\\n                torch.device(self.device_type, self.device_ids[0]),\\n                self.use_side_stream_for_tensor_copies,\\n            )\\n            args, kwargs = moved_inputs[0], moved_kwargs[0]\\n            # Cast inputs to reduced precision if needed.\\n            if self.mixed_precision is not None:\\n                args, kwargs = _cast_forward_inputs(\\n                    self.mixed_precision.param_dtype,\\n                    *args,\\n                    **kwargs,\\n                )\\n            return args, kwargs\\n        else:\\n            # Cast inputs to reduced precision if needed.\\n            # TODO (rohan-varma) test this codepath.\\n            if self.mixed_precision is not None:\\n                inputs, kwargs = _cast_forward_inputs(\\n                    self.mixed_precision.param_dtype,\\n                    *inputs,\\n                    **kwargs,\\n                )\\n            return inputs, kwargs\\n\\n    def _post_forward(self, output):\\n        if self._should_disable_cpp_reducer():\\n            return output\\n\\n        if self._delay_all_reduce_all_params:\\n            self._clear_grad_buffer()\\n            return output\\n\\n        # sync params according to location (before/after forward) user\\n        # specified as part of hook, if hook was specified.\\n        if self._check_sync_bufs_post_fwd():\\n            self._sync_buffers()\\n\\n        if torch.is_grad_enabled() and self.require_backward_grad_sync:\\n            self.require_forward_param_sync = True\\n            # We'll return the output object verbatim since it is a freeform\\n            # object. We need to find any tensors in this object, though,\\n            # because we need to figure out which parameters were used during\\n            # this forward pass, to ensure we short circuit reduction for any\\n            # unused parameters. Only if `find_unused_parameters` is set.\\n            if self.find_unused_parameters and not self.static_graph:\\n                # Do not need to populate this for static graph.\\n                self.reducer.prepare_for_backward(list(_find_tensors(output)))\\n            else:\\n                self.reducer.prepare_for_backward([])\\n        else:\\n            self.require_forward_param_sync = False\\n\\n        # TODO: DDPSink is currently enabled for unused parameter detection and\\n        # static graph training for first iteration.\\n        if (self.find_unused_parameters and not self.static_graph) or (\\n            self.static_graph and not self._static_graph_delay_allreduce_enqueued\\n        ):\\n            (\\n                output_tensor_list,\\n                treespec,\\n                output_is_rref,\\n            ) = _tree_flatten_with_rref(output)\\n            output_placeholders: List[Optional[torch.Tensor]] = [\\n                None for _ in range(len(output_tensor_list))\\n            ]\\n            # Do not touch tensors that have no grad_fn, which can cause issues\\n            # such as https://github.com/pytorch/pytorch/issues/60733\\n            for i, output in enumerate(output_tensor_list):\\n                if torch.is_tensor(output) and output.grad_fn is None:\\n                    output_placeholders[i] = output\\n\\n            # When find_unused_parameters=True, makes tensors which require grad\\n            # run through the DDPSink backward pass. When not all outputs are\\n            # used in loss, this makes those corresponding tensors receive\\n            # undefined gradient which the reducer then handles to ensure\\n            # param.grad field is not touched and we don't error out.\\n            passthrough_tensor_list = _DDPSink.apply(\\n                weakref.ref(self),\\n                *output_tensor_list,\\n            )\\n            for i in range(len(output_placeholders)):\\n                if output_placeholders[i] is None:\\n                    output_placeholders[i] = passthrough_tensor_list[i]\\n\\n            # Reconstruct output data structure.\\n            output = _tree_unflatten_with_rref(\\n                output_placeholders, treespec, output_is_rref\\n            )\\n\\n        # At the end of the forward pass, reset the grad buffer and grad views\\n        self._clear_grad_buffer()\\n        return output\\n\\n    def forward(self, *inputs, **kwargs):\\n        with torch.autograd.profiler.record_function(\\\"DistributedDataParallel.forward\\\"):\\n            inputs, kwargs = self._pre_forward(*inputs, **kwargs)\\n            output = (\\n                self.module.forward(*inputs, **kwargs)\\n                if self._delay_all_reduce_all_params\\n                else self._run_ddp_forward(*inputs, **kwargs)\\n            )\\n            return self._post_forward(output)\\n\\n    def scatter(self, inputs, kwargs, device_ids):\\n        return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)\\n\\n    def to_kwargs(self, inputs, kwargs, device_id):\\n        # Kept for BC\\n        return _to_kwargs(\\n            inputs,\\n            kwargs,\\n            torch.device(self.device_type, device_id),\\n            self.use_side_stream_for_tensor_copies,\\n        )\\n\\n    def gather(self, outputs, output_device):\\n        return gather(outputs, output_device, dim=self.dim)\\n\\n    def train(self, mode=True):\\n        super().train(mode)\\n        return self\\n\\n    # When running in join mode, schedules an allreduce to notify joined ranks\\n    # of whether backwards pass synchronization will run this iteration or not.\\n    def _check_global_requires_backward_grad_sync(self, is_joined_rank):\\n        if not is_joined_rank and self.require_backward_grad_sync:\\n            requires_sync_tensor = torch.ones(1, device=self.device)\\n        else:\\n            requires_sync_tensor = torch.zeros(1, device=self.device)\\n\\n        work = dist.all_reduce(\\n            requires_sync_tensor, group=self.process_group, async_op=True\\n        )\\n\\n        # (kwen2501) This if condition is a plain translation of previous\\n        # behavior, i.e. in the `is_joined_rank=False` case, `work.wait()`\\n        # is not called and it doesn't care about the result. I am guessing\\n        # that it just wants to fire a matching all-reduce and does not want\\n        # the main stream to wait.\\n        if is_joined_rank:\\n            work.wait()\\n            should_sync_backwards = requires_sync_tensor.item() != 0\\n            return should_sync_backwards\\n        else:\\n            return None  # Return value is not/should not be used.\\n\\n    # When running in join mode, checks and performs sync of module buffers if\\n    # the models have buffers that should be synchronized in the forward pass.\\n    def _check_and_sync_module_buffers(self):\\n        if self._check_sync_bufs_pre_fwd():\\n            authoritative_rank = self._find_common_rank(self._distributed_rank, False)\\n            self._sync_module_buffers(authoritative_rank)\\n\\n    # When running in join model, agrees upon a common rank and broadcast model\\n    # parameters to all other ranks.\\n    def _sync_final_model(self, is_last_joiner):\\n        # Agree upon the process that will be the authoritative model copy.\\n        # The current rank is a candidate for being the authoritative copy if\\n        # is_last_joiner=True. We break ties via picking the larger rank.\\n        self._authoritative_rank = self._find_common_rank(\\n            self._distributed_rank, is_last_joiner\\n        )\\n        _sync_module_states(\\n            module=self.module,\\n            process_group=self.process_group,\\n            broadcast_bucket_size=self.broadcast_bucket_size,\\n            src=self._authoritative_rank,\\n            params_and_buffers_to_ignore=self.parameters_to_ignore,\\n            broadcast_buffers=self.broadcast_buffers,\\n        )\\n\\n    # Schedule comm ops to match those scheduled in the reducer's backward\\n    # pass.\\n    def _match_all_reduce_for_bwd_pass(self):\\n        comm_work = []\\n        # Schedule comm in the same order as Reducer schedules them, i.e.\\n        # the order of the buckets. Retrieving the bucket order from the reducer\\n        # ensures that we keep the same order in join mode, such as when bucket\\n        # order is rebuilt dynamically.\\n\\n        # Returns grad_buckets in order, but real tensors are substituted with\\n        # zero tensors of the same shape.\\n        grad_buckets = self.reducer._get_zeros_like_grad_buckets()\\n        for grad_bucket in grad_buckets:\\n            # Joined processes contribute zero gradient. In the case that\\n            # divide_by_initial_world_size=True, we divide grads by the static\\n            # world size, if not, the dividing factor is reduced by the number\\n            # of joined processes.\\n            work = self.reducer._run_comm_hook(grad_bucket)\\n            comm_work.append(work)\\n        for work in comm_work:\\n            work.wait()\\n\\n    # Allreduces the used parameter mapping across ranks.\\n    def _match_unused_params_allreduce(self):\\n        locally_used_param_map = self.reducer._get_local_used_map()\\n        self.process_group.allreduce(locally_used_param_map)\\n\\n    def join(\\n        self,\\n        divide_by_initial_world_size: bool = True,\\n        enable: bool = True,\\n        throw_on_early_termination: bool = False,\\n    ):\\n        r\\\"\\\"\\\"\\n        Context manager for training with uneven inputs across processes in DDP.\\n\\n        This context manager will keep track of already-joined DDP processes,\\n        and \\\"shadow\\\" the forward and backward passes by inserting collective\\n        communication operations to match with the ones created by non-joined\\n        DDP processes. This will ensure each collective call has a corresponding\\n        call by already-joined DDP processes, preventing hangs or errors that\\n        would otherwise happen when training with uneven inputs across\\n        processes. Alternatively, if the flag ``throw_on_early_termination`` is\\n        specified to be ``True``, all trainers will throw an error once one rank\\n        runs out of inputs, allowing these errors to be caught and handled\\n        according to application logic.\\n\\n        Once all DDP processes have joined, the context manager will broadcast\\n        the model corresponding to the last joined process to all processes to\\n        ensure the model is the same across all processes\\n        (which is guaranteed by DDP).\\n\\n        To use this to enable training with uneven inputs across processes,\\n        simply wrap this context manager around your training loop. No further\\n        modifications to the model or data loading is required.\\n\\n        .. warning::\\n            If the model or training loop this context manager is wrapped around\\n            has additional distributed collective operations, such as\\n            ``SyncBatchNorm`` in the model's forward pass, then the flag\\n            ``throw_on_early_termination`` must be enabled. This is because this\\n            context manager is not aware of non-DDP collective communication.\\n            This flag will cause all ranks to throw when any one rank\\n            exhausts inputs, allowing these errors to be caught and recovered\\n            from across all ranks.\\n\\n        Args:\\n            divide_by_initial_world_size (bool): If ``True``, will divide\\n                gradients by the initial ``world_size`` DDP training was launched\\n                with. If ``False``, will compute the effective world size\\n                (number of ranks that have not depleted their inputs yet) and\\n                divide gradients by that during allreduce. Set\\n                ``divide_by_initial_world_size=True`` to ensure every input\\n                sample including the uneven inputs have equal weight in terms of\\n                how much they contribute to the global gradient. This is\\n                achieved by always dividing the gradient by the initial\\n                ``world_size`` even when we encounter uneven inputs. If you set\\n                this to ``False``, we divide the gradient by the remaining\\n                number of nodes. This ensures parity with training on a smaller\\n                ``world_size`` although it also means the uneven inputs would\\n                contribute more towards the global gradient. Typically, you\\n                would want to set this to ``True`` for cases where the last few\\n                inputs of your training job are uneven. In extreme cases, where\\n                there is a large discrepancy in the number of inputs, setting\\n                this to ``False`` might provide better results.\\n            enable (bool): Whether to enable uneven input detection or not. Pass\\n                in ``enable=False`` to disable in cases where you know that\\n                inputs are even across participating processes. Default is\\n                ``True``.\\n            throw_on_early_termination (bool): Whether to throw an error\\n                or continue training when at least one rank has exhausted\\n                inputs. If ``True``, will throw upon the first rank reaching end\\n                of data. If ``False``, will continue training with a smaller\\n                effective world size until all ranks are joined. Note that if\\n                this flag is specified, then the flag\\n                ``divide_by_initial_world_size`` would be ignored. Default\\n                is ``False``.\\n\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"Distributed\\\")\\n            >>> import torch\\n            >>> import torch.distributed as dist\\n            >>> import os\\n            >>> import torch.multiprocessing as mp\\n            >>> import torch.nn as nn\\n            >>> # On each spawned worker\\n            >>> def worker(rank):\\n            >>>     dist.init_process_group(\\\"nccl\\\", rank=rank, world_size=2)\\n            >>>     torch.cuda.set_device(rank)\\n            >>>     model = nn.Linear(1, 1, bias=False).to(rank)\\n            >>>     model = torch.nn.parallel.DistributedDataParallel(\\n            >>>         model, device_ids=[rank], output_device=rank\\n            >>>     )\\n            >>>     # Rank 1 gets one more input than rank 0.\\n            >>>     inputs = [torch.tensor([1]).float() for _ in range(10 + rank)]\\n            >>>     with model.join():\\n            >>>         for _ in range(5):\\n            >>>             for inp in inputs:\\n            >>>                 loss = model(inp).sum()\\n            >>>                 loss.backward()\\n            >>>     # Without the join() API, the below synchronization will hang\\n            >>>     # blocking for rank 1's allreduce to complete.\\n            >>>     torch.cuda.synchronize(device=rank)\\n        \\\"\\\"\\\"\\n        return Join(\\n            [self],\\n            enable,\\n            throw_on_early_termination,\\n            divide_by_initial_world_size=divide_by_initial_world_size,\\n        )\\n\\n    def join_hook(\\n        self,\\n        **kwargs,\\n    ):\\n        r\\\"\\\"\\\"\\n        DDP join hook enables training on uneven inputs by mirroring communications in forward and backward passes.\\n\\n        Arguments:\\n            kwargs (dict): a :class:`dict` containing any keyword arguments\\n                to modify the behavior of the join hook at run time; all\\n                :class:`Joinable` instances sharing the same join context\\n                manager are forwarded the same value for ``kwargs``.\\n\\n        The hook supports the following keyword arguments:\\n            divide_by_initial_world_size (bool, optional):\\n                If ``True``, then gradients are divided by the initial world\\n                size that DDP was launched with.\\n                If ``False``, then gradients are divided by the effective world\\n                size (i.e. the number of non-joined processes), meaning that\\n                the uneven inputs contribute more toward the global gradient.\\n                Typically, this should be set to ``True`` if the degree of\\n                unevenness is small but can be set to ``False`` in extreme\\n                cases for possibly better results.\\n                Default is ``True``.\\n        \\\"\\\"\\\"\\n        divide_by_initial_world_size = kwargs.get(\\\"divide_by_initial_world_size\\\", True)\\n        return _DDPJoinHook(\\n            self, divide_by_initial_world_size=divide_by_initial_world_size\\n        )\\n\\n    @property\\n    def join_device(self):\\n        return self.device\\n\\n    @property\\n    def join_process_group(self):\\n        return self.process_group\\n\\n    def _register_buffer_comm_hook(\\n        self,\\n        state,\\n        hook: Callable,\\n        comm_hook_location=_BufferCommHookLocation.POST_FORWARD,\\n    ):\\n        r\\\"\\\"\\\"\\n        Allow custom registration of hooks that define how buffer are synchronized across ranks.\\n\\n        The hook takes in an optional state and is passed in a Dict[str, Tensor]\\n        corresponding to buffer names and the buffers, and can run arbitrary reductions\\n        on buffers as opposed to DDP's default broadcast from rank 0. This is useful for\\n        example if a counter needs to be summed or averaged across ranks every iteration.\\n\\n        Args:\\n            state (Any): Optional state that is passed to the hook.\\n            hook (Callable): Callable with the following signature:\\n                         ``hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]``\\n            comm_hook_location (_BufferCommHookLocation): Enum value indicating\\n                            where to run the hook.\\n                            _BufferCommHookLocation.PRE_FORWARD means that the\\n                            hook will run _before_ the forward pass, and\\n                            _BufferCommHookLocation.POST_FORWARD means that the\\n                            hook will run _after_ the forward pass.\\n\\n            NOTE: To maximize performance, users can return a\\n                List[torch.futures.Future] from their hook, and DDP will\\n                install and await these hooks appropriately at the end of\\n                the backward pass. This will ensure all buffers are\\n                synchronized by the end of the backward pass. If this\\n                setting is used, it is recommended to pass\\n                comm_hook_location=_BufferCommHookLocation.POST_FORWARD,\\n                which will trigger the hook after the forward pass.\\n                If _BufferCommHookLocation.PRE_FORWARD is used, users must\\n                ensure appropriate synchronization when manipulating GPU\\n                buffers in the forward pass.\\n        \\\"\\\"\\\"\\n        assert callable(hook)\\n        self.buffer_hook = _BufferCommHook(\\n            buffer_comm_hook=hook,\\n            buffer_comm_hook_state=state,\\n            buffer_comm_hook_location=comm_hook_location,\\n        )\\n\\n    def register_comm_hook(self, state: object, hook: Callable):\\n        r\\\"\\\"\\\"\\n        Register communication hook for user-defined DDP aggregation of gradients across multiple workers.\\n\\n        This hook would be very useful for researchers to try out new ideas. For\\n        example, this hook can be used to implement several algorithms like GossipGrad\\n        and gradient compression which involve different communication strategies for\\n        parameter syncs while running Distributed DataParallel training.\\n\\n        Args:\\n            state (object): Passed to the hook to maintain any state information during the training process.\\n                            Examples include error feedback in gradient compression,\\n                            peers to communicate with next in GossipGrad, etc.\\n\\n                            It is locally stored by each worker\\n                            and shared by all the gradient tensors on the worker.\\n            hook (Callable): Callable with the following signature:\\n                             ``hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]``:\\n\\n                             This function is called once the bucket is ready. The\\n                             hook can perform whatever processing is needed and return\\n                             a Future indicating completion of any async work (ex: allreduce).\\n                             If the hook doesn't perform any communication, it still\\n                             must return a completed Future. The Future should hold the\\n                             new value of grad bucket's tensors. Once a bucket is ready,\\n                             c10d reducer would call this hook and use the tensors returned\\n                             by the Future and copy grads to individual parameters.\\n                             Note that the future's return type must be a single tensor.\\n\\n                             We also provide an API called ``get_future`` to retrieve a\\n                             Future associated with the completion of ``c10d.ProcessGroup.Work``.\\n                             ``get_future`` is currently supported for NCCL and also supported for most\\n                             operations on GLOO and MPI, except for peer to peer operations (send/recv).\\n\\n        .. warning ::\\n            Grad bucket's tensors will not be predivided by world_size. User is responsible\\n            to divide by the world_size in case of operations like allreduce.\\n\\n        .. warning ::\\n            DDP communication hook can only be registered once and should be registered\\n            before calling backward.\\n\\n        .. warning ::\\n            The Future object that hook returns should contain a single tensor\\n            that has the same shape with the tensors inside grad bucket.\\n\\n        .. warning ::\\n            ``get_future`` API supports NCCL, and partially GLOO and MPI backends (no support\\n            for peer-to-peer operations like send/recv) and will return a ``torch.futures.Future``.\\n\\n        Example::\\n            Below is an example of a noop hook that returns the same tensor.\\n\\n            >>> # xdoctest: +SKIP('undefined name')\\n            >>> def noop(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:\\n            >>>     fut = torch.futures.Future()\\n            >>>     fut.set_result(bucket.buffer())\\n            >>>     return fut\\n            >>> ddp.register_comm_hook(state=None, hook=noop)\\n\\n        Example::\\n            Below is an example of a Parallel SGD algorithm where gradients are encoded before\\n            allreduce, and then decoded after allreduce.\\n\\n            >>> # xdoctest: +SKIP('undefined name')\\n            >>> def encode_and_decode(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:\\n            >>>     encoded_tensor = encode(bucket.buffer())  # encode gradients\\n            >>>     fut = torch.distributed.all_reduce(encoded_tensor).get_future()\\n            >>>     # Define the then callback to decode.\\n            >>>     def decode(fut):\\n            >>>         decoded_tensor = decode(fut.value()[0])  # decode gradients\\n            >>>         return decoded_tensor\\n            >>>     return fut.then(decode)\\n            >>> ddp.register_comm_hook(state=None, hook=encode_and_decode)\\n        \\\"\\\"\\\"\\n        self._check_comm_hook(hook)\\n        assert self.logger is not None\\n        self.logger._set_comm_hook_name(hook.__qualname__)\\n        self._comm_hooks.append((hook, state))\\n        dist._register_comm_hook(self.reducer, state, hook)\\n\\n    def _register_builtin_comm_hook(self, comm_hook_type):\\n        r\\\"\\\"\\\"\\n        Register a built-in communication hook that specifies how DDP aggregates gradients across multiple workers.\\n\\n        The built-in hooks aim to provide efficient C++ implementations for certain hooks,\\n        which might not be as efficient if implemented in Python using a Python communication hook.\\n\\n        Args:\\n            comm_hook_type (dist.BuiltinCommHookType): type of communication hook, such as ALLREDUCE, FP16_COMPRESS, etc.\\n\\n        .. warning ::\\n            DDP communication hook can only be registered once and should be registered\\n            before calling backward.\\n\\n        Example::\\n            Below is an example of a FP16 compression where gradients are\\n            compressed into 16-bit floating-point numbers before allreduce, and\\n            then decompressed after allreduce.\\n\\n            >>> # xdoctest: +SKIP('undefined name')\\n            >>> ddp._register_builtin_comm_hook(dist.BuiltinCommHookType.FP16_COMPRESS)\\n\\n        \\\"\\\"\\\"\\n        assert self.logger is not None\\n        self.logger._set_comm_hook_name(str(comm_hook_type))\\n        dist._register_builtin_comm_hook(self.reducer, comm_hook_type)\\n\\n    def _register_fused_optim(self, optim: Type, *args, optim_params=None, **kwargs):\\n        r\\\"\\\"\\\"\\n        Register an optimizer in DDP to optimize parameter immediately after its gradient reduction.\\n\\n        Registers an optimizer with DDP such that the optimization for a\\n        parameter will run immediately when that parameter's gradient is\\n        finished with reduction, instead of waiting for all parameters'\\n        gradients to finish reduction. This can result in a training speedup\\n        depending on your workload since the optimizer can run while gradient\\n        reduction for other parameters are still ongoing. In addition, this has\\n        the potential to reduce peak memory consumption during training, as it\\n        only needs to load the per-parameter optimizer states of a single\\n        parameter at a time, instead of loading all per-parameter optimizer\\n        states at once.\\n\\n        Args:\\n            optim (Type): a ``torch.optim.Optimizer`` class to be registered\\n            as a fused optimizer.\\n            *args (Sequence[Any]): Arguments to forward to `optim`.\\n            optim_params (Optional[Iterable[torch.Tensor]]): Set of parameters\\n            to optimize, similar to `params` argument of traditional `torch.optim`\\n            Optimizers. If this is omitted, all DDP model parameters will be\\n            optimized.\\n            **kwargs: (Dict[str, Any]): Keyword arguments to forward to `optim`.\\n\\n        .. warning ::\\n            _register_fused_optim should only be called once on a DDP instance,\\n            and registering multiple fused optimizers for the same DDP model\\n            is not currently supported. Please ping\\n            https://github.com/pytorch/pytorch/issues/71595 if this is necessary\\n            for your use case.\\n\\n        .. warning ::\\n            _register_fused_optim and register_comm_hook currently do not\\n            compose together, meaning that custom DDP communication hooks are\\n            not supported with overlapped optimizers. Please ping\\n            https://github.com/pytorch/pytorch/issues/71595 if this is necessary\\n            for your use case.\\n\\n        .. warning ::\\n            Gradient accumulation and DDP `no_sync` are currently not supported\\n            with overlapped optimizer. Please ping\\n            https://github.com/pytorch/pytorch/issues/71595 if this is necessary\\n            for your use case.\\n\\n        Example::\\n\\n            >>> # xdoctest: +SKIP(\\\"No rendezvous handler\\\")\\n            >>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')\\n            >>> net = torch.nn.parallel.DistributedDataParallel(model, pg)\\n            >>> lr = 1e-2\\n            >>> betas = (0.9, 0.99)\\n            >>> eps = 1e-6\\n            >>> net._register_fused_optim(torch.optim.Adam, lr, betas=betas, eps=eps)\\n            >>> # Example with subset of parameters\\n            >>> params_to_opt = [list(net.parameters())[0]]\\n            >>> net._register_fused_optim(\\n            ...   torch.optim.Adam, lr, optim_params=params_to_opt,  betas=betas, eps=eps\\n            ... )\\n        \\\"\\\"\\\"\\n        # Note: importing in function, otherwise this will cause a circular\\n        # import as optimizer_overlap module needs to import DistributedDataParallel.\\n        from torch.distributed.algorithms._optimizer_overlap import _as_overlapped_optim\\n\\n        overlapped_optim = _as_overlapped_optim(optim, optim_params, *args, **kwargs)\\n        try:\\n            overlapped_optim.register_ddp(self)\\n        except NotImplementedError as e:\\n            raise RuntimeError(\\n                f\\\"{optim} does not support overlapped DDP. Please file an issue to PyTorch or the respective owner of {optim}.\\\"\\n            ) from e\\n\\n    def _distributed_broadcast_coalesced(\\n        self, tensors, buffer_size, authoritative_rank=0\\n    ):\\n        dist._broadcast_coalesced(\\n            self.process_group, tensors, buffer_size, authoritative_rank\\n        )\\n\\n    def _check_sync_bufs_post_fwd(self):\\n        return (\\n            self.will_sync_module_buffers()\\n            and hasattr(self, \\\"buffer_hook\\\")\\n            and self.buffer_hook.buffer_comm_hook_location\\n            == _BufferCommHookLocation.POST_FORWARD\\n        )\\n\\n    def _check_sync_bufs_pre_fwd(self):\\n        return self.will_sync_module_buffers() and (\\n            not hasattr(self, \\\"buffer_hook\\\")\\n            or self.buffer_hook.buffer_comm_hook_location\\n            == _BufferCommHookLocation.PRE_FORWARD\\n        )\\n\\n    def will_sync_module_buffers(self):\\n        return (\\n            self.require_forward_param_sync\\n            and self.broadcast_buffers\\n            and len(self.modules_buffers) > 0\\n        )\\n\\n    def _find_common_rank(self, input_rank, rank_cond):\\n        # -1 indicates that this rank is not under consideration to be the\\n        # common_rank\\n        rank_to_use = torch.tensor(\\n            [input_rank if rank_cond else -1],\\n            device=self.device,\\n        )\\n        dist.all_reduce(rank_to_use, op=ReduceOp.MAX, group=self.process_group)\\n        if rank_to_use.item() == -1:\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"BUG! Expected rank_cond to be true for at least one process.\\\"\\n                \\\" This indicates a bug in PyTorch, please report an issue.\\\",\\n            )\\n        return rank_to_use.item()\\n\\n    def _sync_buffers(self):\\n        with torch.no_grad():\\n            # module buffer sync\\n            # Synchronize buffers across processes.\\n            # If we are running DDP with the join manager, we have to agree\\n            # upon a rank to sync module buffers from, since rank 0 may\\n            # already have been joined and have stale module buffers.\\n            if self._join_config.enable:\\n                authoritative_rank = self._find_common_rank(\\n                    self._distributed_rank, True\\n                )\\n            else:\\n                # The process with rank 0 is considered the authoritative copy.\\n                authoritative_rank = 0\\n            # Update self.modules_buffers incase any buffers were\\n            # reassigned.\\n            self._assign_modules_buffers()\\n            self._sync_module_buffers(authoritative_rank)\\n\\n    def _sync_module_buffers(self, authoritative_rank):\\n        if not hasattr(self, \\\"buffer_hook\\\"):\\n            self._default_broadcast_coalesced(authoritative_rank=authoritative_rank)\\n        else:\\n            hook = self.buffer_hook.buffer_comm_hook\\n            state = self.buffer_hook.buffer_comm_hook_state\\n            futs = hook(state, self.named_module_buffers)\\n            if futs is not None:\\n                self.reducer._install_post_backward_futures(futs)\\n\\n    def _default_broadcast_coalesced(\\n        self, bufs=None, bucket_size=None, authoritative_rank=0\\n    ):\\n        \\\"\\\"\\\"\\n        Broadcasts buffers from rank 0 to rest of workers.\\n\\n        If bufs, bucket_size are None, default values self.modules_buffers\\n        and self.broadcast_bucket_size are used instead.\\n        \\\"\\\"\\\"\\n        if bufs is None:\\n            bufs = self.modules_buffers\\n        if bucket_size is None:\\n            bucket_size = self.broadcast_bucket_size\\n\\n        self._distributed_broadcast_coalesced(bufs, bucket_size, authoritative_rank)\\n\\n    def _passing_sync_batchnorm_handle(self, module):\\n        for layer in module.modules():\\n            if isinstance(layer, torch.nn.modules.SyncBatchNorm):\\n                if self.device_type == \\\"cpu\\\":\\n                    self._log_and_throw(\\n                        ValueError,\\n                        \\\"SyncBatchNorm layers only work with GPU modules\\\",\\n                    )\\n\\n    def _check_comm_hook(self, hook):\\n        if not callable(hook):\\n            self._log_and_throw(TypeError, \\\"Communication hook must be callable.\\\")\\n\\n        sig = inspect.signature(hook)\\n        if (\\n            sig.parameters[\\\"bucket\\\"].annotation != inspect._empty\\n            and sig.parameters[\\\"bucket\\\"].annotation != dist.GradBucket\\n        ):\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"Communication hook: bucket annotation should be dist.GradBucket.\\\",\\n            )\\n\\n        if (\\n            sig.return_annotation != inspect._empty\\n            and sig.return_annotation != torch.futures.Future[torch.Tensor]\\n        ):\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"Communication hook: return annotation should be torch.futures.Future[torch.Tensor].\\\",\\n            )\\n\\n        if hook.__name__ in [\\n            \\\"bf16_compress_hook\\\",\\n            \\\"bf16_compress_wrapper_hook\\\",\\n        ] and (\\n            (torch.version.cuda is None and torch.version.hip is None)\\n            or (\\n                torch.version.cuda is not None\\n                and int(torch.version.cuda.split(\\\".\\\")[0]) < 11\\n            )\\n            or not dist.is_available()\\n            or not dist.is_nccl_available()\\n            or torch.cuda.nccl.version() < (2, 10)\\n        ):\\n            self._log_and_throw(\\n                TypeError,\\n                \\\"BF16 all reduce communication hook required CUDA 11+ and NCCL 2.10+.\\\",\\n            )\\n\\n    @property\\n    def _distributed_rank(self):\\n        return dist.get_rank(self.process_group)\\n\\n    @staticmethod\\n    def _get_data_parallel_params(module, named_params=False):\\n        \\\"\\\"\\\"Return a generator of parameters managed by a given DDP unit.\\\"\\\"\\\"\\n        for param in (\\n            module.parameters() if not named_params else module.named_parameters()\\n        ):\\n            if not hasattr(param, \\\"_ddp_ignored\\\"):\\n                yield param\\n\\n    @staticmethod\\n    def _set_params_and_buffers_to_ignore_for_model(\\n        module, params_and_buffers_to_ignore\\n    ):\\n        \\\"\\\"\\\"\\n        Set parameters and buffers to be ignored by DDP.\\n\\n        Expected format for parameters is the fully qualified name: {module_name}.{param_name}, and\\n        similarly, {module_name}.{buffer_name} for buffers. For example:\\n        params_to_ignore = []\\n        # NB: model here is vanilla PyTorch module, not yet wrapped with DDP.\\n        for module_name, module in model.named_modules():\\n            for param_name, param in module.named_parameters(recurse=False):\\n                if should_ignore(param):\\n                    # Create expected format\\n                    fqn = f\\\"{module_name}.{param_name}\\\"\\n                    params_to_ignore.append(fqn)\\n        torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model(\\n            model,\\n            params_to_ignore\\n        )\\n        \\\"\\\"\\\"\\n        # This is a workaround to set parameters and buffers DDP should ignore\\n        # during synchronization. It will be removed when the API is finalized\\n        # as part of addressing https://github.com/pytorch/pytorch/issues/43690.\\n        module._ddp_params_and_buffers_to_ignore = params_and_buffers_to_ignore\\n        for name, param in module.named_parameters():\\n            if name in params_and_buffers_to_ignore:\\n                param._ddp_ignored = True\\n        for name, buffer in module.named_buffers():\\n            if name in params_and_buffers_to_ignore:\\n                buffer._ddp_ignored = True\\n\\n    def _get_ddp_logging_data(self):\\n        r\\\"\\\"\\\"\\n        Return a dictionary of logging data for debugging and analysis.\\n\\n        This interface can be called after DistributedDataParallel() is\\n        constructed. It returns a dictionary of logging data. It could help\\n        for debugging and analysis. The logging data includes DistributedDataParallel\\n        constructor input parameters, some internal states of DistributedDataParallel\\n        and performance metrics. Simply print the dictionary and see what\\n        these metrics are.\\n        This is a prototype interface and subject to change in the future.\\n        \\\"\\\"\\\"\\n        assert self.logger is not None\\n        ddp_logging_data = self.logger._get_ddp_logging_data()\\n        return {**ddp_logging_data.strs_map, **ddp_logging_data.ints_map}\\n\\n    def _set_ddp_runtime_logging_sample_rate(self, sample_rate):\\n        r\\\"\\\"\\\"\\n        Set sample_rate of collecting runtime stats.\\n\\n        This interface allows users to set sample_rate of collecting\\n        runtime stats. The runtime stats will be recorded for the\\n        first 10 iterations, after 10 iterations runtime stats will be\\n        recorded once every \\\"sample_rate\\\" training iterations. In\\n        default, runtime stats are recorded for the first 10 iterations,\\n        after 10 iterations runtime stats are recorded once every\\n        \\\"kDDPRuntimeLoggingSampleRate=100\\\" training iterations.\\n        This is a prototype interface and subject to change in the future.\\n        \\\"\\\"\\\"\\n        if sample_rate < 1:\\n            self._log_and_throw(\\n                ValueError,\\n                \\\"DDP runtime logging sample rate should be equal or greater than 1\\\",\\n            )\\n        self.reducer._set_ddp_runtime_logging_sample_rate(sample_rate)\\n\\n    def _set_static_graph(self):\\n        \\\"\\\"\\\"\\n        Set static graph for DDP.\\n\\n        It is recommended to set static graph in the DDP constructor, which will\\n        call this private API internally.\\n        \\\"\\\"\\\"\\n        # If self.static_graph has been set, no need to set it again\\n        if self.static_graph:\\n            warnings.warn(\\n                \\\"You've set static_graph to be True, no need to set it again.\\\"\\n            )\\n            return\\n        self.static_graph = True\\n        self._static_graph_delay_allreduce_enqueued = False\\n        self.reducer._set_static_graph()\\n        assert self.logger is not None\\n        self.logger._set_static_graph()\\n        if self.find_unused_parameters:\\n            warnings.warn(\\n                \\\"You passed find_unused_parameters=true to DistributedDataParallel, \\\"\\n                \\\"`_set_static_graph` will detect unused parameters automatically, so \\\"\\n                \\\"you do not need to set find_unused_parameters=true, just be sure these \\\"\\n                \\\"unused parameters will not change during training loop while calling \\\"\\n                \\\"`_set_static_graph`.\\\"\\n            )\\n\\n    def _remove_autograd_hooks(self):\\n        \\\"\\\"\\\"Remove autograd hooks registered by the reducer on the model parameters.\\\"\\\"\\\"\\n        self.reducer._remove_autograd_hooks()\\n\\n    def _check_reducer_finalized(self):\\n        \\\"\\\"\\\"\\n        Check if the reducer has processed all buckets and finalized the backward appropriately.\\n\\n        It is useful to call this method after calling .backward() in your training loop\\n        in order to avoid subsequent hard to debug errors down the road due to the\\n        reducer not finalizing backward.\\n        \\\"\\\"\\\"\\n        self.reducer._check_reducer_finalized()\\n\\n    def _set_sparse_metadata(self, global_unique_ids):\\n        self.reducer._set_sparse_metadata(global_unique_ids)\\n\\n    def _update_process_group(self, new_process_group):\\n        \\\"\\\"\\\"\\n        Dynamically updates the process group for DDP so that we can shrink/expand DDP\\n        world size without having to reinitialize DDP.\\n\\n        NOTE: If you are using custom communications hooks via, register_comm_hook,\\n        you need to update the process groups for those hooks separately.\\n        \\\"\\\"\\\"\\n        # Force a rebuild of buckets for a new process group. This ensures all ranks\\n        # are synchronized in terms of when they will rebuild buckets and also\\n        # re-evaluates previous assumptions of buckets given the world size might have\\n        # changed.\\n        self._has_rebuilt_buckets = False\\n        self.reducer._reset_state()\\n\\n        if not _rank_not_in_group(new_process_group):\\n            self.process_group = new_process_group\\n            self.reducer._update_process_group(new_process_group)\\n\\n    def _set_ddp_sink_clone(self, val: bool):\\n        \\\"\\\"\\\"\\n        Sets whether or not DDPSink should clone the output tensors or not.\\n        The default is True since if the loss is modified in place we run\\n        into the view is modified in-place error.\\n\\n        Although, cloning the tensors can add significant memory and\\n        performance hit if the number and size of tensors are large. As\\n        a result, this can be set to False if you are not modifying the\\n        loss in place.\\n        \\\"\\\"\\\"\\n        self._ddp_sink_clone = val\\n\\n\\nfrom collections import OrderedDict\\nfrom typing import (\\n    cast,\\n    Dict,\\n    Iterator,\\n    List,\\n    Optional,\\n    Sequence,\\n    Set,\\n    TYPE_CHECKING,\\n    TypeVar,\\n    Union,\\n)\\n\\nimport torch\\nfrom torch._utils import _get_device_index\\nfrom torch.nn.modules import Module\\nfrom torch.nn.parallel import comm\\n\\n\\nif TYPE_CHECKING:\\n    from torch.jit import ScriptModule\\n    from torch.jit._state import EnabledProxy\\n\\n\\n__all__ = [\\\"replicate\\\"]\\n\\n\\ndef _is_script_module(module: Module) -> bool:\\n    import torch.jit\\n\\n    return isinstance(module, torch.jit.ScriptModule)\\n\\n\\ndef _is_script_method(module: Module) -> bool:\\n    import torch.jit\\n\\n    return isinstance(module, torch._C.ScriptMethod)\\n\\n\\ndef _init_script_module() -> \\\"ScriptModule\\\":\\n    import torch.jit\\n\\n    return torch.jit.ScriptModule()\\n\\n\\ndef _is_jit_enabled() -> \\\"EnabledProxy\\\":\\n    import torch.jit._state\\n\\n    return torch.jit._state._enabled\\n\\n\\n# Check if we can safely replicate the module.\\n# there are two types of module:\\n# 1. python modules\\n# 2. ScriptModule\\n#\\n# currently a module cannot be replicated properly if the descendants of\\n# any ScriptModule contains python module (type 1 above)\\ndef _replicatable_module(module: Module, memo: Optional[Set[Module]] = None) -> bool:\\n    # module.modules() contains module itself as the first element\\n    def descendant_modules(module: Module) -> Iterator[Module]:\\n        gen = module.modules()\\n        next(gen)\\n        return gen\\n\\n    if not _is_jit_enabled():\\n        return True\\n    if memo is None:\\n        memo = set()\\n\\n    # memoize visited modules\\n    memo.add(module)\\n    if _is_script_module(module):\\n        memo.update(descendant_modules(module))\\n        return all(\\n            _is_script_module(descendant) for descendant in descendant_modules(module)\\n        )\\n\\n    for child in module.children():\\n        # since any unreplicatable module will cause the check to return\\n        # False early, visited modules here can be safely ignored.\\n        if child in memo:\\n            continue\\n        if not _replicatable_module(child, memo):\\n            return False\\n\\n    return True\\n\\n\\ndef _broadcast_coalesced_reshape(\\n    tensors: Sequence[torch.Tensor],\\n    devices: Sequence[Union[int, torch.device]],\\n    detach: bool = False,\\n) -> List[List[torch.Tensor]]:\\n    from torch.nn.parallel._functions import Broadcast\\n\\n    if detach:\\n        return comm.broadcast_coalesced(tensors, devices)\\n    else:\\n        # Use the autograd function to broadcast if not detach\\n        if len(tensors) > 0:\\n            tensor_copies = Broadcast.apply(devices, *tensors)\\n            return [\\n                tensor_copies[i : i + len(tensors)]\\n                for i in range(0, len(tensor_copies), len(tensors))\\n            ]\\n        else:\\n            return []\\n\\n\\nT = TypeVar(\\\"T\\\", bound=Module)\\n\\n\\ndef replicate(\\n    network: T,\\n    devices: Sequence[Union[int, torch.device]],\\n    detach: bool = False,\\n) -> List[T]:\\n    if not _replicatable_module(network):\\n        raise RuntimeError(\\n            \\\"Cannot replicate network where python modules are \\\"\\n            \\\"childrens of ScriptModule\\\"\\n        )\\n\\n    if not devices:\\n        return []\\n\\n    devices = [_get_device_index(x, True) for x in devices]\\n    num_replicas = len(devices)\\n\\n    params = list(network.parameters())\\n    param_indices = {param: idx for idx, param in enumerate(params)}\\n    param_copies = _broadcast_coalesced_reshape(params, devices, detach)\\n\\n    buffers = list(network.buffers())\\n    buffers_rg: List[torch.Tensor] = []\\n    buffers_not_rg: List[torch.Tensor] = []\\n    for buf in buffers:\\n        if buf.requires_grad and not detach:\\n            buffers_rg.append(buf)\\n        else:\\n            buffers_not_rg.append(buf)\\n\\n    buffer_indices_rg = {buf: idx for idx, buf in enumerate(buffers_rg)}\\n    buffer_indices_not_rg = {buf: idx for idx, buf in enumerate(buffers_not_rg)}\\n\\n    buffer_copies_rg = _broadcast_coalesced_reshape(buffers_rg, devices, detach=detach)\\n    buffer_copies_not_rg = _broadcast_coalesced_reshape(\\n        buffers_not_rg, devices, detach=True\\n    )\\n\\n    modules = list(network.modules())\\n    module_copies: List[List[Module]] = [[] for _ in devices]\\n    module_indices: Dict[Module, int] = {}\\n\\n    for i, module in enumerate(modules):\\n        module_indices[module] = i\\n        for j in range(num_replicas):\\n            replica = module._replicate_for_data_parallel()\\n            # This is a temporary fix for DDP. DDP needs to access the\\n            # replicated model parameters. It used to do so through\\n            # `mode.parameters()`. The fix added in #33907 for DP stops the\\n            # `parameters()` API from exposing the replicated parameters.\\n            # Hence, we add a `_former_parameters` dict here to support DDP.\\n            replica._former_parameters = OrderedDict()\\n\\n            module_copies[j].append(replica)\\n\\n    for i, module in enumerate(modules):\\n        for key, child in module._modules.items():\\n            if child is None:\\n                for j in range(num_replicas):\\n                    replica = module_copies[j][i]\\n                    replica._modules[key] = None\\n            else:\\n                module_idx = module_indices[child]\\n                for j in range(num_replicas):\\n                    replica = module_copies[j][i]\\n                    setattr(replica, key, module_copies[j][module_idx])\\n        for key, param in module._parameters.items():\\n            if param is None:\\n                for j in range(num_replicas):\\n                    replica = module_copies[j][i]\\n                    replica._parameters[key] = None\\n            else:\\n                param_idx = param_indices[param]\\n                for j in range(num_replicas):\\n                    replica = module_copies[j][i]\\n                    param_copy = param_copies[j][param_idx]\\n                    # parameters in replicas are no longer leaves,\\n                    # so setattr them as non-parameter attributes\\n                    setattr(replica, key, param_copy)\\n                    # expose the parameter for DDP\\n                    replica._former_parameters[key] = param_copy\\n        for key, buf in module._buffers.items():  # type: ignore[assignment]\\n            if buf is None:\\n                for j in range(num_replicas):\\n                    replica = module_copies[j][i]\\n                    replica._buffers[key] = None\\n            else:\\n                if buf.requires_grad and not detach:\\n                    buffer_copies = buffer_copies_rg\\n                    buffer_idx = buffer_indices_rg[buf]\\n                else:\\n                    buffer_copies = buffer_copies_not_rg\\n                    buffer_idx = buffer_indices_not_rg[buf]\\n                for j in range(num_replicas):\\n                    replica = module_copies[j][i]\\n                    setattr(replica, key, buffer_copies[j][buffer_idx])\\n\\n    return [cast(T, module_copies[j][0]) for j in range(num_replicas)]\\n\\n\\nimport threading\\nfrom typing import Any, cast, Dict, List, Optional, Sequence, Tuple, Union\\n\\nimport torch\\nfrom torch._utils import ExceptionWrapper\\nfrom torch.cuda._utils import _get_device_index\\nfrom torch.nn.modules import Module\\n\\n\\n__all__ = [\\\"get_a_var\\\", \\\"parallel_apply\\\"]\\n\\n\\ndef get_a_var(\\n    obj: Union[torch.Tensor, List[Any], Tuple[Any, ...], Dict[Any, Any]],\\n) -> Optional[torch.Tensor]:\\n    if isinstance(obj, torch.Tensor):\\n        return obj\\n\\n    if isinstance(obj, (list, tuple)):\\n        for result in map(get_a_var, obj):\\n            if isinstance(result, torch.Tensor):\\n                return result\\n    if isinstance(obj, dict):\\n        for result in map(get_a_var, obj.items()):\\n            if isinstance(result, torch.Tensor):\\n                return result\\n    return None\\n\\n\\ndef parallel_apply(\\n    modules: Sequence[Module],\\n    inputs: Sequence[Any],\\n    kwargs_tup: Optional[Sequence[Dict[str, Any]]] = None,\\n    devices: Optional[Sequence[Optional[Union[int, torch.device]]]] = None,\\n) -> List[Any]:\\n    r\\\"\\\"\\\"Apply each `module` in :attr:`modules` in parallel on each of :attr:`devices`.\\n\\n    Args:\\n        modules (Module): modules to be parallelized\\n        inputs (tensor): inputs to the modules\\n        devices (list of int or torch.device): CUDA devices\\n\\n    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and\\n    :attr:`devices` (if given) should all have same length. Moreover, each\\n    element of :attr:`inputs` can either be a single object as the only argument\\n    to a module, or a collection of positional arguments.\\n    \\\"\\\"\\\"\\n    assert len(modules) == len(\\n        inputs\\n    ), f\\\"The number of modules {len(modules)} is not equal to the number of inputs {len(inputs)}\\\"\\n    if kwargs_tup is not None:\\n        assert len(modules) == len(kwargs_tup)\\n    else:\\n        kwargs_tup = (cast(Dict[str, Any], {}),) * len(modules)\\n    if devices is not None:\\n        assert len(modules) == len(devices)\\n    else:\\n        devices = [None] * len(modules)\\n    devices = [_get_device_index(x, True) for x in devices]\\n    streams = [torch.cuda.current_stream(x) for x in devices]\\n    lock = threading.Lock()\\n    results = {}\\n    grad_enabled, autocast_enabled = (\\n        torch.is_grad_enabled(),\\n        torch.is_autocast_enabled(),\\n    )\\n\\n    def _worker(\\n        i: int,\\n        module: Module,\\n        input: Any,\\n        kwargs: Dict[str, Any],\\n        device: Optional[Union[int, torch.device]] = None,\\n        stream: Optional[torch.cuda.Stream] = None,\\n    ) -> None:\\n        torch.set_grad_enabled(grad_enabled)\\n        if device is None:\\n            t = get_a_var(input)\\n            if t is None:\\n                with lock:\\n                    results[i] = ExceptionWrapper(\\n                        where=f\\\"in replica {i}, no device was provided and no tensor input was found; \\\"\\n                        \\\"device cannot be resolved\\\"\\n                    )\\n                return\\n            device = t.get_device()\\n        if stream is None:\\n            stream = torch.cuda.current_stream(device)\\n        try:\\n            with torch.cuda.device(device), torch.cuda.stream(\\n                stream\\n            ), torch.amp.autocast(\\\"cuda\\\", enabled=autocast_enabled):\\n                # this also avoids accidental slicing of `input` if it is a Tensor\\n                if not isinstance(input, (list, tuple)):\\n                    input = (input,)\\n                output = module(*input, **kwargs)\\n            with lock:\\n                results[i] = output\\n        except Exception:\\n            with lock:\\n                results[i] = ExceptionWrapper(\\n                    where=f\\\"in replica {i} on device {device}\\\"\\n                )\\n\\n    if len(modules) > 1:\\n        threads = [\\n            threading.Thread(\\n                target=_worker, args=(i, module, input, kwargs, device, stream)\\n            )\\n            for i, (module, input, kwargs, device, stream) in enumerate(\\n                zip(modules, inputs, kwargs_tup, devices, streams)\\n            )\\n        ]\\n\\n        for thread in threads:\\n            thread.start()\\n        for thread in threads:\\n            thread.join()\\n    else:\\n        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0], streams[0])\\n\\n    outputs = []\\n    for i in range(len(inputs)):\\n        output = results[i]\\n        if isinstance(output, ExceptionWrapper):\\n            output.reraise()\\n        outputs.append(output)\\n    return outputs\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing import Any, Dict, List, Optional, overload, Sequence, Tuple, TypeVar, Union\\nfrom typing_extensions import deprecated\\n\\nimport torch\\nfrom torch.nn.parallel._functions import Gather, Scatter\\n\\n\\n__all__ = [\\\"scatter\\\", \\\"scatter_kwargs\\\", \\\"gather\\\"]\\n\\n\\n@deprecated(\\n    \\\"`is_namedtuple` is deprecated, please use the python checks instead\\\",\\n    category=FutureWarning,\\n)\\ndef is_namedtuple(obj: Any) -> bool:\\n    # Check if type was created from collections.namedtuple or a typing.NamedTuple.\\n    return _is_namedtuple(obj)\\n\\n\\ndef _is_namedtuple(obj: Any) -> bool:\\n    # Check if type was created from collections.namedtuple or a typing.NamedTuple.\\n    return (\\n        isinstance(obj, tuple) and hasattr(obj, \\\"_asdict\\\") and hasattr(obj, \\\"_fields\\\")\\n    )\\n\\n\\nT = TypeVar(\\\"T\\\", dict, list, tuple)\\n\\n\\n# For some reason, 'scatter' returns a tuple when given a single Tensor input but a list otherwise.\\n@overload\\ndef scatter(\\n    inputs: torch.Tensor,\\n    target_gpus: Sequence[Union[int, torch.device]],\\n    dim: int = ...,\\n) -> Tuple[torch.Tensor, ...]:\\n    ...\\n\\n\\n@overload\\ndef scatter(\\n    inputs: T,\\n    target_gpus: Sequence[Union[int, torch.device]],\\n    dim: int = ...,\\n) -> List[T]:\\n    ...\\n\\n\\ndef scatter(inputs, target_gpus, dim=0):\\n    r\\\"\\\"\\\"Slice tensors into approximately equal chunks and distributes them across given GPUs.\\n\\n    Duplicates references to objects that are not tensors.\\n    \\\"\\\"\\\"\\n\\n    def scatter_map(obj):\\n        if isinstance(obj, torch.Tensor):\\n            return Scatter.apply(target_gpus, None, dim, obj)\\n        if _is_namedtuple(obj):\\n            return [type(obj)(*args) for args in zip(*map(scatter_map, obj))]\\n        if isinstance(obj, tuple) and len(obj) > 0:\\n            return list(zip(*map(scatter_map, obj)))\\n        if isinstance(obj, list) and len(obj) > 0:\\n            return [list(i) for i in zip(*map(scatter_map, obj))]\\n        if isinstance(obj, dict) and len(obj) > 0:\\n            return [type(obj)(i) for i in zip(*map(scatter_map, obj.items()))]\\n        return [obj for _ in target_gpus]\\n\\n    # After scatter_map is called, a scatter_map cell will exist. This cell\\n    # has a reference to the actual function scatter_map, which has references\\n    # to a closure that has a reference to the scatter_map cell (because the\\n    # fn is recursive). To avoid this reference cycle, we set the function to\\n    # None, clearing the cell\\n    try:\\n        res = scatter_map(inputs)\\n    finally:\\n        scatter_map = None  # type: ignore[assignment]\\n    return res\\n\\n\\ndef scatter_kwargs(\\n    inputs: Tuple[Any, ...],\\n    kwargs: Optional[Dict[str, Any]],\\n    target_gpus: Sequence[Union[int, torch.device]],\\n    dim: int = 0,\\n) -> Tuple[Tuple[Any, ...], Tuple[Dict[str, Any], ...]]:\\n    r\\\"\\\"\\\"Scatter with support for kwargs dictionary.\\\"\\\"\\\"\\n    scattered_inputs = scatter(inputs, target_gpus, dim) if inputs else []\\n    scattered_kwargs = scatter(kwargs, target_gpus, dim) if kwargs else []\\n    if len(scattered_inputs) < len(scattered_kwargs):\\n        scattered_inputs.extend(\\n            () for _ in range(len(scattered_kwargs) - len(scattered_inputs))\\n        )\\n    elif len(scattered_kwargs) < len(inputs):\\n        scattered_kwargs.extend(\\n            {} for _ in range(len(scattered_inputs) - len(scattered_kwargs))\\n        )\\n    return tuple(scattered_inputs), tuple(scattered_kwargs)\\n\\n\\ndef gather(outputs: Any, target_device: Union[int, torch.device], dim: int = 0) -> Any:\\n    r\\\"\\\"\\\"Gather tensors from different GPUs on a specified device.\\n\\n    This function is useful for gathering the results of a distributed computation.\\n    It takes a sequence of objects, one for each GPU, and returns a single object\\n    on the specified device.\\n\\n    Args:\\n        outputs (Any): A sequence of objects (potentially tensors) to gather.\\n        target_device (Union[int, torch.device]): The device to gather the tensors to.\\n            Use 'cpu' for CPU to avoid a deprecation warning.\\n        dim (int, optional): The dimension along which to gather. Default: 0.\\n\\n    Returns:\\n        Any: A gathered object (potentially tensor) on the specified device.\\n    \\\"\\\"\\\"\\n\\n    def gather_map(outputs):\\n        out = outputs[0]\\n        if isinstance(out, torch.Tensor):\\n            return Gather.apply(target_device, dim, *outputs)\\n        if out is None:\\n            return None\\n        if isinstance(out, dict):\\n            if not all(len(out) == len(d) for d in outputs):\\n                raise ValueError(\\\"All dicts must have the same number of keys\\\")\\n            return type(out)((k, gather_map([d[k] for d in outputs])) for k in out)\\n        if _is_namedtuple(out):\\n            return type(out)._make(map(gather_map, zip(*outputs)))\\n        return type(out)(map(gather_map, zip(*outputs)))\\n\\n    # Recursive function calls like this create reference cycles.\\n    # Setting the function to None clears the refcycle.\\n    try:\\n        res = gather_map(outputs)\\n    finally:\\n        gather_map = None  # type: ignore[assignment]\\n    return res\\n\\n\\n# mypy: allow-untyped-defs\\nfrom typing_extensions import deprecated\\n\\nfrom torch.nn.parallel.data_parallel import data_parallel, DataParallel\\nfrom torch.nn.parallel.distributed import DistributedDataParallel\\nfrom torch.nn.parallel.parallel_apply import parallel_apply\\nfrom torch.nn.parallel.replicate import replicate\\nfrom torch.nn.parallel.scatter_gather import gather, scatter\\n\\n\\n__all__ = [\\n    \\\"replicate\\\",\\n    \\\"scatter\\\",\\n    \\\"parallel_apply\\\",\\n    \\\"gather\\\",\\n    \\\"data_parallel\\\",\\n    \\\"DataParallel\\\",\\n    \\\"DistributedDataParallel\\\",\\n]\\n\\n\\n@deprecated(\\n    \\\"`torch.nn.parallel.DistributedDataParallelCPU` is deprecated, \\\"\\n    \\\"please use `torch.nn.parallel.DistributedDataParallel` instead.\\\",\\n    category=FutureWarning,\\n)\\nclass DistributedDataParallelCPU(DistributedDataParallel):\\n    pass\\n\\n\\n# mypy: allow-untyped-defs\\nimport operator\\nimport warnings\\nfrom itertools import chain\\nfrom typing import Any, Dict, Generic, List, Optional, Sequence, Tuple, TypeVar, Union\\n\\nimport torch\\nfrom torch._utils import (\\n    _get_all_device_indices,\\n    _get_available_device_type,\\n    _get_device_index,\\n    _get_devices_properties,\\n)\\nfrom torch.nn.modules import Module\\nfrom torch.nn.parallel.parallel_apply import parallel_apply\\nfrom torch.nn.parallel.replicate import replicate\\nfrom torch.nn.parallel.scatter_gather import gather, scatter_kwargs\\n\\n\\n__all__ = [\\\"DataParallel\\\", \\\"data_parallel\\\"]\\n\\n\\ndef _check_balance(device_ids: Sequence[Union[int, torch.device]]) -> None:\\n    imbalance_warn = \\\"\\\"\\\"\\n    There is an imbalance between your GPUs. You may want to exclude GPU {} which\\n    has less than 75% of the memory or cores of GPU {}. You can do so by setting\\n    the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES\\n    environment variable.\\\"\\\"\\\"\\n    device_ids = [_get_device_index(x, True) for x in device_ids]\\n    dev_props = _get_devices_properties(device_ids)\\n\\n    def warn_imbalance(get_prop):\\n        values = [get_prop(props) for props in dev_props]\\n        min_pos, min_val = min(enumerate(values), key=operator.itemgetter(1))\\n        max_pos, max_val = max(enumerate(values), key=operator.itemgetter(1))\\n        if min_val / max_val < 0.75:\\n            warnings.warn(\\n                imbalance_warn.format(device_ids[min_pos], device_ids[max_pos])\\n            )\\n            return True\\n        return False\\n\\n    if warn_imbalance(lambda props: props.total_memory):\\n        return\\n    if warn_imbalance(lambda props: props.multi_processor_count):\\n        return\\n\\n\\nT = TypeVar(\\\"T\\\", bound=Module)\\n\\n\\nclass DataParallel(Module, Generic[T]):\\n    r\\\"\\\"\\\"Implements data parallelism at the module level.\\n\\n    This container parallelizes the application of the given :attr:`module` by\\n    splitting the input across the specified devices by chunking in the batch\\n    dimension (other objects will be copied once per device). In the forward\\n    pass, the module is replicated on each device, and each replica handles a\\n    portion of the input. During the backwards pass, gradients from each replica\\n    are summed into the original module.\\n\\n    The batch size should be larger than the number of GPUs used.\\n\\n    .. warning::\\n        It is recommended to use :class:`~torch.nn.parallel.DistributedDataParallel`,\\n        instead of this class, to do multi-GPU training, even if there is only a single\\n        node. See: :ref:`cuda-nn-ddp-instead` and :ref:`ddp`.\\n\\n    Arbitrary positional and keyword inputs are allowed to be passed into\\n    DataParallel but some types are specially handled. tensors will be\\n    **scattered** on dim specified (default 0). tuple, list and dict types will\\n    be shallow copied. The other types will be shared among different threads\\n    and can be corrupted if written to in the model's forward pass.\\n\\n    The parallelized :attr:`module` must have its parameters and buffers on\\n    ``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`\\n    module.\\n\\n    .. warning::\\n        In each forward, :attr:`module` is **replicated** on each device, so any\\n        updates to the running module in ``forward`` will be lost. For example,\\n        if :attr:`module` has a counter attribute that is incremented in each\\n        ``forward``, it will always stay at the initial value because the update\\n        is done on the replicas which are destroyed after ``forward``. However,\\n        :class:`~torch.nn.DataParallel` guarantees that the replica on\\n        ``device[0]`` will have its parameters and buffers sharing storage with\\n        the base parallelized :attr:`module`. So **in-place** updates to the\\n        parameters or buffers on ``device[0]`` will be recorded. E.g.,\\n        :class:`~torch.nn.BatchNorm2d` and :func:`~torch.nn.utils.spectral_norm`\\n        rely on this behavior to update the buffers.\\n\\n    .. warning::\\n        Forward and backward hooks defined on :attr:`module` and its submodules\\n        will be invoked ``len(device_ids)`` times, each with inputs located on\\n        a particular device. Particularly, the hooks are only guaranteed to be\\n        executed in correct order with respect to operations on corresponding\\n        devices. For example, it is not guaranteed that hooks set via\\n        :meth:`~torch.nn.Module.register_forward_pre_hook` be executed before\\n        `all` ``len(device_ids)`` :meth:`~torch.nn.Module.forward` calls, but\\n        that each such hook be executed before the corresponding\\n        :meth:`~torch.nn.Module.forward` call of that device.\\n\\n    .. warning::\\n        When :attr:`module` returns a scalar (i.e., 0-dimensional tensor) in\\n        :func:`forward`, this wrapper will return a vector of length equal to\\n        number of devices used in data parallelism, containing the result from\\n        each device.\\n\\n    .. note::\\n        There is a subtlety in using the\\n        ``pack sequence -> recurrent network -> unpack sequence`` pattern in a\\n        :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.\\n        See :ref:`pack-rnn-unpack-with-data-parallelism` section in FAQ for\\n        details.\\n\\n\\n    Args:\\n        module (Module): module to be parallelized\\n        device_ids (list of int or torch.device): CUDA devices (default: all devices)\\n        output_device (int or torch.device): device location of output (default: device_ids[0])\\n\\n    Attributes:\\n        module (Module): the module to be parallelized\\n\\n    Example::\\n\\n        >>> # xdoctest: +SKIP\\n        >>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])\\n        >>> output = net(input_var)  # input_var can be on any device, including CPU\\n    \\\"\\\"\\\"\\n\\n    # TODO: update notes/cuda.rst when this class handles 8+ GPUs well\\n\\n    def __init__(\\n        self,\\n        module: T,\\n        device_ids: Optional[Sequence[Union[int, torch.device]]] = None,\\n        output_device: Optional[Union[int, torch.device]] = None,\\n        dim: int = 0,\\n    ) -> None:\\n        super().__init__()\\n        torch._C._log_api_usage_once(\\\"torch.nn.parallel.DataParallel\\\")\\n        device_type = _get_available_device_type()\\n        if device_type is None:\\n            self.module = module\\n            self.device_ids = []\\n            return\\n\\n        if device_ids is None:\\n            device_ids = _get_all_device_indices()\\n\\n        if device_ids is None:\\n            raise RuntimeError(\\\"no available devices were found\\\")\\n\\n        if output_device is None:\\n            output_device = device_ids[0]\\n\\n        self.dim = dim\\n        self.module = module\\n        self.device_ids = [_get_device_index(x, True) for x in device_ids]\\n        self.output_device = _get_device_index(output_device, True)\\n        self.src_device_obj = torch.device(device_type, self.device_ids[0])\\n\\n        if device_type == \\\"cuda\\\":\\n            _check_balance(self.device_ids)\\n\\n        if len(self.device_ids) == 1:\\n            self.module.to(self.src_device_obj)\\n\\n    def forward(self, *inputs: Any, **kwargs: Any) -> Any:\\n        with torch.autograd.profiler.record_function(\\\"DataParallel.forward\\\"):\\n            if not self.device_ids:\\n                return self.module(*inputs, **kwargs)\\n\\n            for t in chain(self.module.parameters(), self.module.buffers()):\\n                if t.device != self.src_device_obj:\\n                    raise RuntimeError(\\n                        \\\"module must have its parameters and buffers \\\"\\n                        f\\\"on device {self.src_device_obj} (device_ids[0]) but found one of \\\"\\n                        f\\\"them on device: {t.device}\\\"\\n                    )\\n\\n            inputs, module_kwargs = self.scatter(inputs, kwargs, self.device_ids)\\n            # for forward function without any inputs, empty list and dict will be created\\n            # so the module can be executed on one device which is the first one in device_ids\\n            if not inputs and not module_kwargs:\\n                inputs = ((),)\\n                module_kwargs = ({},)\\n\\n            if len(self.device_ids) == 1:\\n                return self.module(*inputs[0], **module_kwargs[0])\\n            replicas = self.replicate(self.module, self.device_ids[: len(inputs)])\\n            outputs = self.parallel_apply(replicas, inputs, module_kwargs)\\n            return self.gather(outputs, self.output_device)\\n\\n    def replicate(\\n        self, module: T, device_ids: Sequence[Union[int, torch.device]]\\n    ) -> List[T]:\\n        return replicate(module, device_ids, not torch.is_grad_enabled())\\n\\n    def scatter(\\n        self,\\n        inputs: Tuple[Any, ...],\\n        kwargs: Optional[Dict[str, Any]],\\n        device_ids: Sequence[Union[int, torch.device]],\\n    ) -> Any:\\n        return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)\\n\\n    def parallel_apply(\\n        self, replicas: Sequence[T], inputs: Sequence[Any], kwargs: Any\\n    ) -> List[Any]:\\n        return parallel_apply(\\n            replicas, inputs, kwargs, self.device_ids[: len(replicas)]\\n        )\\n\\n    def gather(self, outputs: Any, output_device: Union[int, torch.device]) -> Any:\\n        return gather(outputs, output_device, dim=self.dim)\\n\\n\\ndef data_parallel(\\n    module: Module,\\n    inputs: Any,\\n    device_ids: Optional[Sequence[Union[int, torch.device]]] = None,\\n    output_device: Optional[Union[int, torch.device]] = None,\\n    dim: int = 0,\\n    module_kwargs: Optional[Any] = None,\\n) -> torch.Tensor:\\n    r\\\"\\\"\\\"Evaluate module(input) in parallel across the GPUs given in device_ids.\\n\\n    This is the functional version of the DataParallel module.\\n\\n    Args:\\n        module (Module): the module to evaluate in parallel\\n        inputs (Tensor): inputs to the module\\n        device_ids (list of int or torch.device): GPU ids on which to replicate module\\n        output_device (list of int or torch.device): GPU location of the output  Use -1 to indicate the CPU.\\n            (default: device_ids[0])\\n    Returns:\\n        a Tensor containing the result of module(input) located on\\n        output_device\\n    \\\"\\\"\\\"\\n    if not isinstance(inputs, tuple):\\n        inputs = (inputs,) if inputs is not None else ()\\n\\n    device_type = _get_available_device_type()\\n\\n    if device_type is None:\\n        raise RuntimeError(\\\"device type could not be determined\\\")\\n\\n    if device_ids is None:\\n        device_ids = _get_all_device_indices()\\n\\n    if device_ids is None:\\n        raise RuntimeError(\\\"no available devices were found\\\")\\n\\n    if output_device is None:\\n        output_device = device_ids[0]\\n\\n    device_ids = [_get_device_index(x, True) for x in device_ids]\\n    output_device = _get_device_index(output_device, True)\\n    src_device_obj = torch.device(device_type, device_ids[0])\\n\\n    for t in chain(module.parameters(), module.buffers()):\\n        if t.device != src_device_obj:\\n            raise RuntimeError(\\n                \\\"module must have its parameters and buffers \\\"\\n                f\\\"on device {src_device_obj} (device_ids[0]) but found one of \\\"\\n                f\\\"them on device: {t.device}\\\"\\n            )\\n\\n    inputs, module_kwargs = scatter_kwargs(inputs, module_kwargs, device_ids, dim)\\n    # for module without any inputs, empty list and dict will be created\\n    # so the module can be executed on one device which is the first one in device_ids\\n    if not inputs and not module_kwargs:\\n        inputs = ((),)\\n        module_kwargs = ({},)\\n\\n    assert module_kwargs is not None\\n\\n    if len(device_ids) == 1:\\n        return module(*inputs[0], **module_kwargs[0])\\n    used_device_ids = device_ids[: len(inputs)]\\n    replicas = replicate(module, used_device_ids)\\n    outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)\\n    return gather(outputs, output_device, dim)\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Dynamic Modules.\\n\\nThis package is in the process of being deprecated.\\nPlease, use `torch.ao.nn.qat.dynamic` instead.\\n\\\"\\\"\\\"\\nfrom torch.nn.qat import dynamic, modules  # noqa: F403\\nfrom torch.nn.qat.modules import *  # noqa: F403\\n\\n\\n__all__ = [\\n    \\\"Linear\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"Embedding\\\",\\n    \\\"EmbeddingBag\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/qat`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/qat/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\nfrom torch.ao.nn.qat.modules.linear import Linear\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/qat`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/qat/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.qat.modules.conv import Conv1d, Conv2d, Conv3d\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Modules.\\n\\nThis package is in the process of being deprecated.\\nPlease, use `torch.ao.nn.qat.modules` instead.\\n\\\"\\\"\\\"\\nfrom torch.ao.nn.qat.modules.conv import Conv1d, Conv2d, Conv3d\\nfrom torch.ao.nn.qat.modules.embedding_ops import Embedding, EmbeddingBag\\nfrom torch.ao.nn.qat.modules.linear import Linear\\nfrom torch.nn.qat.modules import conv, embedding_ops, linear\\n\\n\\n__all__ = [\\n    \\\"Linear\\\",\\n    \\\"Conv1d\\\",\\n    \\\"Conv2d\\\",\\n    \\\"Conv3d\\\",\\n    \\\"Embedding\\\",\\n    \\\"EmbeddingBag\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/qat`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/qat/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.qat.modules.embedding_ops import Embedding, EmbeddingBag\\n\\n\\n__all__ = [\\\"Embedding\\\", \\\"EmbeddingBag\\\"]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Dynamic Modules.\\n\\nThis package is in the process of being deprecated.\\nPlease, use `torch.ao.nn.qat.dynamic` instead.\\n\\\"\\\"\\\"\\nfrom torch.nn.qat.dynamic.modules import *  # noqa: F403\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/qat/dynamic`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/qat/dynamic/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\nfrom torch.ao.nn.qat.dynamic.modules.linear import Linear\\n\\n\\nfrom torch.nn.qat.dynamic.modules.linear import Linear\\n\\n\\n__all__ = [\\\"Linear\\\"]\\n\\n\\nfrom torch.ao.nn.intrinsic import (\\n    BNReLU2d,\\n    BNReLU3d,\\n    ConvBn1d,\\n    ConvBn2d,\\n    ConvBn3d,\\n    ConvBnReLU1d,\\n    ConvBnReLU2d,\\n    ConvBnReLU3d,\\n    ConvReLU1d,\\n    ConvReLU2d,\\n    ConvReLU3d,\\n    LinearBn1d,\\n    LinearReLU,\\n)\\nfrom torch.ao.nn.intrinsic.modules.fused import _FusedModule  # noqa: F401\\n\\n# Include the subpackages in case user imports from it directly\\nfrom torch.nn.intrinsic import modules, qat, quantized  # noqa: F401\\n\\n\\n__all__ = [\\n    \\\"ConvBn1d\\\",\\n    \\\"ConvBn2d\\\",\\n    \\\"ConvBn3d\\\",\\n    \\\"ConvBnReLU1d\\\",\\n    \\\"ConvBnReLU2d\\\",\\n    \\\"ConvBnReLU3d\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n    \\\"LinearReLU\\\",\\n    \\\"BNReLU2d\\\",\\n    \\\"BNReLU3d\\\",\\n    \\\"LinearBn1d\\\",\\n]\\n\\n\\n# to ensure customers can use the module below\\n# without importing it directly\\nfrom torch.nn.intrinsic.quantized import dynamic, modules  # noqa: F401\\nfrom torch.nn.intrinsic.quantized.modules import *  # noqa: F403\\n\\n\\n__all__ = [\\n    \\\"BNReLU2d\\\",\\n    \\\"BNReLU3d\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.ao.nn.intrinsic.quantized import ConvReLU1d, ConvReLU2d, ConvReLU3d\\n\\n\\n__all__ = [\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n]\\n\\n\\nfrom torch.ao.nn.intrinsic.quantized import LinearReLU\\n\\n\\n__all__ = [\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.ao.nn.intrinsic.quantized import BNReLU2d, BNReLU3d\\n\\n\\n__all__ = [\\n    \\\"BNReLU2d\\\",\\n    \\\"BNReLU3d\\\",\\n]\\n\\n\\nfrom torch.nn.intrinsic.quantized.modules.bn_relu import BNReLU2d, BNReLU3d\\nfrom torch.nn.intrinsic.quantized.modules.conv_relu import (\\n    ConvReLU1d,\\n    ConvReLU2d,\\n    ConvReLU3d,\\n)\\nfrom torch.nn.intrinsic.quantized.modules.linear_relu import LinearReLU\\n\\n\\n__all__ = [\\n    \\\"LinearReLU\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n    \\\"BNReLU2d\\\",\\n    \\\"BNReLU3d\\\",\\n]\\n\\n\\nfrom torch.nn.intrinsic.quantized.dynamic.modules import *  # noqa: F403\\n\\n\\nfrom torch.ao.nn.intrinsic.quantized.dynamic import LinearReLU\\n\\n\\n__all__ = [\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.nn.intrinsic.quantized.dynamic.modules.linear_relu import LinearReLU\\n\\n\\n__all__ = [\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.ao.nn.intrinsic import (\\n    BNReLU2d,\\n    BNReLU3d,\\n    ConvBn1d,\\n    ConvBn2d,\\n    ConvBn3d,\\n    ConvBnReLU1d,\\n    ConvBnReLU2d,\\n    ConvBnReLU3d,\\n    ConvReLU1d,\\n    ConvReLU2d,\\n    ConvReLU3d,\\n    LinearBn1d,\\n    LinearReLU,\\n)\\nfrom torch.ao.nn.intrinsic.modules.fused import _FusedModule  # noqa: F401\\n\\n\\n__all__ = [\\n    \\\"BNReLU2d\\\",\\n    \\\"BNReLU3d\\\",\\n    \\\"ConvBn1d\\\",\\n    \\\"ConvBn2d\\\",\\n    \\\"ConvBn3d\\\",\\n    \\\"ConvBnReLU1d\\\",\\n    \\\"ConvBnReLU2d\\\",\\n    \\\"ConvBnReLU3d\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n    \\\"LinearBn1d\\\",\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.nn.intrinsic.modules.fused import (\\n    _FusedModule,\\n    BNReLU2d,\\n    BNReLU3d,\\n    ConvBn1d,\\n    ConvBn2d,\\n    ConvBn3d,\\n    ConvBnReLU1d,\\n    ConvBnReLU2d,\\n    ConvBnReLU3d,\\n    ConvReLU1d,\\n    ConvReLU2d,\\n    ConvReLU3d,\\n    LinearBn1d,\\n    LinearReLU,\\n)\\n\\n\\n__all__ = [\\n    \\\"BNReLU2d\\\",\\n    \\\"BNReLU3d\\\",\\n    \\\"ConvBn1d\\\",\\n    \\\"ConvBn2d\\\",\\n    \\\"ConvBn3d\\\",\\n    \\\"ConvBnReLU1d\\\",\\n    \\\"ConvBnReLU2d\\\",\\n    \\\"ConvBnReLU3d\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n    \\\"LinearBn1d\\\",\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.nn.intrinsic.qat.modules import *  # noqa: F403\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Intrinsic QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/intrinsic/qat/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.intrinsic.qat import LinearBn1d\\n\\n\\n__all__ = [\\n    \\\"LinearBn1d\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Intrinsic QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/intrinsic/qat/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.intrinsic.qat import LinearReLU\\n\\n\\n__all__ = [\\n    \\\"LinearReLU\\\",\\n]\\n\\n\\nfrom torch.nn.intrinsic.qat.modules.conv_fused import (\\n    ConvBn1d,\\n    ConvBn2d,\\n    ConvBn3d,\\n    ConvBnReLU1d,\\n    ConvBnReLU2d,\\n    ConvBnReLU3d,\\n    ConvReLU1d,\\n    ConvReLU2d,\\n    ConvReLU3d,\\n    freeze_bn_stats,\\n    update_bn_stats,\\n)\\nfrom torch.nn.intrinsic.qat.modules.linear_fused import LinearBn1d\\nfrom torch.nn.intrinsic.qat.modules.linear_relu import LinearReLU\\n\\n\\n__all__ = [\\n    \\\"LinearReLU\\\",\\n    \\\"LinearBn1d\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvReLU3d\\\",\\n    \\\"ConvBn1d\\\",\\n    \\\"ConvBn2d\\\",\\n    \\\"ConvBn3d\\\",\\n    \\\"ConvBnReLU1d\\\",\\n    \\\"ConvBnReLU2d\\\",\\n    \\\"ConvBnReLU3d\\\",\\n    \\\"update_bn_stats\\\",\\n    \\\"freeze_bn_stats\\\",\\n]\\n\\n\\n# flake8: noqa: F401\\nr\\\"\\\"\\\"Intrinsic QAT Modules.\\n\\nThis file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and\\nis kept here for compatibility while the migration process is ongoing.\\nIf you are adding a new entry/functionality, please, add it to the\\nappropriate file under the `torch/ao/nn/intrinsic/qat/modules`,\\nwhile adding an import statement here.\\n\\\"\\\"\\\"\\n\\nfrom torch.ao.nn.intrinsic.qat import (\\n    ConvBn1d,\\n    ConvBn2d,\\n    ConvBn3d,\\n    ConvBnReLU1d,\\n    ConvBnReLU2d,\\n    ConvBnReLU3d,\\n    ConvReLU1d,\\n    ConvReLU2d,\\n    ConvReLU3d,\\n    freeze_bn_stats,\\n    update_bn_stats,\\n)\\n\\n\\n__all__ = [\\n    # Modules\\n    \\\"ConvBn1d\\\",\\n    \\\"ConvBnReLU1d\\\",\\n    \\\"ConvReLU1d\\\",\\n    \\\"ConvBn2d\\\",\\n    \\\"ConvBnReLU2d\\\",\\n    \\\"ConvReLU2d\\\",\\n    \\\"ConvBn3d\\\",\\n    \\\"ConvBnReLU3d\\\",\\n    \\\"ConvReLU3d\\\",\\n    # Utilities\\n    \\\"freeze_bn_stats\\\",\\n    \\\"update_bn_stats\\\",\\n]\\n\\n\\n# mypy: allow-untyped-defs\\n\\\"\\\"\\\"Defines bias subclasses that work with scaled_dot_product_attention\\\"\\\"\\\"\\nfrom enum import auto, IntEnum\\nfrom typing import Optional\\nfrom warnings import warn\\n\\nimport torch\\nimport torch.nn.functional as F\\nfrom torch.backends.cuda import (\\n    can_use_efficient_attention,\\n    can_use_flash_attention,\\n    is_flash_attention_available,\\n    SDPAParams,\\n)\\nfrom torch.nn.attention import _raise_kernel_warnings\\nfrom torch.nn.attention._utils import (\\n    _calculate_scale,\\n    _input_requires_grad,\\n    _postprocess_flash_output,\\n    _validate_sdpa_input,\\n)\\n\\n\\n__all__ = [\\\"causal_upper_left\\\", \\\"causal_lower_right\\\", \\\"CausalVariant\\\", \\\"CausalBias\\\"]\\n\\n\\ntorch._dynamo.allow_in_graph(is_flash_attention_available)\\ntorch._dynamo.allow_in_graph(can_use_flash_attention)\\ntorch._dynamo.allow_in_graph(can_use_efficient_attention)\\ntorch._dynamo.allow_in_graph(SDPAParams)\\n\\n\\nclass CausalVariant(IntEnum):\\n    r\\\"\\\"\\\"\\n    Enum for causal variants used in attention mechanisms.\\n\\n    Defines two types of causal biases:\\n\\n    `UPPER_LEFT`: Represents upper-left triangular bias for standard causal attention.\\n    The equivalent pytorch code for constructing this bias is:\\n\\n    .. code-block:: python\\n\\n        torch.tril(torch.ones(size, dtype=torch.bool))\\n\\n    For instance, with `shape=(3,4)`, the materialized bias tensor will be:\\n\\n    .. code-block:: text\\n\\n        [[1, 0, 0, 0],\\n         [1, 1, 0, 0],\\n         [1, 1, 1, 0]]\\n\\n\\n    `LOWER_RIGHT`: Represents lower-right triangular bias, the include values are aligned to the lower\\n    right corner of the matrix.\\n\\n    The equivalent pytorch code for constructing this bias is:\\n\\n    .. code-block:: python\\n\\n        diagonal_offset = size[1] - size[0]\\n        torch.tril(\\n            torch.ones(size, dtype=torch.bool),\\n            diagonal=diagonal_offset,\\n        )\\n\\n    For instance, with `shape=(3,4)`, the materialized bias tensor will be:\\n\\n    .. code-block:: text\\n\\n        [[1, 1, 0, 0],\\n         [1, 1, 1, 0],\\n         [1, 1, 1, 1]]\\n\\n    Note that these variants are equivalent to each other when the sequence lengths of the query and key/value\\n    tensors are equal since the triangular matrix is square.\\n\\n    .. warning:: This enum is a prototype and subject to change.\\n    \\\"\\\"\\\"\\n\\n    UPPER_LEFT = auto()\\n    LOWER_RIGHT = auto()\\n\\n\\nclass CausalBias(torch.Tensor):\\n    \\\"\\\"\\\"\\n    A bias representing causal attention patterns. For an overview of the bias structure, see the :class:`CausalVariant` enum.\\n\\n    This class is used for defining causal (triangular) attention biases. For construing the bias, there exist\\n    two factory functions: :func:`causal_upper_left` and :func:`causal_lower_right`.\\n\\n    Example:\\n\\n    .. code-block:: python\\n\\n        from torch.nn.attention.bias import causal_lower_right\\n\\n        bsz, num_heads, seqlen_q, seqlen_kv, head_dim = 32, 8, 4, 12, 8\\n\\n        # Create a lower-right causal bias\\n        attn_bias = causal_lower_right(seqlen_q, seqlen_kv)\\n\\n        q = torch.randn(bsz, num_heads, seqlen_q, head_dim, device=\\\"cuda\\\", dtype=torch.float16)\\n        k = torch.randn(bsz, num_heads, seqlen_kv, head_dim, device=\\\"cuda\\\", dtype=torch.float16)\\n        v = torch.randn(bsz, num_heads, seqlen_kv, head_dim, device=\\\"cuda\\\", dtype=torch.float16)\\n\\n        out = F.scaled_dot_product_attention(q, k, v, attn_bias)\\n\\n    .. warning:: This class is a prototype and subject to change.\\n    \\\"\\\"\\\"\\n\\n    def __init__(self, variant: CausalVariant, seq_len_q: int, seq_len_kv: int):\\n        \\\"\\\"\\\"\\n        Initializes the CausalBias instance with a specified variant and sequence lengths.\\n\\n        Args:\\n            variant (CausalVariant): The type of causal bias to use (either UPPER_LEFT or LOWER_RIGHT).\\n            seq_len_q (int): The sequence length of the query tensor.\\n            seq_len_kv (int): The sequence length of the key/value tensor.\\n\\n        Raises a warning if the LOWER_RIGHT variant is used with seq_len_q > seq_len_kv, as it may produce NaNs.\\n        \\\"\\\"\\\"\\n        assert isinstance(variant, CausalVariant)\\n        self.variant = variant\\n        self.seq_len_q = seq_len_q\\n        self.seq_len_kv = seq_len_kv\\n        if seq_len_q > seq_len_kv and variant == CausalVariant.LOWER_RIGHT:\\n            warn(\\n                \\\"Lower right causal bias will produce NaNs in the output when seq_len_q > seq_len_kv!\\\"\\n            )\\n\\n    def _upper_left(self, device: torch.device) -> torch.Tensor:\\n        \\\"\\\"\\\"Upper left causal bias\\\"\\\"\\\"\\n        return torch.tril(\\n            torch.ones(self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool)\\n        )\\n\\n    def _lower_right(self, device: torch.device) -> torch.Tensor:\\n        \\\"\\\"\\\"Lower right causal bias\\\"\\\"\\\"\\n        diagonal_offset = self.seq_len_kv - self.seq_len_q\\n        return torch.tril(\\n            torch.ones(\\n                self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool\\n            ),\\n            diagonal=diagonal_offset,\\n        )\\n\\n    def _materialize(self, device: Optional[torch.device] = None) -> torch.Tensor:\\n        \\\"\\\"\\\"\\n        Materializes the causal bias into a tensor form.\\n\\n        Depending on the variant, this method generates either an upper-left or lower-right\\n        triangular matrix to represent the causal bias.\\n\\n        Args:\\n            device (Optional[torch.device]): The device on which to create the tensor. Defaults to CPU.\\n\\n        Returns:\\n            torch.Tensor: The materialized bias tensor.\\n        \\\"\\\"\\\"\\n        if device is None:\\n            device = torch.device(\\\"cpu\\\")\\n        if self.variant == CausalVariant.UPPER_LEFT:\\n            return self._upper_left(device)\\n        elif self.variant == CausalVariant.LOWER_RIGHT:\\n            return self._lower_right(device)\\n\\n    @staticmethod\\n    def _dispatch(\\n        query: torch.Tensor,\\n        key: torch.Tensor,\\n        value: torch.Tensor,\\n        attn_mask: \\\"CausalBias\\\",\\n        dropout_p: float = 0.0,\\n        is_causal: bool = False,\\n        scale: Optional[float] = None,\\n        enable_gqa: bool = False,\\n    ) -> torch.Tensor:\\n        r\\\"\\\"\\\"\\n        Handles the logic for computing attention with the specified causal bias.\\n\\n        Args:\\n            query (Tensor): Query tensor; shape :math:`(N, ..., L, E)`.\\n            key (Tensor): Key tensor; shape :math:`(N, ..., S, E)`.\\n            value (Tensor): Value tensor; shape :math:`(N, ..., S, Ev)`.\\n            attn_mask (CausalBias): The type of causal attention to apply.\\n                A boolean mask where a value of True indicates that the element *should* take part in attention.\\n                A float mask of the same type as query, key, value that is added to the attention score.\\n            dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied\\n            is_causal (bool): If true, assumes upper left causal attention masking and errors if both attn_mask and is_causal\\n                are set.\\n            scale (optional float): Scaling factor applied prior to softmax. If None, the default value is set\\n                to :math:`\\\\frac{1}{\\\\sqrt{E}}`.\\n            enable_gqa (optional bool): If set to True, Grouped Query Attention (GQA) is enabled, by default it is set to False.\\n\\n        Returns:\\n            output (Tensor): Attention output; shape :math:`(N, ..., L, Ev)`.\\n\\n        Raises:\\n            ValueError: If the causal bias variant is not a CausalVariant type.\\n\\n        \\\"\\\"\\\"\\n        if is_causal:\\n            raise ValueError(\\\"CausalBias should not be used with causal=True\\\")\\n\\n        if (\\n            attn_mask.seq_len_q == attn_mask.seq_len_kv\\n            or attn_mask.variant == CausalVariant.UPPER_LEFT\\n        ):\\n            return F.scaled_dot_product_attention(\\n                query,\\n                key,\\n                value,\\n                attn_mask=None,\\n                dropout_p=dropout_p,\\n                is_causal=True,\\n                scale=scale,\\n                enable_gqa=enable_gqa,\\n            )\\n        elif attn_mask.variant == CausalVariant.LOWER_RIGHT:\\n            _validate_sdpa_input(query, key, value, None, dropout_p, is_causal, scale)\\n            sdpa_params = SDPAParams(\\n                query, key, value, None, dropout_p, is_causal, enable_gqa\\n            )\\n            if can_use_flash_attention(sdpa_params):\\n                needs_padding = query.size(-1) % 8 != 0\\n                og_head_size = query.size(-1)\\n                og_scale = _calculate_scale(og_head_size, scale)\\n                if needs_padding:\\n                    query = torch.nn.functional.pad(query, (0, 8 - query.size(-1) % 8))\\n                    key = torch.nn.functional.pad(key, (0, 8 - key.size(-1) % 8))\\n                    value = torch.nn.functional.pad(value, (0, 8 - value.size(-1) % 8))\\n                out = torch.ops.aten._scaled_dot_product_flash_attention(\\n                    query,\\n                    key,\\n                    value,\\n                    dropout_p,\\n                    is_causal=True,  # TODO: Flash accepts causal = True and for this particular op it means lower right\\n                    return_debug_mask=False,\\n                    scale=og_scale,\\n                )[0]\\n                return _postprocess_flash_output(out, og_head_size)\\n            if can_use_efficient_attention(sdpa_params):\\n                compute_log_sumexp = False\\n                if _input_requires_grad(query, key, value):\\n                    compute_log_sumexp = True\\n                return torch.ops.aten._efficient_attention_forward(\\n                    query.transpose(1, 2),\\n                    key.transpose(1, 2),\\n                    value.transpose(1, 2),\\n                    bias=None,\\n                    cu_seqlens_q=None,\\n                    cu_seqlens_k=None,\\n                    max_seqlen_q=None,\\n                    max_seqlen_k=None,\\n                    dropout_p=dropout_p,\\n                    custom_mask_type=int(attn_mask.variant),\\n                    compute_log_sumexp=compute_log_sumexp,\\n                    scale=scale,\\n                    seqlen_k=None,\\n                )[0].transpose(1, 2)\\n            else:\\n                _raise_kernel_warnings(sdpa_params)\\n                # We cant use efficient attention the only support for lower right is via materialization\\n                return F.scaled_dot_product_attention(\\n                    query,\\n                    key,\\n                    value,\\n                    attn_mask=attn_mask._materialize(query.device),\\n                    dropout_p=dropout_p,\\n                    is_causal=False,\\n                    scale=scale,\\n                    enable_gqa=enable_gqa,\\n                )\\n        else:\\n            raise ValueError(\\n                f\\\"CausalBias.variant must be a CausalVariant type, but found: {attn_mask.variant}\\\"\\n            )\\n\\n    @classmethod\\n    def __torch_function__(cls, func, types, args=(), kwargs=None):\\n        \\\"\\\"\\\"Defines the behavior of torch.nn.functional.scaled_dot_product_attention when the attn_bias is an AttnBias\\\"\\\"\\\"\\n        if kwargs is None:\\n            kwargs = {}\\n        if func != torch.nn.functional.scaled_dot_product_attention:\\n            raise NotImplementedError(\\n                \\\"CausalBias only supports scaled_dot_product_attention\\\"\\n            )\\n        return cls._dispatch(*args, **kwargs)\\n\\n    def __repr__(self):\\n        return self._materialize().__repr__()\\n\\n\\ndef causal_upper_left(*size) -> CausalBias:\\n    \\\"\\\"\\\"\\n    Creates an upper-left triangular causal bias.\\n\\n    This function generates a upper-left triangular matrix to represent causal attention bias with a\\n    diagonal offset set so that the inclusive values are aligned to the upper left corner of the matrix.\\n    This equivalent to the `is_causal=True` argument in `scaled_dot_product_attention`.\\n\\n    The equivalent pytorch code for constructing this bias is:\\n\\n    .. code-block:: python\\n\\n        torch.tril(torch.ones(size, dtype=torch.bool))\\n\\n    For instance, with `shape=(3,4)`, the materialized bias tensor will be:\\n\\n    .. code-block:: text\\n\\n        [[1, 0, 0, 0],\\n         [1, 1, 0, 0],\\n         [1, 1, 1, 0]]\\n\\n    Args:\\n        size: The size of the bias matrix.\\n\\n    Returns:\\n        CausalBias: The UPPER_LEFT triangular causal bias variant.\\n    \\\"\\\"\\\"\\n    assert len(size) == 2, \\\"causal_upper_left only supports 2D tensors\\\"\\n    seq_len_q, seq_len_kv = size\\n    return CausalBias(CausalVariant.UPPER_LEFT, seq_len_q, seq_len_kv)\\n\\n\\ndef causal_lower_right(*size) -> CausalBias:\\n    \\\"\\\"\\\"\\n    Creates a lower-right triangular causal bias.\\n\\n    This function generates a lower-right triangular matrix to represent causal attention bias with a\\n    diagonal offset set so that the inclusive values are aligned to the lower right corner of the matrix.\\n\\n    The equivalent pytorch code for constructing this bias is:\\n\\n    .. code-block:: python\\n\\n        diagonal_offset = size[1] - size[0]\\n        torch.tril(\\n            torch.ones(size, dtype=torch.bool),\\n            diagonal=diagonal_offset,\\n        )\\n\\n    For instance, with `shape=(3,4)`, the materialized bias tensor will be:\\n\\n    .. code-block:: text\\n\\n        [[1, 1, 0, 0],\\n         [1, 1, 1, 0],\\n         [1, 1, 1, 1]]\\n\\n    Args:\\n        size: The size of the bias matrix.\\n\\n    Returns:\\n        CausalBias: The LOWER_RIGHT triangular causal bias variant.\\n    \\\"\\\"\\\"\\n    assert len(size) == 2, \\\"causal_lower_right only supports 2D tensors\\\"\\n    seq_len_q, seq_len_kv = size\\n    return CausalBias(CausalVariant.LOWER_RIGHT, seq_len_q, seq_len_kv)\\n\\n\\n# mypy: allow-untyped-defs\\n\\\"\\\"\\\"Defines utilities for interacting with scaled_dot_product_attention\\\"\\\"\\\"\\nimport math\\nfrom typing import List, Optional, Union\\n\\nimport torch\\n\\n\\n__all__: List[str] = []\\n\\n\\ndef _input_requires_grad(*tensors: torch.Tensor) -> bool:\\n    \\\"\\\"\\\"Returns True if any of the tensors requires grad\\\"\\\"\\\"\\n    return any(t.requires_grad for t in tensors)\\n\\n\\ndef _postprocess_flash_output(inpt_tensor: torch.Tensor, og_size: int) -> torch.Tensor:\\n    \\\"\\\"\\\"Handles the unpad of the last dimension\\\"\\\"\\\"\\n    if inpt_tensor.size(-1) != og_size:\\n        return inpt_tensor[..., :og_size]\\n    return inpt_tensor\\n\\n\\ndef _calculate_scale(head_dim_size: int, scale: Optional[float]) -> float:\\n    \\\"\\\"\\\"\\n    For FlashAttention we pad the head dimension to be a multiple of 8 so we need to scale the output\\n    by the original head size and not the padded.\\n    \\\"\\\"\\\"\\n    if scale is not None:\\n        return scale\\n    return 1.0 / math.sqrt(head_dim_size)\\n\\n\\n_SUPPORTED_HEAD_DIMS = [2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]\\n\\n\\ndef _supported_head_dim(n: Union[int, torch.SymInt]) -> bool:\\n    \\\"\\\"\\\"Returns true if the head dim is supported by FlexAttention\\\"\\\"\\\"\\n    return n in _SUPPORTED_HEAD_DIMS\\n\\n\\ndef _validate_sdpa_input(\\n    query: torch.Tensor,\\n    key: torch.Tensor,\\n    value: torch.Tensor,\\n    attn_mask: Optional[torch.Tensor] = None,\\n    dropout_p=0.0,\\n    is_causal=False,\\n    scale=None,\\n):\\n    if query.dtype != key.dtype or query.dtype != value.dtype:\\n        raise ValueError(\\n            f\\\"Expected query, key, and value to have the same dtype, \\\"\\n            f\\\"but got query.dtype: {query.dtype}, key.dtype: {key.dtype}, \\\"\\n            f\\\"and value.dtype: {value.dtype} instead.\\\"\\n        )\\n    if query.device != key.device or query.device != value.device:\\n        raise ValueError(\\n            f\\\"Expected query, key, and value to have the same device type, \\\"\\n            f\\\"but got query.device: {query.device}, key.device: {key.device}, \\\"\\n            f\\\"and value.device: {value.device} instead.\\\"\\n        )\\n    if query.dim() < 2 or key.dim() < 2 or value.dim() < 2:\\n        raise ValueError(\\n            f\\\"Expected query, key, and value to all be  at least 2 dimensional, but got query.dim: \\\"\\n            f\\\"{query.dim()}, key.dim: {key.dim()} and value.dim: {value.dim()} instead.\\\"\\n        )\\n\\n\\n# mypy: allow-untyped-decorators\\n# mypy: allow-untyped-defs\\n# flake8: noqa C101\\n\\\"\\\"\\\"This module implements the user facing API for flex_attention in PyTorch.\\\"\\\"\\\"\\nimport functools\\nimport inspect\\nimport itertools\\nimport math\\nimport operator\\nfrom contextlib import nullcontext\\nfrom enum import Enum\\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\\n\\nimport torch\\nfrom torch import Tensor\\nfrom torch._higher_order_ops.flex_attention import (\\n    flex_attention as flex_attention_hop,\\n    TransformGetItemToIndex,\\n)\\nfrom torch._higher_order_ops.utils import _set_compilation_env\\nfrom torch.fx.experimental.proxy_tensor import (\\n    _temp_remove_pre_dispatch_torch_function_mode,\\n)\\nfrom torch.nn.attention._utils import _supported_head_dim, _validate_sdpa_input\\nfrom torch.utils._pytree import tree_map_only\\n\\n\\n__all__ = [\\n    \\\"BlockMask\\\",\\n    \\\"flex_attention\\\",\\n    \\\"create_block_mask\\\",\\n    \\\"create_mask\\\",\\n    \\\"or_masks\\\",\\n    \\\"and_masks\\\",\\n    \\\"noop_mask\\\",\\n]\\n\\n_score_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor, Tensor], Tensor]\\n_mask_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]\\n\\n\\nclass _ModificationType(Enum):\\n    \\\"\\\"\\\"Enum for the type of modification function.\\n    - SCORE_MOD: score_mod function which accepts a score as the first argument\\n    - mask_mod: mask function which does not accept a score and is only used for generating\\n    block mask\\n    \\\"\\\"\\\"\\n\\n    SCORE_MOD = 1\\n    MASK_MOD = 2\\n    UNKNOWN = 3\\n\\n\\ndef _get_mod_type(fn: Callable) -> _ModificationType:\\n    \\\"\\\"\\\"Get the type of modification function.\\n    This function inspects the number of positional arguments of the function to determine\\n    the type of modification function. If the function has 5 positional arguments, it is\\n    considered as a score_mod function. If the function has 4 positional arguments, it is\\n    considered as a mask function.\\n    \\\"\\\"\\\"\\n    num_positional_args = sum(\\n        1\\n        for param in inspect.signature(fn).parameters.values()\\n        if param.default == inspect.Parameter.empty\\n    )\\n    assert num_positional_args == 5 or num_positional_args == 4\\n    if num_positional_args == 5:\\n        return _ModificationType.SCORE_MOD\\n    elif num_positional_args == 4:\\n        return _ModificationType.MASK_MOD\\n    else:\\n        return _ModificationType.UNKNOWN\\n\\n\\n# Need to define it here so that Dynamo doesn't skip it\\ndef _vmap_for_bhqkv(\\n    fn: Callable,\\n    prefix: Tuple[Optional[int], ...],\\n    suffix: Tuple[Optional[int], ...] = (),\\n    out_dims: Union[int, List[Optional[int]]] = 0,\\n    group_dim: bool = False,\\n):\\n    \\\"\\\"\\\"Used to vmap both score_mods and mask_mods over 4-dimensional/5-dimension inputs.\\n    Mapping over the [b, hq, q_idx, kv_idx] or [b, hkv, g, q_idx, kv_idx] dimensions.\\n\\n    Args:\\n        fn (callable): The function to vmap.\\n        prefix (tuple): The prefix of the vmap. For score mod functions,\\n                        this should be set to (0,). For mask_mods = ()\\n        suffix (tuple): We need to add (0,) if gradOut is being mapped over,\\n                        and (None,) * len(other_buffers).\\n        out_dims (tuple): For forward cases, keep this as the default 0 since\\n                          we are only returning 1 output. For backwards, the joint\\n                          graph returns grads for B, H, Q_idx, KV_idx and other_buffers,\\n                          so we set this to (0, None, None, None, None) + (None,) * len(other_buffers).\\n\\n    Returns:\\n        callable: The vmapped function.\\n    \\\"\\\"\\\"\\n    # We vamp a function 4 times, broadcasting the [b, h, q_idx, kv_idx] dimensions\\n    dimensions: List[Tuple[None | int, None | int, None | int, None | int]] = []\\n    dimensions = [\\n        (None, None, None, 0),\\n        (None, None, 0, None),\\n        (None, 0, None, None),\\n    ]\\n\\n    if group_dim:\\n        dimensions += [\\n            (None, 0, None, None),\\n        ]\\n\\n    dimensions += [\\n        (0, None, None, None),\\n    ]\\n\\n    for dims in dimensions:\\n        fn = torch.vmap(fn, in_dims=prefix + dims + suffix, out_dims=out_dims)\\n    return fn\\n\\n\\ndef _identity(\\n    score: Tensor,\\n    batch: Tensor,\\n    head: Tensor,\\n    token_q: Tensor,\\n    token_kv: Tensor,\\n) -> Tensor:\\n    return score\\n\\n\\ndef noop_mask(\\n    batch: Tensor,\\n    head: Tensor,\\n    token_q: Tensor,\\n    token_kv: Tensor,\\n) -> Tensor:\\n    \\\"\\\"\\\"Returns a noop mask_mod\\\"\\\"\\\"\\n    return batch.new_ones(size=(), dtype=torch.bool, device=batch.device)\\n\\n\\n_DEFAULT_SPARSE_BLOCK_SIZE = 128\\n_LARGE_SPARSE_BLOCK_SIZE = 1 << 30\\n\\n\\ndef _ordered_to_dense(num_blocks_in_row: Tensor, col_indices: Tensor):\\n    num_rows = col_indices.shape[-2]\\n    num_cols = col_indices.shape[-1]\\n    batch_dims = num_blocks_in_row.shape[:-1]\\n    device = num_blocks_in_row.device\\n\\n    def create_dense_one(kv_num_blocks, kv_indices):\\n        dense_mask = kv_indices.new_zeros(num_rows, num_cols + 1, dtype=torch.int32)\\n\\n        row_indices = torch.arange(num_rows, dtype=torch.int, device=device).unsqueeze(\\n            -1\\n        )\\n        col_range = torch.arange(num_cols, dtype=torch.int, device=device)\\n        index_mask = col_range < kv_num_blocks.unsqueeze(-1)\\n\\n        # We write to one spot \\\"out of bounds\\\"\\n        valid_indices = torch.where(index_mask, kv_indices, num_cols)\\n\\n        # set the values in 'a' to 1 where the indices are valid\\n        dense_mask[row_indices, valid_indices] = 1\\n        return dense_mask[:, :num_cols].contiguous()\\n\\n    create_dense_batched = create_dense_one\\n    for _ in range(len(batch_dims)):\\n        create_dense_batched = torch.vmap(create_dense_batched, in_dims=(0, 0))\\n\\n    out = create_dense_batched(num_blocks_in_row, col_indices)\\n    return out\\n\\n\\ndef _dense_to_ordered(dense_mask) -> Tuple:\\n    dense_mask = dense_mask.to(dtype=torch.int32)\\n    num_blocks_in_row = dense_mask.sum(dim=-1)\\n    col_indices = torch.argsort(dense_mask, dim=-1, descending=True, stable=True)\\n    return (\\n        num_blocks_in_row.to(torch.int32).contiguous(),\\n        col_indices.to(torch.int32).contiguous(),\\n    )\\n\\n\\ndef _transpose_ordered(num_blocks_in_row: Tensor, col_indices: Tensor):\\n    dense = _ordered_to_dense(num_blocks_in_row, col_indices)\\n    return _dense_to_ordered(dense.transpose(-2, -1))\\n\\n\\nclass BlockMask:\\n    r\\\"\\\"\\\"\\n    BlockMask is our format for representing a block-sparse attention mask.\\n    It is somewhat of a cross in-between BCSR and a non-sparse format.\\n\\n    Basics\\n    ------\\n    A block-sparse mask means that instead of representing the sparsity of\\n    individual elements in the mask, a KV_BLOCK_SIZE x Q_BLOCK_SIZE block is\\n    considered sparse only if every element within that block is sparse.\\n    This aligns well with hardware, which generally expects to perform\\n    contiguous loads and computation.\\n\\n    This format is primarily optimized for 1. simplicity, and 2. kernel\\n    efficiency. Notably, it is *not* optimized for size, as this mask is always\\n    reduced by a factor of KV_BLOCK_SIZE * Q_BLOCK_SIZE. If the size is a\\n    concern, the tensors can be reduced in size by increasing the block size.\\n\\n    The essentials of our format are:\\n\\n    num_blocks_in_row: Tensor[ROWS]:\\n    Describes the number of blocks present in each row.\\n\\n    col_indices: Tensor[ROWS, MAX_BLOCKS_IN_COL]:\\n    `col_indices[i]` is the sequence of block positions for row i. The values of\\n    this row after `col_indices[i][num_blocks_in_row[i]]` are undefined.\\n\\n    For example, to reconstruct the original tensor from this format:\\n\\n    .. code-block:: python\\n\\n        dense_mask = torch.zeros(ROWS, COLS)\\n        for row in range(ROWS):\\n            for block_idx in range(num_blocks_in_row[row]):\\n                dense_mask[row, col_indices[row, block_idx]] = 1\\n\\n    Notably, this format makes it easier to implement a reduction along the\\n    *rows* of the mask.\\n\\n    Details\\n    -------\\n    The basics of our format require only kv_num_blocks and kv_indices. But, we\\n    have up to 8 tensors on this object. This represents 4 pairs:\\n\\n    1. (kv_num_blocks, kv_indices): Used for the forwards pass of attention, as\\n    we reduce along the KV dimension.\\n\\n    2. [OPTIONAL] (full_kv_num_blocks, full_kv_indices): This is optional and\\n    purely an optimization. As it turns out, applying masking to every block\\n    is quite expensive! If we specifically know which blocks are \\\"full\\\" and\\n    don't require masking at all, then we can skip applying mask_mod to these\\n    blocks. This requires the user to split out a separate mask_mod from the\\n    score_mod. For causal masks, this is about a 15% speedup.\\n\\n    3. [GENERATED] (q_num_blocks, q_indices): Required for the backwards pass,\\n    as computing dKV requires iterating along the mask along the Q dimension. These are autogenerated from 1.\\n\\n    4. [GENERATED] (full_q_num_blocks, full_q_indices): Same as above, but for\\n    the backwards pass. These are autogenerated from 2.\\n    \\\"\\\"\\\"\\n    kv_num_blocks: Tensor\\n    kv_indices: Tensor\\n    full_kv_num_blocks: Optional[Tensor]\\n    full_kv_indices: Optional[Tensor]\\n    q_num_blocks: Optional[Tensor]\\n    q_indices: Optional[Tensor]\\n    full_q_num_blocks: Optional[Tensor]\\n    full_q_indices: Optional[Tensor]\\n    BLOCK_SIZE: Tuple[int, int]\\n    mask_mod: _mask_mod_signature\\n\\n    def __init__(\\n        self,\\n        kv_num_blocks: Tensor,\\n        kv_indices: Tensor,\\n        full_kv_num_blocks: Optional[Tensor],\\n        full_kv_indices: Optional[Tensor],\\n        q_num_blocks: Optional[Tensor],\\n        q_indices: Optional[Tensor],\\n        full_q_num_blocks: Optional[Tensor],\\n        full_q_indices: Optional[Tensor],\\n        BLOCK_SIZE: Tuple[int, int],\\n        mask_mod: _mask_mod_signature,\\n    ):\\n        if kv_indices.dim() < 2:\\n            raise RuntimeError(\\\"BlockMask must have at least 2 dimensions\\\")\\n        assert kv_num_blocks is not None, \\\"kv_num_blocks must be provided\\\"\\n        assert kv_indices is not None, \\\"kv_indices must be provided\\\"\\n        assert q_num_blocks is not None, \\\"q_num_blocks must be provided\\\"\\n        assert q_indices is not None, \\\"q_indices must be provided\\\"\\n        assert (full_kv_num_blocks is None) == (\\n            full_kv_indices is None\\n        ), \\\"full_kv_num_blocks and full_kv_indices must be both provided or omitted\\\"\\n        assert (full_q_num_blocks is None) == (\\n            full_q_indices is None\\n        ), \\\"full_q_num_blocks and full_q_indices must be both provided or omitted\\\"\\n\\n        self.kv_num_blocks = kv_num_blocks\\n        self.kv_indices = kv_indices\\n        self.full_kv_num_blocks = full_kv_num_blocks\\n        self.full_kv_indices = full_kv_indices\\n        self.q_num_blocks = q_num_blocks\\n        self.q_indices = q_indices\\n        self.full_q_num_blocks = full_q_num_blocks\\n        self.full_q_indices = full_q_indices\\n        self.BLOCK_SIZE = BLOCK_SIZE\\n        self.mask_mod = mask_mod\\n\\n    @classmethod\\n    def from_kv_blocks(\\n        cls,\\n        kv_num_blocks: Tensor,\\n        kv_indices: Tensor,\\n        full_kv_num_blocks: Optional[Tensor] = None,\\n        full_kv_indices: Optional[Tensor] = None,\\n        BLOCK_SIZE: Union[int, Tuple[int, int]] = _DEFAULT_SPARSE_BLOCK_SIZE,\\n        mask_mod: Optional[_mask_mod_signature] = None,\\n    ):\\n        \\\"\\\"\\\"\\n        Creates a BlockMask instance from key-value block information.\\n\\n        Args:\\n            kv_num_blocks (Tensor): Number of kv_blocks in each Q_BLOCK_SIZE row tile.\\n            kv_indices (Tensor): Indices of key-value blocks in each Q_BLOCK_SIZE row tile.\\n            full_kv_num_blocks (Optional[Tensor]): Number of full kv_blocks in each Q_BLOCK_SIZE row tile.\\n            full_kv_indices (Optional[Tensor]): Indices of full key-value blocks in each Q_BLOCK_SIZE row tile.\\n            BLOCK_SIZE (Union[int, Tuple[int, int]]): Size of KV_BLOCK_SIZE x Q_BLOCK_SIZE tiles.\\n            mask_mod (Optional[Callable]): Function to modify the mask.\\n\\n        Returns:\\n            BlockMask: Instance with full Q information generated via _transposed_ordered\\n\\n        Raises:\\n            RuntimeError: If kv_indices has < 2 dimensions.\\n            AssertionError: If only one of full_kv_* args is provided.\\n        \\\"\\\"\\\"\\n        if kv_indices.dim() < 2:\\n            raise RuntimeError(\\\"BlockMask must have at least 2 dimensions\\\")\\n\\n        assert (full_kv_num_blocks is None) == (\\n            full_kv_indices is None\\n        ), \\\"full_kv_num_blocks and full_kv_indices must be both provided or omitted\\\"\\n\\n        # Generate q_num_blocks and q_indices\\n        q_num_blocks, q_indices = _transpose_ordered(kv_num_blocks, kv_indices)\\n        if full_kv_num_blocks is not None:\\n            assert full_kv_indices is not None\\n            full_q_num_blocks, full_q_indices = _transpose_ordered(\\n                full_kv_num_blocks, full_kv_indices\\n            )\\n        else:\\n            full_q_num_blocks, full_q_indices = None, None\\n\\n        if isinstance(BLOCK_SIZE, int):\\n            BLOCK_SIZE = (BLOCK_SIZE, BLOCK_SIZE)\\n\\n        mask_mod = mask_mod if mask_mod is not None else noop_mask\\n\\n        return cls(\\n            kv_num_blocks=kv_num_blocks,\\n            kv_indices=kv_indices,\\n            full_kv_num_blocks=full_kv_num_blocks,\\n            full_kv_indices=full_kv_indices,\\n            q_num_blocks=q_num_blocks,\\n            q_indices=q_indices,\\n            full_q_num_blocks=full_q_num_blocks,\\n            full_q_indices=full_q_indices,\\n            BLOCK_SIZE=BLOCK_SIZE,\\n            mask_mod=mask_mod,\\n        )\\n\\n    def as_tuple(self, flatten: bool = True):\\n        \\\"\\\"\\\"\\n        Returns a tuple of the attributes of the BlockMask.\\n\\n        Args:\\n            flatten (bool): If True, it will flatten the tuple of (KV_BLOCK_SIZE, Q_BLOCK_SIZE)\\n        \\\"\\\"\\\"\\n        block_size = (\\n            (self.BLOCK_SIZE[0], self.BLOCK_SIZE[1]) if flatten else (self.BLOCK_SIZE,)\\n        )\\n\\n        return (\\n            self.kv_num_blocks,\\n            self.kv_indices,\\n            self.full_kv_num_blocks,\\n            self.full_kv_indices,\\n            self.q_num_blocks,\\n            self.q_indices,\\n            self.full_q_num_blocks,\\n            self.full_q_indices,\\n            *block_size,\\n            self.mask_mod,\\n        )\\n\\n    def __str__(self):\\n        s = f\\\"BlockMask(shape={self.shape}, sparsity={self.sparsity():.2f}%, \\\\n\\\"\\n        mask_str = self.to_string().strip()\\n        s += mask_str\\n        s += \\\"\\\\n)\\\"\\n        return s\\n\\n    def __getitem__(self, index) -> \\\"BlockMask\\\":\\n        \\\"\\\"\\\"\\n        Returns a new BlockMask instance by getting the mask for the given index position.\\n\\n        Args:\\n            index: Index to apply to all attributes.\\n\\n        Example Usage:\\n            .. code-block:: python\\n\\n                def causal_mask(b, h, q_idx, kv_idx):\\n                    return q_idx >= kv_idx\\n\\n                block_mask = create_block_mask(causal_mask, 4, 2, 512, 512, device=\\\"cuda\\\")\\n                assert block_mask.kv_num_blocks.shape == (4,2,4)\\n                assert block_mask.kv_indices.shape == (4,2,4,4)\\n\\n                # Index on batch dimension\\n                new_block_mask = block_mask[0]\\n                assert new_block_mask.kv_num_blocks.shape == (2,4)\\n                assert new_block_mask.kv_indices.shape == (2,4,4)\\n\\n                # Index on batch and head dimension\\n                new_block_mask = block_mask[0, 1]\\n                assert new_block_mask.kv_num_blocks.shape == (4,)\\n                assert new_block_mask.kv_indices.shape == (4,4)\\n\\n                # slicing on batch and head dimension\\n                new_block_mask = block_mask[0:2, 1:2]\\n                assert new_block_mask.kv_num_blocks.shape == (2,1,4)\\n                assert new_block_mask.kv_indices.shape == (2,1,4,4)\\n\\n                # slicing on batch, head, and query dimension\\n                new_block_mask = block_mask[0:2, 1:2, torch.tensor([1], dtype=torch.int32)]\\n                assert new_block_mask.kv_num_blocks.shape == (2,1,1)\\n                assert new_block_mask.kv_indices.shape == (2,1,1,4)\\n        \\\"\\\"\\\"\\n        new_kv_num_blocks = self.kv_num_blocks[index]\\n        new_kv_indices = self.kv_indices[index]\\n        if self.full_kv_num_blocks is not None:\\n            assert self.full_kv_indices is not None\\n            new_full_kv_num_blocks = self.full_kv_num_blocks[index]\\n            new_full_kv_indices = self.full_kv_indices[index]\\n        else:\\n            new_full_kv_num_blocks = None\\n            new_full_kv_indices = None\\n        return BlockMask.from_kv_blocks(\\n            new_kv_num_blocks,\\n            new_kv_indices,\\n            new_full_kv_num_blocks,\\n            new_full_kv_indices,\\n            BLOCK_SIZE=self.BLOCK_SIZE,\\n            mask_mod=None,\\n        )\\n\\n    def __repr__(self):\\n        def shape_or_none(x: Optional[torch.Tensor]):\\n            return x.shape if x is not None else None\\n\\n        return (\\n            f\\\"BlockMask(\\\\n\\\"\\n            f\\\"    kv_num_blocks={self.kv_num_blocks.shape},\\\\n\\\"\\n            f\\\"    kv_indices={self.kv_indices.shape},\\\\n\\\"\\n            f\\\"    full_kv_num_blocks={shape_or_none(self.full_kv_num_blocks )},\\\\n\\\"\\n            f\\\"    full_kv_indices={shape_or_none(self.full_kv_indices)},\\\\n\\\"\\n            f\\\"    q_num_blocks={shape_or_none(self.q_num_blocks)},\\\\n\\\"\\n            f\\\"    q_indices={shape_or_none(self.q_indices)},\\\\n\\\"\\n            f\\\"    full_q_num_blocks={shape_or_none(self.full_q_num_blocks)},\\\\n\\\"\\n            f\\\"    full_q_indices={shape_or_none(self.full_q_indices)},\\\\n\\\"\\n            f\\\"    BLOCK_SIZE={self.BLOCK_SIZE},\\\\n\\\"\\n            f\\\"    shape={self.shape},\\\\n\\\"\\n            f\\\"    sparsity={self.sparsity():.2f}%,\\\\n\\\"\\n            f\\\"    mask_mod={self.mask_mod.__name__ if hasattr(self.mask_mod, '__name__') else self.mask_mod}\\\\n\\\"\\n            f\\\")\\\"\\n        )\\n\\n    @property\\n    def shape(self):\\n        \\\"\\\"\\\"Returns the shape of the mask.\\\"\\\"\\\"\\n        *batch_dims, q_length, _ = self.kv_indices.shape\\n        q_length = self.kv_indices.shape[-2] * self.BLOCK_SIZE[0]\\n        kv_length = self.kv_indices.shape[-1] * self.BLOCK_SIZE[1]\\n        return tuple(batch_dims + [q_length, kv_length])\\n\\n    def numel(self):\\n        \\\"\\\"\\\"Returns the number of elements (not accounting for sparsity) in the mask.\\\"\\\"\\\"\\n        shape = self.shape\\n\\n        def _prod(xs):\\n            return functools.reduce(operator.mul, xs, 1)\\n\\n        return _prod(shape)\\n\\n    def sparsity(self) -> float:\\n        \\\"\\\"\\\"Computes the percentage of blocks that are sparse (i.e. not computed)\\\"\\\"\\\"\\n        total_size = self.numel()\\n        computed_blocks = self.kv_num_blocks.sum()\\n        if self.full_kv_num_blocks is not None:\\n            computed_blocks += self.full_kv_num_blocks.sum()\\n\\n        computed_size = computed_blocks.item() * self.BLOCK_SIZE[0] * self.BLOCK_SIZE[1]\\n        dense_ratio = computed_size / total_size\\n        return 100 * (1 - dense_ratio)\\n\\n    def to_dense(self) -> Tensor:\\n        \\\"\\\"\\\"Returns a dense block that is equivalent to the block mask.\\\"\\\"\\\"\\n        partial_dense = _ordered_to_dense(self.kv_num_blocks, self.kv_indices)\\n        if self.full_kv_num_blocks is not None:\\n            assert self.full_kv_indices is not None\\n            return partial_dense | _ordered_to_dense(\\n                self.full_kv_num_blocks, self.full_kv_indices\\n            )\\n        return partial_dense\\n\\n    def to_string(self, grid_size=(20, 20), limit=4):\\n        \\\"\\\"\\\"Returns a string representation of the block mask. Quite nifty.\\n\\n        If grid_size is None, prints out an uncompressed version. Warning, it can be quite big!\\n        \\\"\\\"\\\"\\n        dense_mask = self.to_dense()\\n        *batch_dims, num_rows, num_cols = dense_mask.shape\\n        if isinstance(grid_size, int):\\n            max_rows = grid_size\\n            max_cols = grid_size\\n        elif grid_size == -1:\\n            max_rows = num_rows\\n            max_cols = num_cols\\n        else:\\n            max_rows, max_cols = grid_size\\n\\n        def create_block_vis(*batch_idx):\\n            descriptors = []\\n\\n            descriptors.append(f\\\"{batch_idx}\\\")\\n\\n            vis = \\\", \\\".join(reversed(descriptors)) + \\\"\\\\n\\\"\\n\\n            def summarize_section(section):\\n                percentage = section.float().mean().item()\\n                if percentage == 1:\\n                    return \\\"█\\\"\\n                elif percentage == 0:\\n                    return \\\" \\\"\\n                else:\\n                    return \\\"░\\\"\\n\\n            def cdiv(a, b):\\n                return (a + (b - 1)) // b\\n\\n            row_step = max(1, cdiv(num_rows, max_rows))\\n            col_step = max(1, cdiv(num_cols, max_cols))\\n\\n            for r in range(0, num_rows, row_step):\\n                for c in range(0, num_cols, col_step):\\n                    cur_mask = dense_mask\\n                    for idx in batch_idx:\\n                        cur_mask = cur_mask[idx]\\n                    char = summarize_section(\\n                        cur_mask[r : r + row_step, c : c + col_step]\\n                    )\\n                    vis += char * 2\\n                vis += \\\"\\\\n\\\"\\n            return vis\\n\\n        total_vis = []\\n        for idx, batch_idx in enumerate(\\n            itertools.product(*[range(i) for i in batch_dims])\\n        ):\\n            if idx == limit:\\n                total_vis.append(\\\"...\\\")\\n                total_vis.append(\\\"To print out more, set BlockMask.to_string(limit=N)\\\")\\n                total_vis.append(\\n                    \\\"You can also index (BlockMask[batch, head]) to choose a specific batch or head\\\"\\n                )\\n                break\\n            block_vis = create_block_vis(*batch_idx)\\n            total_vis.append(block_vis)\\n\\n        return \\\"\\\\n\\\".join(total_vis)\\n\\n    def to(self, device: Union[torch.device, str]) -> \\\"BlockMask\\\":\\n        \\\"\\\"\\\"Moves the BlockMask to the specified device.\\n\\n        Args:\\n            device (torch.device or str): The target device to move the BlockMask to.\\n                Can be a torch.device object or a string (e.g., 'cpu', 'cuda:0').\\n\\n        Returns:\\n            BlockMask: A new BlockMask instance with all tensor components moved\\n            to the specified device.\\n\\n        Note:\\n            This method does not modify the original BlockMask in-place.\\n            Instead, it returns a new BlockMask instance where invidual tensor attributes\\n            may or may not be moved to the specified device, depending on their\\n            current device placement.\\n        \\\"\\\"\\\"\\n        mapped_attributes = tree_map_only(\\n            torch.Tensor,\\n            lambda x: x.to(device),\\n            self.as_tuple(flatten=False),\\n        )\\n        return BlockMask(*mapped_attributes)\\n\\n\\ndef _broadcast_to_dim(x, dim):\\n    while x.dim() < dim:\\n        x = x.unsqueeze(0)\\n    return x\\n\\n\\ndef _round_up_to_multiple(x, multiple):\\n    return (x + multiple - 1) // multiple * multiple\\n\\n\\ndef _convert_mask_to_block_mask(\\n    mask: Tensor,\\n    KV_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,\\n    Q_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,\\n    separate_full_blocks: bool = False,\\n) -> Tuple[Tensor, Optional[Tensor]]:\\n    assert mask.dtype == torch.bool\\n    mask = _broadcast_to_dim(mask, 4)\\n    B, H, Q, KV = mask.shape\\n    assert Q % Q_BLOCK_SIZE == 0\\n    assert KV % KV_BLOCK_SIZE == 0\\n    mask = mask.view(\\n        B, H, Q // Q_BLOCK_SIZE, Q_BLOCK_SIZE, KV // KV_BLOCK_SIZE, KV_BLOCK_SIZE\\n    )  # [B, H, Q//Q_BLOCK_SIZE, Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE, KV_BLOCK_SIZE]\\n    mask = mask.permute(\\n        0, 1, 2, 4, 3, 5\\n    )  # [B, H, Q//Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE, Q_BLOCK_SIZE, KV_BLOCK_SIZE]\\n    mask_block_sum = mask.sum(\\n        dim=[-2, -1]\\n    )  # [B, H, Q//Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE]\\n    if separate_full_blocks:\\n        full_block_sum = Q_BLOCK_SIZE * KV_BLOCK_SIZE\\n        full_blocks = mask_block_sum == full_block_sum\\n        partial_blocks = (mask_block_sum > 0) & (mask_block_sum < full_block_sum)\\n        partial_blocks = partial_blocks.to(dtype=torch.int8)\\n        full_blocks = full_blocks.to(dtype=torch.int8)\\n        return partial_blocks, full_blocks\\n    else:\\n        partial_blocks = mask_block_sum > 0\\n        partial_blocks = partial_blocks.to(dtype=torch.int8)\\n        return partial_blocks, None\\n\\n\\ndef or_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:\\n    \\\"\\\"\\\"Returns a mask_mod that's the union of provided mask_mods\\\"\\\"\\\"\\n    if not all(callable(arg) for arg in mask_mods):\\n        raise RuntimeError(f\\\"All inputs should be callable mask_mods: {mask_mods}\\\")\\n\\n    def or_mask(b, h, q_idx, kv_idx):\\n        result = b.new_zeros((), dtype=torch.bool)\\n        for mask in mask_mods:\\n            result = result | mask(b, h, q_idx, kv_idx)\\n        return result\\n\\n    return or_mask\\n\\n\\ndef and_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:\\n    \\\"\\\"\\\"Returns a mask_mod that's the intersection of provided mask_mods\\\"\\\"\\\"\\n    if not all(callable(arg) for arg in mask_mods):\\n        raise RuntimeError(f\\\"All inputs should be callable mask_mods: {mask_mods}\\\")\\n\\n    def and_mask(b, h, q_idx, kv_idx):\\n        result = b.new_ones((), dtype=torch.bool)\\n        for mask in mask_mods:\\n            result = result & mask(b, h, q_idx, kv_idx)\\n        return result\\n\\n    return and_mask\\n\\n\\ndef _convert_block_mask_to_mask(\\n    block_mask,\\n    KV_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,\\n    Q_BLOCK_SIZE=_DEFAULT_SPARSE_BLOCK_SIZE,\\n) -> Tensor:\\n    assert block_mask.dim() == 4\\n    B, H, Q, KV = block_mask.shape\\n    block_mask = block_mask.expand(Q_BLOCK_SIZE, KV_BLOCK_SIZE, *block_mask.shape)\\n    block_mask = block_mask.permute(2, 3, 4, 0, 5, 1).reshape(\\n        B, H, Q * Q_BLOCK_SIZE, KV * KV_BLOCK_SIZE\\n    )\\n    return block_mask\\n\\n\\ndef _create_sparse_block_from_block_mask(\\n    block_mask: Tuple[Tensor, Optional[Tensor]],\\n    mask_mod: Optional[Callable],\\n    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,\\n    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,\\n) -> BlockMask:\\n    partial_blocks, full_blocks = block_mask\\n\\n    partial_bm = _dense_to_ordered(partial_blocks)\\n    if full_blocks is not None:\\n        full_bm = _dense_to_ordered(full_blocks)\\n    else:\\n        full_bm = (None, None)\\n\\n    return BlockMask.from_kv_blocks(\\n        partial_bm[0],\\n        partial_bm[1],\\n        full_bm[0],\\n        full_bm[1],\\n        BLOCK_SIZE=(KV_BLOCK_SIZE, Q_BLOCK_SIZE),\\n        mask_mod=mask_mod,\\n    )\\n\\n\\ndef create_mask(\\n    mod_fn: Union[_score_mod_signature, _mask_mod_signature],\\n    B: Optional[int],\\n    H: Optional[int],\\n    Q_LEN: int,\\n    KV_LEN: int,\\n    device: str = \\\"cuda\\\",\\n    _compile: bool = False,\\n) -> Tensor:\\n    r\\\"\\\"\\\"This function creates a mask tensor from a mod_fn function.\\n\\n    Args:\\n        mod_fn (Union[_score_mod_signature, _mask_mod_signature]): Function to modify attention scores.\\n        B (int): Batch size.\\n        H (int): Number of query heads.\\n        Q_LEN (int): Sequence length of query.\\n        KV_LEN (int): Sequence length of key/value.\\n        device (str): Device to run the mask creation on.\\n\\n    Returns:\\n        mask (Tensor): A mask tensor with shape (B, H, M, N).\\n    \\\"\\\"\\\"\\n    if B is None:\\n        B = 1\\n    if H is None:\\n        H = 1\\n    b = torch.arange(0, B, device=device)\\n    h = torch.arange(0, H, device=device)\\n    m = torch.arange(0, Q_LEN, device=device)\\n    n = torch.arange(0, KV_LEN, device=device)\\n    # TODO: fix this\\n    # Lack instantiation support for __torch_function__ mode support under compile\\n    if _compile:\\n        ctx = nullcontext()\\n    else:\\n        ctx = TransformGetItemToIndex()  # type: ignore[assignment]\\n    mod_type = _get_mod_type(mod_fn)\\n\\n    with ctx:\\n        if mod_type == _ModificationType.SCORE_MOD:\\n            score_mod = mod_fn\\n            score_mod = _vmap_for_bhqkv(score_mod, prefix=(0,))  # first input is score\\n            out = score_mod(torch.zeros(B, H, Q_LEN, KV_LEN, device=device), b, h, m, n)\\n            mask = torch.where(torch.isneginf(out), False, True)\\n            return mask\\n        elif mod_type == _ModificationType.MASK_MOD:\\n            mask_mod = mod_fn\\n            mask_mod = _vmap_for_bhqkv(mask_mod, prefix=())\\n            mask = mask_mod(b, h, m, n)\\n            return mask\\n        else:\\n            raise AssertionError\\n\\n\\ndef _create_block_mask_inner(\\n    mask_mod: Callable,\\n    B: int,\\n    H: int,\\n    Q_LEN: int,\\n    KV_LEN: int,\\n    device: str,\\n    KV_BLOCK_SIZE: int,\\n    Q_BLOCK_SIZE: int,\\n):\\n    r\\\"\\\"\\\"Work around for being unable to instantiate __torch_function__ mode under compile.\\n    `create_block_mask` will compile this inner function and wrap the call to this\\n    with the __torch_function__ mode.\\n    \\\"\\\"\\\"\\n    mask_tensor = create_mask(mask_mod, B, H, Q_LEN, KV_LEN, device, _compile=True)\\n    partial_block_mask, full_block_mask = _convert_mask_to_block_mask(\\n        mask_tensor,\\n        KV_BLOCK_SIZE=KV_BLOCK_SIZE,\\n        Q_BLOCK_SIZE=Q_BLOCK_SIZE,\\n        separate_full_blocks=True,\\n    )\\n    return partial_block_mask, full_block_mask\\n\\n\\ndef create_block_mask(\\n    mask_mod: _mask_mod_signature,\\n    B: Optional[int],\\n    H: Optional[int],\\n    Q_LEN: int,\\n    KV_LEN: int,\\n    device: str = \\\"cuda\\\",\\n    BLOCK_SIZE: Union[int, Tuple[int, int]] = _DEFAULT_SPARSE_BLOCK_SIZE,\\n    _compile=False,\\n) -> BlockMask:\\n    r\\\"\\\"\\\"This function creates a block mask tuple from a mask_mod function.\\n\\n    Args:\\n        mask_mod (Callable): mask_mod function. This is a callable that defines the\\n            masking pattern for the attention mechanism. It takes four arguments:\\n            b (batch size), h (number of heads), q_idx (query index), and kv_idx (key/value index).\\n            It should return a boolean tensor indicating which attention connections are allowed (True)\\n            or masked out (False).\\n        B (int): Batch size.\\n        H (int): Number of query heads.\\n        Q_LEN (int): Sequence length of query.\\n        KV_LEN (int): Sequence length of key/value.\\n        device (str): Device to run the mask creation on.\\n        KV_BLOCK_SIZE (int): Block size of block mask for each query.\\n        Q_BLOCK_SIZE (int): Block size of block mask for each key/value.\\n        _compile (bool): Whether to compile the mask creation.\\n\\n    Returns:\\n        BlockMask:  A BlockMask object that contains the block mask information.\\n\\n    Example Usage:\\n        .. code-block:: python\\n\\n            def causal_mask(b, h, q_idx, kv_idx):\\n                return q_idx >= kv_idx\\n\\n            block_mask = create_block_mask(causal_mask, 1, 1, 8192, 8192, device=\\\"cuda\\\")\\n            query = torch.randn(1, 1, 8192, 64, device=\\\"cuda\\\", dtype=torch.float16)\\n            key = torch.randn(1, 1, 8192, 64, device=\\\"cuda\\\", dtype=torch.float16)\\n            value = torch.randn(1, 1, 8192, 64, device=\\\"cuda\\\", dtype=torch.float16)\\n            output = flex_attention(query, key, value, block_mask=block_mask)\\n    \\\"\\\"\\\"\\n    mod_type = _get_mod_type(mask_mod)\\n    assert (\\n        mod_type == _ModificationType.MASK_MOD\\n    ), f\\\"create-block_mask requires a mask_mod function! Got {mask_mod}\\\"\\n    inner_func = _create_block_mask_inner\\n    if B is None:\\n        B = 1\\n    if H is None:\\n        H = 1\\n    if isinstance(BLOCK_SIZE, int):\\n        Q_BLOCK_SIZE = BLOCK_SIZE\\n        KV_BLOCK_SIZE = BLOCK_SIZE\\n    else:\\n        Q_BLOCK_SIZE, KV_BLOCK_SIZE = BLOCK_SIZE\\n\\n    if Q_LEN < 128:\\n        Q_BLOCK_SIZE = Q_LEN\\n    else:\\n        Q_LEN = _round_up_to_multiple(Q_LEN, Q_BLOCK_SIZE)\\n    KV_LEN = _round_up_to_multiple(KV_LEN, KV_BLOCK_SIZE)\\n    if _compile:\\n        inner_func = torch.compile(inner_func, fullgraph=True, dynamic=False)\\n    with TransformGetItemToIndex():\\n        partial_block_mask, full_block_mask = inner_func(\\n            mask_mod, B, H, Q_LEN, KV_LEN, device, KV_BLOCK_SIZE, Q_BLOCK_SIZE\\n        )\\n        block_mask = _create_sparse_block_from_block_mask(\\n            (partial_block_mask, full_block_mask), mask_mod\\n        )\\n    return block_mask\\n\\n\\ndef _create_empty_block_mask(query: Tensor, key: Tensor) -> BlockMask:\\n    r\\\"\\\"\\\"Default block mask for flex attention.\\n    If users don't specify any block sparse mask info, we create this\\n    empty block sparse mask. Which creates a BlockMask with 1 block that is the full length\\n    of the query and key tensors.\\n    \\\"\\\"\\\"\\n    device = query.device\\n    return BlockMask.from_kv_blocks(\\n        kv_num_blocks=torch.ones([1, 1, 1], dtype=torch.int32, device=device),\\n        kv_indices=torch.zeros([1, 1, 1, 1], dtype=torch.int32, device=device),\\n        BLOCK_SIZE=_LARGE_SPARSE_BLOCK_SIZE,\\n    )\\n\\n\\ndef _apply_kernel_options(\\n    query: Tensor, key: Tensor, value: Tensor, return_lse: bool, kernel_options\\n):\\n    kernel_options = {} if kernel_options is None else dict(kernel_options)\\n\\n    kernel_options.setdefault(\\\"ROWS_GUARANTEED_SAFE\\\", False)\\n    kernel_options.setdefault(\\\"PRESCALE_QK\\\", False)\\n\\n    # If foward kernel needs to return logsumexp is decided by this rule internally.\\n    assert \\\"OUTPUT_LOGSUMEXP\\\" not in kernel_options\\n    kernel_options[\\\"OUTPUT_LOGSUMEXP\\\"] = True\\n    if not return_lse:\\n        any_inputs_require_grad = (\\n            query.requires_grad or key.requires_grad or value.requires_grad\\n        )\\n        output_logsumexp = any_inputs_require_grad and torch.is_grad_enabled()\\n        kernel_options[\\\"OUTPUT_LOGSUMEXP\\\"] = output_logsumexp\\n\\n    return kernel_options\\n\\n\\ndef _validate_embed_dim(query: Tensor, key: Tensor, value: Tensor):\\n    if query.size(-1) != key.size(-1):\\n        raise ValueError(\\n            f\\\"Expect query and key/value to have the same embedding dimension \\\"\\n            f\\\"but got E={query.size(-1)} and E={key.size(-1)}.\\\"\\n        )\\n    # TODO this config segfaults with Triton without:\\n    # https://github.com/triton-lang/triton/pull/4540\\n    if not (\\n        _supported_head_dim(query.size(-1)) and _supported_head_dim(value.size(-1))\\n    ):\\n        raise ValueError(\\n            f\\\"NYI: Currently non power of 2 embedding dimension are not supported. \\\"\\n            f\\\"Got E={query.size(-1)} and Ev={value.size(-1)}.\\\"\\n        )\\n    if value.size(-1) > query.size(-1):\\n        raise ValueError(\\n            f\\\"NYI: Currently value embedding dimension must be less than or equal to query embedding dimension. \\\"\\n            f\\\"Got Ev={value.size(-1)} and E={query.size(-1)}.\\\"\\n        )\\n\\n\\ndef flex_attention(\\n    query: Tensor,\\n    key: Tensor,\\n    value: Tensor,\\n    score_mod: Optional[_score_mod_signature] = None,\\n    block_mask: Optional[BlockMask] = None,\\n    scale: Optional[float] = None,\\n    enable_gqa: bool = False,\\n    return_lse: bool = False,\\n    kernel_options: Optional[Dict[str, Any]] = None,\\n) -> Union[Tensor, Tuple[Tensor, Tensor]]:\\n    r\\\"\\\"\\\"This function implements scaled dot product attention with an arbitrary attention score modification function.\\n\\n    This function computes the scaled dot product attention between query, key, and value tensors with a user-defined\\n    attention score modification function. The attention score modification function will be applied after the attention\\n    scores have been calculated between the query and key tensors. The attention scores are calculated as follows:\\n\\n    The ``score_mod`` function should have the following signature:\\n\\n    .. code-block:: python\\n\\n        def score_mod(\\n            score: Tensor,\\n            batch: Tensor,\\n            head: Tensor,\\n            q_idx: Tensor,\\n            k_idx: Tensor\\n        ) -> Tensor:\\n\\n    Where:\\n        - ``score``: A scalar tensor representing the attention score,\\n          with the same data type and device as the query, key, and value tensors.\\n        - ``batch``, ``head``, ``q_idx``, ``k_idx``: Scalar tensors indicating\\n          the batch index, query head index, query index, and key/value index, respectively.\\n          These should have the ``torch.int`` data type and be located on the same device as the score tensor.\\n\\n    Args:\\n        query (Tensor): Query tensor; shape :math:`(B, Hq, L, E)`.\\n        key (Tensor): Key tensor; shape :math:`(B, Hkv, S, E)`.\\n        value (Tensor): Value tensor; shape :math:`(B, Hkv, S, Ev)`.\\n        score_mod (Optional[Callable]): Function to modify attention scores. By default no score_mod is applied.\\n        block_mask (Optional[BlockMask]): BlockMask object that controls the blocksparsity pattern of the attention.\\n        scale (Optional[float]): Scaling factor applied prior to softmax. If none, the default value is set to :math:`\\\\frac{1}{\\\\sqrt{E}}`.\\n        enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA) and broadcasts key/value heads to query heads.\\n        return_lse (bool): Whether to return the logsumexp of the attention scores. Default is False.\\n        kernel_options (Optional[Dict[str, Any]]): Options to pass into the Triton kernels.\\n\\n    Returns:\\n        output (Tensor): Attention output; shape :math:`(B, Hq, L, Ev)`.\\n\\n    Shape legend:\\n        - :math:`N: \\\\text{Batch size} ... : \\\\text{Any number of other batch dimensions (optional)}`\\n        - :math:`S: \\\\text{Source sequence length}`\\n        - :math:`L: \\\\text{Target sequence length}`\\n        - :math:`E: \\\\text{Embedding dimension of the query and key}`\\n        - :math:`Ev: \\\\text{Embedding dimension of the value}`\\n\\n    .. warning::\\n        `torch.nn.attention.flex_attention` is a prototype feature in PyTorch.\\n        Please look forward to a more stable implementation in a future version of PyTorch.\\n        Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype\\n\\n    \\\"\\\"\\\"\\n    # Some basic input validation\\n    _validate_sdpa_input(query, key, value)\\n    _validate_embed_dim(query, key, value)\\n    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:\\n        raise NotImplementedError(\\\"NYI: query, key, and value must be 4D tensors\\\")\\n    if (not enable_gqa) and query.size(-3) != key.size(-3):\\n        raise ValueError(\\n            f\\\"Expect query and key/value to have the same number of heads \\\"\\n            f\\\"but got Hq={query.size(-3)} and Hkv={key.size(-3)}. \\\"\\n            f\\\"Try setting enable_gqa=True for GQA.\\\"\\n        )\\n    if enable_gqa:\\n        Hq = query.size(1)\\n        Hkv = key.size(1)\\n        if Hq % Hkv != 0:\\n            raise ValueError(\\n                f\\\"Expect number of query heads to be a multiple of kv heads for GQA \\\"\\n                f\\\"but got Hq={Hq} and Hkv={Hkv}.\\\"\\n            )\\n\\n    if score_mod is None:\\n        score_mod = _identity\\n    if block_mask is None:\\n        block_mask = _create_empty_block_mask(query, key)\\n    if scale is None:\\n        scale = 1.0 / math.sqrt(query.size(-1))\\n\\n    kernel_options = _apply_kernel_options(\\n        query,\\n        key,\\n        value,\\n        return_lse,\\n        kernel_options,\\n    )\\n\\n    if torch.compiler.is_dynamo_compiling():\\n        # mark head_dim and number of heads to be static\\n        for x in [query, key, value]:\\n            torch._dynamo.mark_static(x, -3)\\n            torch._dynamo.mark_static(x, -1)\\n        out, lse = flex_attention_hop(\\n            query, key, value, score_mod, block_mask.as_tuple(), scale, kernel_options\\n        )\\n        if return_lse:\\n            return out, lse * math.log(2)\\n        else:\\n            return out\\n\\n    if not torch._dynamo.is_dynamo_supported():\\n        raise RuntimeError(\\\"flex_attention requires dynamo support\\\")\\n\\n    # Dynamo is expecting a callable with \\\"__code__\\\" attribute.\\n    # We cannot directly pass hop to it. So we wrap it in a dummy function.\\n    def _flex_attention_hop_wrapper(*args, **kwargs):\\n        return flex_attention_hop(*args, **kwargs)\\n\\n    with _set_compilation_env():\\n        with torch._dynamo.utils.disable_cache_limit():\\n            with _temp_remove_pre_dispatch_torch_function_mode():\\n                out, lse = torch.compile(\\n                    _flex_attention_hop_wrapper, backend=\\\"eager\\\", fullgraph=True\\n                )(\\n                    query,\\n                    key,\\n                    value,\\n                    score_mod,\\n                    block_mask.as_tuple(),\\n                    scale,\\n                    kernel_options,\\n                )\\n                if return_lse:\\n                    return out, lse * math.log(2)\\n                else:\\n                    return out\\n\\n\\n# mypy: allow-untyped-defs\\n\\\"\\\"\\\" This module contains functions and classes that alter the behavior of torch.nn.functional.scaled_dot_product_attention \\\"\\\"\\\"\\nimport contextlib\\nfrom typing import List, Union\\nfrom warnings import warn\\n\\nfrom torch._C import _SDPBackend as SDPBackend\\nfrom torch.backends.cuda import (\\n    can_use_efficient_attention,\\n    can_use_flash_attention,\\n    cudnn_sdp_enabled,\\n    enable_cudnn_sdp,\\n    enable_flash_sdp,\\n    enable_math_sdp,\\n    enable_mem_efficient_sdp,\\n    flash_sdp_enabled,\\n    math_sdp_enabled,\\n    mem_efficient_sdp_enabled,\\n    SDPAParams,\\n)\\n\\n\\n__all__: List[str] = [\\\"SDPBackend\\\", \\\"sdpa_kernel\\\", \\\"WARN_FOR_UNFUSED_KERNELS\\\"]\\n\\n# Note: [SDPA warnings]\\n# TODO: Consider using this for sdpa regardless of subclasses\\n# This only effects users of bias subclasses\\n# If this is set to True, we will warn the user if they are not using the fused kernels\\n# As well, it will raise warnings for all the reasons why the fused kernels can't be run.\\n# To set this to True, run\\n# torch.nn.attention.WARN_FOR_UNFUSED_KERNELS = True\\nWARN_FOR_UNFUSED_KERNELS = False\\n\\n\\n# Hacks for Sphinx documentation:\\n# https://stackoverflow.com/questions/38765577/overriding-sphinx-autodoc-alias-of-for-import-of-private-class\\nSDPBackend = SDPBackend\\nr\\\"\\\"\\\"An enum-like class that contains the different backends for scaled dot product attention.\\n    This backend class is designed to be used with the sdpa_kernel context manager.\\n\\n    The following Enums are available:\\n        - ERROR: An error occurred when trying to determine the backend.\\n        - MATH: The math backend for scaled dot product attention.\\n        - FLASH_ATTENTION: The flash attention backend for scaled dot product attention.\\n        - EFFICIENT_ATTENTION: The efficient attention backend for scaled dot product attention.\\n        - CUDNN_ATTENTION: The cuDNN backend for scaled dot product attention.\\n\\n    See :func:`torch.nn.attention.sdpa_kernel` for more details.\\n\\n    .. warning:: This class is in beta and subject to change.\\n\\\"\\\"\\\"\\nSDPBackend.__module__ = __name__\\nSDPBackend.__name__ = \\\"SDPBackend\\\"\\n\\n\\ndef _raise_kernel_warnings(params: SDPAParams) -> None:\\n    \\\"\\\"\\\"\\n    If WARN_FOR_UNFUSED_KERNELS is set to True, this will raise warnings\\n    for all the reasons why the fused kernels can't be run. If using subclasses\\n    \\\"\\\"\\\"\\n    if WARN_FOR_UNFUSED_KERNELS:\\n        if not can_use_efficient_attention(params):\\n            warn(\\\"Efficient attention can't be used because:\\\")\\n            can_use_efficient_attention(params, True)\\n        if not can_use_flash_attention(params):\\n            warn(\\\"Flash attention can't be used because:\\\")\\n            can_use_flash_attention(params, True)\\n\\n\\n@contextlib.contextmanager\\ndef sdpa_kernel(backends: Union[List[SDPBackend], SDPBackend]):\\n    r\\\"\\\"\\\"\\n    Context manager to select which backend to use for scaled dot product attention.\\n\\n    .. warning:: This function is beta and subject to change.\\n\\n    Args:\\n        backend (Union[List[SDPBackend], SDPBackend]): A backend or list of backends for scaled dot product attention.\\n\\n    Example:\\n\\n    .. code-block:: python\\n\\n        from torch.nn.functional import scaled_dot_product_attention\\n        from torch.nn.attention import SDPBackend, sdpa_kernel\\n        # Only enable flash attention backend\\n        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):\\n            scaled_dot_product_attention(...)\\n\\n        # Enable the Math or Efficient attention backends\\n        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):\\n            scaled_dot_product_attention(...)\\n\\n    This context manager can be used to select which backend to use for scaled dot product attention.\\n    Upon exiting the context manager, the previous state of the flags will be restored, enabling all backends.\\n    \\\"\\\"\\\"\\n    assert isinstance(\\n        backends, (list, SDPBackend)\\n    ), \\\"Backend must be an instance of SDPBackend or a list of SDPBackend instances\\\"\\n\\n    if isinstance(backends, SDPBackend):\\n        backends = [backends]\\n\\n    backends = set(backends)\\n    previous_cudnn: bool = cudnn_sdp_enabled()\\n    previous_flash: bool = flash_sdp_enabled()\\n    previous_mem_efficient: bool = mem_efficient_sdp_enabled()\\n    previous_math: bool = math_sdp_enabled()\\n    try:\\n        enable_cudnn = SDPBackend.CUDNN_ATTENTION in backends\\n        enable_flash = SDPBackend.FLASH_ATTENTION in backends\\n        enable_mem_efficient = SDPBackend.EFFICIENT_ATTENTION in backends\\n        enable_math = SDPBackend.MATH in backends\\n\\n        enable_cudnn_sdp(enable_cudnn)\\n        enable_flash_sdp(enable_flash)\\n        enable_mem_efficient_sdp(enable_mem_efficient)\\n        enable_math_sdp(enable_math)\\n        yield {}\\n    finally:\\n        enable_cudnn_sdp(previous_cudnn)\\n        enable_flash_sdp(previous_flash)\\n        enable_mem_efficient_sdp(previous_mem_efficient)\\n        enable_math_sdp(previous_math)\\n\\n\\ndef _get_flash_version() -> str:\\n    \\\"\\\"\\\"This returns the closest matching tag for the flash attention backend\\\"\\\"\\\"\\n    return \\\"2.5.7\\\"\",\"difficulty\":\"easy\",\"domain\":\"Code Repository Understanding\",\"length\":\"long\",\"question\":\"This is the troch.nn modeule. In this module, there exists an implementation of flexible attention mechanisms, in this implementation, what is the default value used for the BLOCK_SIZE parameter when creating a BlockMask from key-value block information if no specific block size is provided, and what is the corresponding _ModificationType enum value that represents a score modification function?\",\"sub_domain\":\"Code repo QA\"}","display_format":"text","language":"","answer_status":"published","assets":[],"source_url":"https://huggingface.co/datasets/zai-org/LongBench-v2","history":"initial import","indexing_mode":"noindex","subproblems":[],"grids":[]}