def _reduction(
a: Tensor,
prim: Callable,
*,
has_identity: bool = True,
accepts_dim_tuple: bool = True, # to handle min/argmin that accept single dim only
dims: Optional[DimsType] = None,
keepdims: bool = False,
dtype: Optional[torch.dtype] = None, # should be specified for ops that support it
out: Optional[Tensor] = None,
output_dtype_kind: REDUCTION_OUTPUT_TYPE_KIND,
): # it is usually SAME, but I want
# ref writers to actually think about what to put here
assert isinstance(a, TensorLike)
if out is not None:
assert isinstance(out, TensorLike)
if dtype is not None:
# TODO - this is true for eager mode currently, but it's wrong behavior for complex norms
if dtype != out.dtype:
raise RuntimeError(
"dtype argument and out dtype must match in reduction"
)
if not accepts_dim_tuple:
assert dims is None or isinstance(dims, int)
if isinstance(dims, int):
dims = (dims,) # type: ignore[assignment]
dims = utils.reduction_dims(a.shape, dims)
if not has_identity:
valid_shape = all(a.shape[i] for i in range(a.ndim) if i in dims)
if not valid_shape:
raise RuntimeError(
"reducing over zero-size dimension for reduction operation without identity"
)
# even though some reductions, like amin or amax, don't strictly require type promotion,
# all the math ops (including comparisons) are still defined only for a computation type,
# so promotion will still happen. We are doing it explicitly here
inp_dtype = dtype if dtype is not None else a.dtype
computation_dtype = utils._get_computation_dtype(inp_dtype)
a_converted = prims.convert_element_type(a, computation_dtype)
result = prim(a_converted, dims)
if keepdims:
output_shape = [a.shape[i] if i not in dims else 1 for i in range(a.ndim)]
broadcast_dims = [i for i in range(a.ndim) if i not in dims]
result = prims.broadcast_in_dim(result, output_shape, broadcast_dims)
if out is not None:
if dtype is None:
if output_dtype_kind == REDUCTION_OUTPUT_TYPE_KIND.SAME:
if out.dtype != a.dtype:
raise RuntimeError("Expected the dtype for input and out to match")
elif output_dtype_kind == REDUCTION_OUTPUT_TYPE_KIND.ALWAYS_BOOL:
if out.dtype != torch.bool:
raise RuntimeError("Expected the dtype for input and out to match")
out = _maybe_resize_out(out, result.shape)
return copy_to(out, result, allow_cross_device=False) # type: ignore[arg-type]
if output_dtype_kind == REDUCTION_OUTPUT_TYPE_KIND.SAME:
result_dtype = dtype if dtype else a.dtype
result = prims.convert_element_type(result, result_dtype)
return result
def isclose(
a: TensorLikeType,
b: TensorLikeType,
rtol: float = 1e-05,
atol: float = 1e-08,
equal_nan: bool = False,
) -> TensorLikeType:
if a.dtype != b.dtype:
msg = "Attempting to compare tensors of different dtypes {0} and {1}!".format(
a.dtype, b.dtype)
raise ValueError(a, b)
if rtol < 0:
msg = "rtol must be greater than or equal to zero, but got {0}!".format(
rtol)
if atol < 0:
msg = "atol must be greater than or equal to zero, but got {0}!".format(
atol)
close = eq(a, b)
if equal_nan and (utils.is_float_dtype(a.dtype)
or utils.is_complex_dtype(a.dtype)):
close = logical_or(close, logical_and(isnan(a), isnan(b)))
# Note: In case of zero tolerances the closeness inequality degenerates to an equality check.
# In this case, the short-circuit prevents false positives as detailed in the paragraph below.
if atol == 0 and rtol == 0:
return close
# Note [closeness error computation]
# atol and rtol are provided as doubles, so the computation
# rtol * other will produce a float or complex tensor.
# When the difference (self - other) is compared to it then the
# tensor representing the difference will also be cast to float or complex.
# However, since (self - other) in uint8 is very likely to produce a
# negative value, this moves the cast forward so the difference is
# always computed in a float or complex type.
# If the values of the integer tensors cannot be exactly represented
# by the default scalar type then this may cause an incorrect result.
if not utils.is_float_dtype(a.dtype) and not utils.is_complex_dtype(
a.dtype):
a = prims.convert_element_type(a, torch.get_default_dtype())
b = prims.convert_element_type(b, torch.get_default_dtype())
allowed_error = add(atol, abs(mul(b, rtol)))
actual_error = abs(sub(a, b))
# Computes finite closeness
result = logical_or(
close,
logical_and(isfinite(actual_error), le(actual_error, allowed_error)))
return result
def _convert(x):
if isinstance(x, TensorLike):
return prims.convert_element_type(x, dtype)
elif isinstance(x, Number):
typ = utils.dtype_to_type(dtype)
return typ(x)
return x
def vector_norm(
x: TensorLikeType,
ord: float = 2.0,
dim: Optional[DimsType] = None,
keepdim: bool = False,
*,
dtype: Optional[torch.dtype] = None,
) -> Tensor:
# Checks
check_fp_or_complex(x.dtype, "linalg.vector_norm")
if isinstance(dim, int):
dim = [dim] # type: ignore[assignment]
elif not isinstance(dim, List) and dim is not None:
# refs.amin just accepts List rather than DimType (Tuple)
dim = list(dim) # type: ignore[assignment]
if x.numel() == 0 and (ord < 0.0 or ord == float("inf")):
check(
dim is not None and len(dim) != 0,
lambda: f"linalg.vector_norm cannot compute the {ord} norm on an empty tensor "
"because the operation does not have an identity",
)
shape = x.shape
assert dim is not None # mypy does not seem to be able to see through check?
for d in dim:
check(
shape[d] != 0,
lambda: f"linalg.vector_norm cannot compute the {ord} norm on the "
f"dimension {d} because this dimension is empty and the "
"operation does not have an identity",
)
check_norm_dtype(dtype, x.dtype, "linalg.vector_norm")
computation_dtype, result_dtype = reduction_dtypes(
x, REDUCTION_OUTPUT_TYPE_KIND.COMPLEX_TO_FLOAT, dtype
)
to_result_dtype = partial(prims.convert_element_type, dtype=result_dtype)
# Implementation
if ord == 0.0:
return refs.sum(refs.ne(x, 0.0), dim=dim, keepdim=keepdim, dtype=result_dtype)
elif ord == float("inf"):
return to_result_dtype(refs.amax(torch.abs(x), dim=dim, keepdim=keepdim))
elif ord == float("-inf"):
return to_result_dtype(refs.amin(torch.abs(x), dim=dim, keepdim=keepdim))
else:
# From here on the computation dtype is important as the reduction is non-trivial
x = prims.convert_element_type(x, computation_dtype)
reduce_sum = partial(refs.sum, dim=dim, keepdim=keepdim)
# Avoid computing a sqrt in abs and then squaring (more stable)
# This could potentially be done for complex dtypes as
# x = torch.real(torch.conj(x) * x))
# and it should be more stable, but it's not clear whether it'll be faster on, say
# CPU (abs is 1 vectorised operation), so leaving it just for real dtypes for now
if not (ord % 2.0 == 0.0 and is_float_dtype(x.dtype)):
x = torch.abs(x)
return to_result_dtype(torch.pow(reduce_sum(torch.pow(x, ord)), 1.0 / ord))
def _maybe_promote_tensor_fft(t: TensorLikeType,
require_complex: bool = False) -> TensorLikeType:
"""Helper to promote a tensor to a dtype supported by the FFT primitives"""
cur_type = t.dtype
new_type = _promote_type_fft(cur_type, require_complex)
if cur_type == new_type:
return t
return prims.convert_element_type(t, new_type)
def _maybe_convert_to_dtype(
a: Union[TensorLikeType, NumberType, Sequence],
dtype: torch.dtype) -> Union[TensorLikeType, NumberType, Sequence]:
if isinstance(a, TensorLike):
if a.dtype != dtype:
return prims.convert_element_type(a, dtype)
return a
if isinstance(a, Number):
return utils.dtype_to_type(dtype)(a)
if isinstance(a, Sequence):
return tuple(_maybe_convert_to_dtype(x, dtype) for x in a)
raise ValueError(
"Received type {0} that is neither a tensor or a number!".format(
type(a)))
def _maybe_convert_to_dtype(
a: Union[TensorLikeType, NumberType, Sequence],
dtype: torch.dtype) -> Union[TensorLikeType, NumberType, Sequence]:
if isinstance(a, TensorLike):
if a.dtype != dtype:
# NOTE: this is incorrect on the CPU
# See https://github.com/pytorch/pytorch/issues/77553
return prims.convert_element_type(a, dtype)
return a
if isinstance(a, Number):
return utils.dtype_to_type(dtype)(a)
if isinstance(a, Sequence):
return tuple(_maybe_convert_to_dtype(x, dtype) for x in a)
raise ValueError(
"Received type {0} that is neither a tensor or a number!".format(
type(a)))
def _maybe_convert_to_dtype(
a: Union[TensorLikeType, NumberType, Sequence, None], dtype: torch.dtype
) -> Union[TensorLikeType, NumberType, Sequence, None]:
import torch._prims as prims
if isinstance(a, TensorLike):
if a.dtype != dtype:
# NOTE: this is incorrect on the CPU
# See https://github.com/pytorch/pytorch/issues/77553
return prims.convert_element_type(a, dtype)
return a
if isinstance(a, Number):
return utils.dtype_to_type(dtype)(a)
if isinstance(a, Sequence):
return tuple(_maybe_convert_to_dtype(x, dtype) for x in a)
# Passthrough None because some functions wrapped with type promotion
# wrapper might have optional args
if a is None:
return None
raise ValueError(
"Received type {0} that is neither a tensor or a number!".format(type(a))
)
def matrix_norm(
A: TensorLikeType,
ord: Union[float, str] = "fro",
dim: DimsType = (-2, -1),
keepdim: bool = False,
*,
dtype: Optional[torch.dtype] = None,
) -> TensorLikeType:
# shape
check_is_matrix(A, "linalg.matrix_norm")
# dim
dim = utils.canonicalize_dims(A.ndim, dim)
if isinstance(dim, int):
dim = (dim, ) # type: ignore[assignment]
check(
len(dim) == 2,
lambda: "linalg.matrix_norm: dim must be a 2-tuple. Got {dim}")
check(
dim[0] != dim[1],
lambda:
"linalg.matrix_norm: dims must be different. Got ({dim[0]}, {dim[1]})",
)
# dtype arg
check_norm_dtype(dtype, A.dtype, "linalg.matrix_norm")
if isinstance(ord, str):
# ord
check(
ord in ("fro", "nuc"),
lambda: "linalg.matrix_norm: Order {ord} not supported.",
)
# dtype
check_fp_or_complex(A.dtype,
"linalg.matrix_norm",
allow_low_precision_dtypes=ord != "nuc")
if ord == "fro":
return vector_norm(A, 2, dim, keepdim, dtype=dtype)
else: # ord == "nuc"
if dtype is not None:
A = prims.convert_element_type(A, dtype)
perm = backshift_permutation(dim[0], dim[1], A.ndim)
result = torch.sum(svdvals(prims.transpose(A, perm)), -1, keepdim)
if keepdim:
inv_perm = inverse_permutation(perm)
result = prims.transpose(torch.unsqueeze(result, -1), inv_perm)
return result
else:
# ord
abs_ord = abs(ord)
check(
abs_ord in (2, 1, float("inf")),
lambda: "linalg.matrix_norm: Order {ord} not supported.",
)
# dtype
check_fp_or_complex(A.dtype,
"linalg.matrix_norm",
allow_low_precision_dtypes=ord != 2)
max_min = partial(torch.amax if ord > 0.0 else torch.amin,
keepdim=keepdim)
if abs_ord == 2.0:
if dtype is not None:
A = prims.convert_element_type(A, dtype)
perm = backshift_permutation(dim[0], dim[1], A.ndim)
result = max_min(svdvals(prims.transpose(A, perm)), dim=-1)
if keepdim:
inv_perm = inverse_permutation(perm)
result = prims.transpose(torch.unsqueeze(result, -1), inv_perm)
return result
else: # 1, -1, inf, -inf
dim0, dim1 = dim
if abs_ord == float("inf"):
dim0, dim1 = dim1, dim0
if not keepdim and (dim0 < dim1):
dim1 -= 1
return max_min(
vector_norm(A, 1.0, dim=dim0, keepdim=keepdim, dtype=dtype),
dim1)