[docs]classIdentity(Module):r"""A placeholder identity operator that is argument-insensitive. Args: args: any argument (unused) kwargs: any keyword argument (unused) Shape: - Input: :math:`(*)`, where :math:`*` means any number of dimensions. - Output: :math:`(*)`, same shape as the input. Examples:: >>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 20]) """def__init__(self,*args,**kwargs):super(Identity,self).__init__()defforward(self,input:Tensor)->Tensor:returninput
[docs]classLinear(Module):r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b` This module supports :ref:`TensorFloat32<tf32_on_ampere>`. Args: in_features: size of each input sample out_features: size of each output sample bias: If set to ``False``, the layer will not learn an additive bias. Default: ``True`` Shape: - Input: :math:`(*, H_{in})` where :math:`*` means any number of dimensions including none and :math:`H_{in} = \text{in\_features}`. - Output: :math:`(*, H_{out})` where all but the last dimension are the same shape as the input and :math:`H_{out} = \text{out\_features}`. Attributes: weight: the learnable weights of the module of shape :math:`(\text{out\_features}, \text{in\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in\_features}}` bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`k = \frac{1}{\text{in\_features}}` Examples:: >>> m = nn.Linear(20, 30) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 30]) """__constants__=['in_features','out_features']in_features:intout_features:intweight:Tensordef__init__(self,in_features:int,out_features:int,bias:bool=True,device=None,dtype=None)->None:factory_kwargs={'device':device,'dtype':dtype}super(Linear,self).__init__()self.in_features=in_featuresself.out_features=out_featuresself.weight=Parameter(torch.empty((out_features,in_features),**factory_kwargs))ifbias:self.bias=Parameter(torch.empty(out_features,**factory_kwargs))else:self.register_parameter('bias',None)self.reset_parameters()defreset_parameters(self)->None:# Setting a=sqrt(5) in kaiming_uniform is the same as initializing with# uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see# https://github.com/pytorch/pytorch/issues/57109init.kaiming_uniform_(self.weight,a=math.sqrt(5))ifself.biasisnotNone:fan_in,_=init._calculate_fan_in_and_fan_out(self.weight)bound=1/math.sqrt(fan_in)iffan_in>0else0init.uniform_(self.bias,-bound,bound)defforward(self,input:Tensor)->Tensor:returnF.linear(input,self.weight,self.bias)defextra_repr(self)->str:return'in_features={}, out_features={}, bias={}'.format(self.in_features,self.out_features,self.biasisnotNone)
# This class exists solely to avoid triggering an obscure error when scripting# an improperly quantized attention layer. See this issue for details:# https://github.com/pytorch/pytorch/issues/58969# TODO: fail fast on quantization API usage error, then remove this class# and replace uses of it with plain LinearclassNonDynamicallyQuantizableLinear(Linear):def__init__(self,in_features:int,out_features:int,bias:bool=True,device=None,dtype=None)->None:super().__init__(in_features,out_features,bias=bias,device=device,dtype=dtype)classBilinear(Module):r"""Applies a bilinear transformation to the incoming data: :math:`y = x_1^T A x_2 + b` Args: in1_features: size of each first input sample in2_features: size of each second input sample out_features: size of each output sample bias: If set to False, the layer will not learn an additive bias. Default: ``True`` Shape: - Input1: :math:`(N, *, H_{in1})` where :math:`H_{in1}=\text{in1\_features}` and :math:`*` means any number of additional dimensions. All but the last dimension of the inputs should be the same. - Input2: :math:`(N, *, H_{in2})` where :math:`H_{in2}=\text{in2\_features}`. - Output: :math:`(N, *, H_{out})` where :math:`H_{out}=\text{out\_features}` and all but the last dimension are the same shape as the input. Attributes: weight: the learnable weights of the module of shape :math:`(\text{out\_features}, \text{in1\_features}, \text{in2\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in1\_features}}` bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in1\_features}}` Examples:: >>> m = nn.Bilinear(20, 30, 40) >>> input1 = torch.randn(128, 20) >>> input2 = torch.randn(128, 30) >>> output = m(input1, input2) >>> print(output.size()) torch.Size([128, 40]) """__constants__=['in1_features','in2_features','out_features']in1_features:intin2_features:intout_features:intweight:Tensordef__init__(self,in1_features:int,in2_features:int,out_features:int,bias:bool=True,device=None,dtype=None)->None:factory_kwargs={'device':device,'dtype':dtype}super(Bilinear,self).__init__()self.in1_features=in1_featuresself.in2_features=in2_featuresself.out_features=out_featuresself.weight=Parameter(torch.empty((out_features,in1_features,in2_features),**factory_kwargs))ifbias:self.bias=Parameter(torch.empty(out_features,**factory_kwargs))else:self.register_parameter('bias',None)self.reset_parameters()defreset_parameters(self)->None:bound=1/math.sqrt(self.weight.size(1))init.uniform_(self.weight,-bound,bound)ifself.biasisnotNone:init.uniform_(self.bias,-bound,bound)defforward(self,input1:Tensor,input2:Tensor)->Tensor:returnF.bilinear(input1,input2,self.weight,self.bias)defextra_repr(self)->str:return'in1_features={}, in2_features={}, out_features={}, bias={}'.format(self.in1_features,self.in2_features,self.out_features,self.biasisnotNone)
[docs]classLazyLinear(LazyModuleMixin,Linear):r"""A :class:`torch.nn.Linear` module where `in_features` is inferred. In this module, the `weight` and `bias` are of :class:`torch.nn.UninitializedParameter` class. They will be initialized after the first call to ``forward`` is done and the module will become a regular :class:`torch.nn.Linear` module. The ``in_features`` argument of the :class:`Linear` is inferred from the ``input.shape[-1]``. Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation on lazy modules and their limitations. Args: out_features: size of each output sample bias: If set to ``False``, the layer will not learn an additive bias. Default: ``True`` Attributes: weight: the learnable weights of the module of shape :math:`(\text{out\_features}, \text{in\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in\_features}}` bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`k = \frac{1}{\text{in\_features}}` """cls_to_become=Linear# type: ignore[assignment]weight:UninitializedParameterbias:UninitializedParameter# type: ignore[assignment]def__init__(self,out_features:int,bias:bool=True,device=None,dtype=None)->None:factory_kwargs={'device':device,'dtype':dtype}# bias is hardcoded to False to avoid creating tensor# that will soon be overwritten.super().__init__(0,0,False)self.weight=UninitializedParameter(**factory_kwargs)self.out_features=out_featuresifbias:self.bias=UninitializedParameter(**factory_kwargs)defreset_parameters(self)->None:ifnotself.has_uninitialized_params()andself.in_features!=0:super().reset_parameters()definitialize_parameters(self,input)->None:# type: ignore[override]ifself.has_uninitialized_params():withtorch.no_grad():self.in_features=input.shape[-1]self.weight.materialize((self.out_features,self.in_features))ifself.biasisnotNone:self.bias.materialize((self.out_features,))self.reset_parameters()
# TODO: PartialLinear - maybe in sparse?
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.