def prelu(a: TensorLikeType, weight: TensorLikeType) -> TensorLikeType:
"""
Reference implementation of torch.nn.functional.prelu
"""
check(
isinstance(a, TensorLike),
lambda: f"prelu: Expected `a` to be tensor, but got: {type(a)}",
)
check(
isinstance(weight, TensorLike),
lambda: f"prelu: Expected `weight` to be tensor, but got: {type(weight)}",
)
if weight.numel() != 1:
check(a.ndim > 0, lambda: "Not allow zero-dim input tensor.")
channel_size = a.shape[1] if a.ndim >= 2 else 1
check(
weight.numel() == channel_size,
lambda: f"Mismatch of parameter numbers and input channel size. Found parameter numbers ="
f" {weight.numel()} and channel size = {channel_size}.",
)
check(
weight.ndim == 0 or weight.ndim == 1,
lambda: f"prelu: Expected `weight` to be a scalar or 1D tensor, but got: "
f"ndim = {weight.ndim}",
)
weight = prims.broadcast_in_dim(
weight, a.shape, tuple() if weight.ndim == 0 else (1,)
)
return refs.where(a > 0, a, a * weight)
def _reshape_meta(a: TensorLikeType, shape: ShapeType):
assert isinstance(a, TensorLike)
utils.validate_shape(shape)
# Validates the tensor and the requested shape have the
# same number of elements
numel = reduce(operator.mul, shape)
if numel != a.numel():
msg = "Attempting to reshape a tensor with {0} elements to a shape with {1} elements!".format(
a.numel(), numel)
raise ValueError(msg)
return TensorMeta(a,
shape=shape,
strides=utils.make_contiguous_strides_for(shape))
def _maybe_resize_out(out: TensorLikeType, shape):
if out.numel() == 0:
return prims.resize(out, shape)
if out.numel() != reduce(operator.mul, shape, 1):
msg = (
"An output with one or more elements was resized since it had shape {0} "
"which does not match the required output shape {1}. "
"This behavior is deprecated, and in a future PyTorch release outputs will not "
"be resized unless they have zero elements. "
"You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0)."
.format(str(out.shape), str(shape)))
warnings.warn(msg)
return prims.resize(out, shape)
return out
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 _reshape_meta(a: TensorLikeType, shape: Sequence):
assert isinstance(a, TensorLike)
utils.validate_shape(shape)
# Validates the tensor and the requested shape have the
# same number of elements
numel = reduce(lambda acc, x: acc * x, shape)
assert a.numel() == numel
def _copy_to_meta(a: TensorLikeType, b: TensorLikeType):
assert isinstance(a, TensorLike)
assert isinstance(b, TensorLike)
# Validates the cast is safe
# TODO: move this as an option on the reference
# a_typ = utils.dtype_to_type(a.dtype)
# b_typ = utils.dtype_to_type(b.dtype)
# if a_typ is not utils.get_higher_type(a_typ, b_typ):
# raise RuntimeError(str(b.dtype), " can't be cast safely to ", str(a.dtype), "!")
# Validates the tensors have the same number of elements
if a.numel() != b.numel():
msg = "Attempting to copy {0} elements to a tensor with {1} elements!".format(
b.numel(), a.numel())
raise RuntimeError(msg)
return a
def _reshape_view_helper(a: TensorLikeType, shape: ShapeType, *,
allow_copy: bool) -> TensorLikeType:
# NOTE: Reshape may be given a shape with a -1 length
# This indicates that the dimension's length should be inferred
# Creates a valid shape
for idx in range(len(shape)):
if shape[idx] == -1:
# Verifies there's only one dimension of length -1 in the shape
if shape.count(-1) > 1:
msg = "Can only infer the length of one dimension, but got shape {0}!".format(
str(shape))
raise ValueError(msg)
# TODO: improve error message
if a.numel() > 0:
length = reduce(operator.floordiv,
(x for x in shape if x != -1), a.numel())
else:
msg = "Cannot reshape a tensor of zero elements into shape {0} because the unspecified length is ambiguous!".format(
str(shape))
raise ValueError(msg)
shape = list(shape)
shape[idx] = length
break
# Short-circuits if shape is the same
utils.validate_shape(shape)
if tuple(a.shape) == tuple(shape):
return prims.view_of(a)
numel = reduce(operator.mul, shape) if len(shape) > 0 else 1
if a.numel() != numel:
msg = "Attempting to reshape a tensor with shape {0} and {1} elements to a shape {2} with {3} elements!".format(
str(a.shape), a.numel(), str(shape), numel)
raise ValueError(msg)
# Special-cases tensors with no elements
if a.numel() == 0:
return as_strided(a, shape, utils.make_contiguous_strides_for(shape))
# Special-cases reshaping zero dim tensors
if a.ndim == 0:
_a = a
for length in shape:
assert length == 1
_a = unsqueeze(_a, -1)
return _a
# Special-cases reshaping to zero dim tensors
if len(shape) == 0:
_a = a
for length in a.shape:
assert length == 1
_a = squeeze(_a, -1)
return _a
# Handles general case: a 1+D tensor reshaped into a distinct 1+D shape
# NOTE [Reshape Algorithm]
# This algorithm works by attempting to greedily construct the desired dimensions in
# the output shape, left to right. It does this by, conceptually, accumulating
# dimensions of the original tensor, also left to right, until the dimension
# can be constructed using prims.split_dim.
# The algorithm also has special handling for tail squeezes/unsqueezes, like
# if a reshape from (5, 5) to (5, 5, 1) or vice versa.
#
# This algorithm does not flatten the original tensor and then split dims as appropriate
# because that would create copies more often than this algorithm. flatten is the only
# operation below which can create a view or a copy, and while it prefers creating
# views it may sometimes create a copy if the tensor's strides do not permit a view.
# As a result, this algorithm tries to minimize flattening.
#
# Note that a better version of this algorithm may exist. Regions which could be
# flattened without creating a copy can be identified in advance, and that might
# allow fewer flatten calls or faster short-circuiting to make a copy.
idx = 0
a_ = a
for length in shape:
# Handles tail unsqueezes
if idx >= a_.ndim:
assert length == 1
a_ = unsqueeze(a_, -1)
idx = idx + 1
continue
# Skips dimensions that are already the correct length
if length == a_.shape[idx]:
idx = idx + 1
continue
# Gathers enough original dimensions such that this new dimension can be created
# Note that this accumulation will terminate because we've verified a and the shape
# specify the same number of elements above
accum = a_.shape[idx]
end = idx
while accum % length != 0:
end = end + 1
accum = accum * a_.shape[end]
if end != idx:
# NOTE: in this case multiple dimensions must be flatten to create the desired dimension
# This flattening is why reshape sometimes creates a copy -- because flattening
# may return a view of a copy
# Checks if collapse can be a view and short-circuits to copying reshape if it can't
new_shape, new_strides = prims._collapse_view_helper(
a_, idx, end + 1)
if new_shape is None:
if allow_copy:
return prims.reshape(a, shape)
msg = "Cannot view a tensor with shape {0} and strides {1} as a tensor with shape {2}!".format(
a.shape, a.stride(), shape)
raise ValueError(msg)
a_ = flatten(a_, idx, end)
# Splits the (possibly flattened) dimension to create the desired dim length
if accum != length:
a_ = prims.split_dim(a_, idx, length)
idx = idx + 1
# Squeezes tail
while idx < a_.ndim:
assert a_.shape[idx] == 1
a_ = squeeze(a_, idx)
return a_
def _resize_meta(a: TensorLikeType, shape: Union[torch.Size, List[int],
Tuple[int, ...]]):
assert a.numel() == 0
return TensorMeta(a,
shape=shape,
strides=utils.make_contiguous_strides_for(shape))