def __deepcopy__(self, memo):
return handle_torch_function(torch.Tensor.__deepcopy__, (self, memo),
self, memo)
def quux(a):
"""Used to test that errors raised in user implementations get propagated"""
if type(a) is not Tensor and has_torch_function((a,)):
return handle_torch_function(quux, (a,), a)
return a
def __ipow__(self, other):
relevant_args = (self, other)
from torch.overrides import has_torch_function, handle_torch_function
if type(self) is not Tensor and type(other) is not Tensor and has_torch_function(relevant_args):
return handle_torch_function(Tensor.__ipow__, relevant_args, self, other)
return NotImplemented
def __hash__(self):
if has_torch_function_unary(self):
return handle_torch_function(Tensor.__hash__, (self,), self)
return id(self)
def bar(a):
"""A function with one argument"""
if type(a) is not Tensor and has_torch_function((a,)):
return handle_torch_function(bar, (a,), a)
return a
def __rdiv__(self, other):
if has_torch_function_variadic(self, other):
return handle_torch_function(Tensor.__rdiv__, (self, other), self, other)
return self.reciprocal() * other
def __ipow__(self, other): # type: ignore[misc]
if has_torch_function_variadic(self, other):
return handle_torch_function(Tensor.__ipow__, (self, other), self, other)
return NotImplemented
def __getitem__(self, key):
return handle_torch_function(torch.Tensor.__getitem__, (self, key), self, key)
def __sub__(self, other):
return handle_torch_function(torch.Tensor.__sub__, (self, other), self, other)
def __rtruediv__(self, other):
return handle_torch_function(torch.Tensor.__rdiv__, (self, other), self, other)
def tanh(self):
return handle_torch_function(torch.Tensor.tanh, (self,), self)
def grad(self):
relevant_args = (self,)
from torch.overrides import has_torch_function, handle_torch_function
if type(self) is not Tensor and has_torch_function(relevant_args):
return handle_torch_function(Tensor.grad.__delete__, relevant_args, self)
del self._grad
def detach(self):
return handle_torch_function(torch.Tensor.detach, (self, ), self)
def clone(self, *, memory_format=torch.preserve_format):
return handle_torch_function(torch.Tensor.clone, (self, ),
self,
memory_format=memory_format)
def resize_as(self, tensor):
if has_torch_function_variadic(self, tensor):
return handle_torch_function(Tensor.resize_as, (self, tensor), self, tensor)
warnings.warn("non-inplace resize_as is deprecated")
from torch.autograd._functions import Resize
return Resize.apply(self, tensor.size())
def multi_head_attention_forward(
query: Tensor,
key: Tensor,
value: Tensor,
embed_dim_to_check: int,
num_heads: int,
in_proj_weight: Tensor,
in_proj_bias: Tensor,
bias_k: Optional[Tensor],
bias_v: Optional[Tensor],
add_zero_attn: bool,
dropout_p: float,
out_proj_weight: Tensor,
out_proj_bias: Tensor,
training: bool = True,
key_padding_mask: Optional[Tensor] = None,
need_weights: bool = True,
attn_mask: Optional[Tensor] = None,
use_separate_proj_weight: bool = False,
q_proj_weight: Optional[Tensor] = None,
k_proj_weight: Optional[Tensor] = None,
v_proj_weight: Optional[Tensor] = None,
static_k: Optional[Tensor] = None,
static_v: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
r"""
Args:
query, key, value: map a query and a set of key-value pairs to an output.
See "Attention Is All You Need" for more details.
embed_dim_to_check: total dimension of the model.
num_heads: parallel attention heads.
in_proj_weight, in_proj_bias: input projection weight and bias.
bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
add_zero_attn: add a new batch of zeros to the key and
value sequences at dim=1.
dropout_p: probability of an element to be zeroed.
out_proj_weight, out_proj_bias: the output projection weight and bias.
training: apply dropout if is ``True``.
key_padding_mask: if provided, specified padding elements in the key will
be ignored by the attention. This is an binary mask. When the value is True,
the corresponding value on the attention layer will be filled with -inf.
need_weights: output attn_output_weights.
attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
the batches while a 3D mask allows to specify a different mask for the entries of each batch.
use_separate_proj_weight: the function accept the proj. weights for query, key,
and value in different forms. If false, in_proj_weight will be used, which is
a combination of q_proj_weight, k_proj_weight, v_proj_weight.
q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
static_k, static_v: static key and value used for attention operators.
Shape:
Inputs:
- query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
the embedding dimension.
- key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
the embedding dimension.
- value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
the embedding dimension.
- key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
will be unchanged. If a BoolTensor is provided, the positions with the
value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
- attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
is provided, it will be added to the attention weight.
- static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
- static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
Outputs:
- attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
E is the embedding dimension.
- attn_output_weights: :math:`(N, L, S)` where N is the batch size,
L is the target sequence length, S is the source sequence length.
"""
if not torch.jit.is_scripting():
tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k,
bias_v, out_proj_weight, out_proj_bias)
if any([type(t) is not Tensor
for t in tens_ops]) and has_torch_function(tens_ops):
return handle_torch_function(
multi_head_attention_forward,
tens_ops,
query,
key,
value,
embed_dim_to_check,
num_heads,
in_proj_weight,
in_proj_bias,
bias_k,
bias_v,
add_zero_attn,
dropout_p,
out_proj_weight,
out_proj_bias,
training=training,
key_padding_mask=key_padding_mask,
need_weights=need_weights,
attn_mask=attn_mask,
use_separate_proj_weight=use_separate_proj_weight,
q_proj_weight=q_proj_weight,
k_proj_weight=k_proj_weight,
v_proj_weight=v_proj_weight,
static_k=static_k,
static_v=static_v)
tgt_len, bsz, embed_dim = query.size()
assert embed_dim == embed_dim_to_check
# allow MHA to have different sizes for the feature dimension
assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
head_dim = embed_dim // num_heads
assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
scaling = float(head_dim)**-0.5
if not use_separate_proj_weight:
if (query is key or torch.equal(
query, key)) and (key is value or torch.equal(key, value)):
# self-attention
q, k, v = linear(query, in_proj_weight, in_proj_bias).chunk(3,
dim=-1)
elif (key is value or torch.equal(key, value)):
# encoder-decoder attention
# This is inline in_proj function with in_proj_weight and in_proj_bias
_b = in_proj_bias
_start = 0
_end = embed_dim
_w = in_proj_weight[_start:_end, :]
if _b is not None:
_b = _b[_start:_end]
q = linear(query, _w, _b)
if key is None:
assert value is None
k = None
v = None
else:
# This is inline in_proj function with in_proj_weight and in_proj_bias
_b = in_proj_bias
_start = embed_dim
_end = None
_w = in_proj_weight[_start:, :]
if _b is not None:
_b = _b[_start:]
k, v = linear(key, _w, _b).chunk(2, dim=-1)
else:
# This is inline in_proj function with in_proj_weight and in_proj_bias
_b = in_proj_bias
_start = 0
_end = embed_dim
_w = in_proj_weight[_start:_end, :]
if _b is not None:
_b = _b[_start:_end]
q = linear(query, _w, _b)
# This is inline in_proj function with in_proj_weight and in_proj_bias
_b = in_proj_bias
_start = embed_dim
_end = embed_dim * 2
_w = in_proj_weight[_start:_end, :]
if _b is not None:
_b = _b[_start:_end]
k = linear(key, _w, _b)
# This is inline in_proj function with in_proj_weight and in_proj_bias
_b = in_proj_bias
_start = embed_dim * 2
_end = None
_w = in_proj_weight[_start:, :]
if _b is not None:
_b = _b[_start:]
v = linear(value, _w, _b)
else:
q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
len1, len2 = q_proj_weight_non_opt.size()
assert len1 == embed_dim and len2 == query.size(-1)
k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
len1, len2 = k_proj_weight_non_opt.size()
assert len1 == embed_dim and len2 == key.size(-1)
v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
len1, len2 = v_proj_weight_non_opt.size()
assert len1 == embed_dim and len2 == value.size(-1)
if in_proj_bias is not None:
q = linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
k = linear(key, k_proj_weight_non_opt,
in_proj_bias[embed_dim:(embed_dim * 2)])
v = linear(value, v_proj_weight_non_opt,
in_proj_bias[(embed_dim * 2):])
else:
q = linear(query, q_proj_weight_non_opt, in_proj_bias)
k = linear(key, k_proj_weight_non_opt, in_proj_bias)
v = linear(value, v_proj_weight_non_opt, in_proj_bias)
q = q * scaling
if attn_mask is not None:
assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
if attn_mask.dtype == torch.uint8:
warnings.warn(
"Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
)
attn_mask = attn_mask.to(torch.bool)
if attn_mask.dim() == 2:
attn_mask = attn_mask.unsqueeze(0)
if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
raise RuntimeError(
'The size of the 2D attn_mask is not correct.')
elif attn_mask.dim() == 3:
if list(attn_mask.size()) != [
bsz * num_heads,
query.size(0), key.size(0)
]:
raise RuntimeError(
'The size of the 3D attn_mask is not correct.')
else:
raise RuntimeError(
"attn_mask's dimension {} is not supported".format(
attn_mask.dim()))
# attn_mask's dim is 3 now.
# convert ByteTensor key_padding_mask to bool
if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
warnings.warn(
"Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
)
key_padding_mask = key_padding_mask.to(torch.bool)
if bias_k is not None and bias_v is not None:
if static_k is None and static_v is None:
k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
if attn_mask is not None:
attn_mask = pad(attn_mask, (0, 1))
if key_padding_mask is not None:
key_padding_mask = pad(key_padding_mask, (0, 1))
else:
assert static_k is None, "bias cannot be added to static key."
assert static_v is None, "bias cannot be added to static value."
else:
assert bias_k is None
assert bias_v is None
q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
if k is not None:
k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
if v is not None:
v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
if static_k is not None:
assert static_k.size(0) == bsz * num_heads
assert static_k.size(2) == head_dim
k = static_k
if static_v is not None:
assert static_v.size(0) == bsz * num_heads
assert static_v.size(2) == head_dim
v = static_v
src_len = k.size(1)
if key_padding_mask is not None:
assert key_padding_mask.size(0) == bsz
assert key_padding_mask.size(1) == src_len
if add_zero_attn:
src_len += 1
k = torch.cat([
k,
torch.zeros(
(k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)
],
dim=1)
v = torch.cat([
v,
torch.zeros(
(v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)
],
dim=1)
if attn_mask is not None:
attn_mask = pad(attn_mask, (0, 1))
if key_padding_mask is not None:
key_padding_mask = pad(key_padding_mask, (0, 1))
attn_output_weights = torch.bmm(q, k.transpose(1, 2))
assert list(
attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
attn_output_weights.masked_fill_(attn_mask, float('-inf'))
else:
attn_output_weights += attn_mask
if key_padding_mask is not None:
attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len,
src_len)
attn_output_weights = attn_output_weights.masked_fill(
key_padding_mask.unsqueeze(1).unsqueeze(2),
float('-inf'),
)
attn_output_weights = attn_output_weights.view(bsz * num_heads,
tgt_len, src_len)
attn_output_weights = softmax(attn_output_weights, dim=-1)
attn_output_weights = dropout(attn_output_weights,
p=dropout_p,
training=training)
attn_output = torch.bmm(attn_output_weights, v)
assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
attn_output = attn_output.transpose(0, 1).contiguous().view(
tgt_len, bsz, embed_dim)
attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
if need_weights:
# average attention weights over heads
attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len,
src_len)
return attn_output, attn_output_weights.sum(dim=1) / num_heads
else:
return attn_output, None
def __rsub__(self, other):
if has_torch_function_variadic(self, other):
return handle_torch_function(Tensor.__rsub__, (self, other), self, other)
return _C._VariableFunctions.rsub(self, other)
def __len__(self):
if has_torch_function_unary(self):
return handle_torch_function(Tensor.__len__, (self,), self)
if self.dim() == 0:
raise TypeError("len() of a 0-d tensor")
return self.shape[0]
def __format__(self, format_spec):
if has_torch_function_unary(self):
return handle_torch_function(Tensor.__format__, (self,), self, format_spec)
if self.dim() == 0:
return self.item().__format__(format_spec)
return object.__format__(self, format_spec)
def grad(self):
if has_torch_function_unary(self):
# TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185
return handle_torch_function(Tensor.grad.__delete__, (self,), self) # type: ignore[attr-defined]
del self._grad
def __rmatmul__(self, other):
if has_torch_function_variadic(self, other):
return handle_torch_function(Tensor.__rmatmul__, (self, other), self, other)
return torch.matmul(other, self)
def _reduce_ex_internal(self, proto):
if has_torch_function_unary(self):
return handle_torch_function(Tensor.__reduce_ex__, (self,), self, proto)
check_serializing_named_tensor(self)
# See Note [Don't serialize hooks]
torch.utils.hooks.warn_if_has_hooks(self)
backward_hooks: Dict[Any, Any] = OrderedDict()
# Note: Numpy array is chosen to be the rebuild component for XLA Tensor.
# We considered a few options:
# 1. CPU tensor can't be used here.
# Otherwise in torch.load CPU storage is reconstructed with randomly
# initialized data, moved onto XLA device, and then storage is updated
# to the serialized content. This works perfectly for CPU/CUDA but not XLA.
# XLA tensor is disconnected with storage so it doesn't get the update.
# 2. Python list is not a good fit due to performance reason.
# `tolist()` converts every single element in the tensor into python objects
# and serialize them one by one.
if self.device.type == 'xla':
arg_xla = (self.cpu().numpy(),
self.dtype,
str(self.device),
self.requires_grad)
return (torch._utils._rebuild_xla_tensor, arg_xla)
if self.device.type == 'mlc':
arg_mlc = (self.cpu().numpy(),
self.dtype,
str(self.device),
self.requires_grad)
return (torch._utils._rebuild_mlc_tensor, arg_mlc)
if self.is_quantized:
# quantizer_params can be different type based on torch attribute
quantizer_params: Union[Tuple[torch.qscheme, float, int], Tuple[Any, Tensor, Tensor, int]]
if self.qscheme() == torch.per_tensor_affine:
quantizer_params = (torch.per_tensor_affine,
self.q_scale(),
self.q_zero_point())
elif self.qscheme() in (torch.per_channel_affine, torch.per_channel_affine_float_qparams):
# convert scales and zero points to tuple to avoid recursive calls
# when/if we get multi-axis quantized tensors in the future, the shape
# is recoverable from the main tensor shape
quantizer_params = (torch.per_channel_affine,
self.q_per_channel_scales(),
self.q_per_channel_zero_points(),
self.q_per_channel_axis())
else:
raise RuntimeError(f"Serialization is not supported for tensors of type {self.qscheme()}")
args_qtensor = (self.storage(),
self.storage_offset(),
tuple(self.size()),
self.stride(),
quantizer_params,
self.requires_grad,
backward_hooks)
return (torch._utils._rebuild_qtensor, args_qtensor)
elif self.is_sparse:
if self.layout == torch.sparse_coo:
args_sparse = (self.layout,
(self._indices(),
self._values(),
self.size()))
else:
raise NotImplementedError(
'sparse tensor __reduce_ex__ for layout `%s`' % (self.layout))
return (torch._utils._rebuild_sparse_tensor, args_sparse)
else:
args = (self.storage(),
self.storage_offset(),
tuple(self.size()),
self.stride(),
self.requires_grad,
backward_hooks) # previously was self._backward_hooks
return (torch._utils._rebuild_tensor_v2, args)
def __cuda_array_interface__(self):
"""Array view description for cuda tensors.
See:
https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
"""
if has_torch_function_unary(self):
# TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185
return handle_torch_function(Tensor.__cuda_array_interface__.__get__, (self,), self) # type: ignore[attr-defined]
# raise AttributeError for unsupported tensors, so that
# hasattr(cpu_tensor, "__cuda_array_interface__") is False.
if not self.is_cuda:
raise AttributeError(
"Can't get __cuda_array_interface__ on non-CUDA tensor type: %s "
"If CUDA data is required use tensor.cuda() to copy tensor to device memory." %
self.type()
)
if self.is_sparse:
raise AttributeError(
"Can't get __cuda_array_interface__ on sparse type: %s "
"Use Tensor.to_dense() to convert to a dense tensor first." %
self.type()
)
# RuntimeError, matching tensor.__array__() behavior.
if self.requires_grad:
raise RuntimeError(
"Can't get __cuda_array_interface__ on Variable that requires grad. "
"If gradients aren't required, use var.detach() to get Variable that doesn't require grad."
)
# CUDA devices are little-endian and tensors are stored in native byte
# order. 1-byte entries are endian-agnostic.
typestr = {
torch.complex64: "<c8",
torch.complex128: "<c16",
torch.float16: "<f2",
torch.float32: "<f4",
torch.float64: "<f8",
torch.uint8: "|u1",
torch.int8: "|i1",
torch.int16: "<i2",
torch.int32: "<i4",
torch.int64: "<i8",
}[self.dtype]
itemsize = self.storage().element_size()
shape = tuple(self.shape)
if self.is_contiguous():
# __cuda_array_interface__ v2 requires the strides to be omitted
# (either not set or set to None) for C-contiguous arrays.
strides = None
else:
strides = tuple(s * itemsize for s in self.stride())
data_ptr = self.data_ptr() if self.numel() > 0 else 0
data = (data_ptr, False) # read-only is false
return dict(typestr=typestr, shape=shape, strides=strides, data=data, version=2)
def __repr__(self):
if has_torch_function_unary(self):
return handle_torch_function(Tensor.__repr__, (self,), self)
# All strings are unicode in Python 3.
return torch._tensor_str._str(self)
def baz(a, b):
"""A function with multiple arguments"""
if type(a) is not Tensor or type(b) is not Tensor and has_torch_function((a, b)):
return handle_torch_function(baz, (a, b), a, b)
return a + b
def norm(self, p="fro", dim=None, keepdim=False, dtype=None):
r"""See :func:`torch.norm`"""
if has_torch_function_unary(self):
return handle_torch_function(Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype)
return torch.norm(self, p, dim, keepdim, dtype=dtype)
def __rsub__(self, other):
relevant_args = (self, other)
from torch.overrides import has_torch_function, handle_torch_function
if type(self) is not Tensor and type(other) is not Tensor and has_torch_function(relevant_args):
return handle_torch_function(Tensor.__rsub__, relevant_args, self, other)
return _C._VariableFunctions.rsub(self, other)
def resize(self, *sizes):
if has_torch_function_unary(self):
return handle_torch_function(Tensor.resize, (self,), self, *sizes)
warnings.warn("non-inplace resize is deprecated")
from torch.autograd._functions import Resize
return Resize.apply(self, sizes)
def __hash__(self):
relevant_args = (self,)
from torch.overrides import has_torch_function, handle_torch_function
if type(self) is not Tensor and has_torch_function(relevant_args):
return handle_torch_function(Tensor.__hash__, relevant_args, self)
return id(self)
def requires_grad_(self, requires_grad=True):
return handle_torch_function(torch.Tensor.requires_grad_,
(self, requires_grad), self,
requires_grad)
