def _concatenate_meta(tensors: Sequence[TensorLikeType],
dim: int) -> TensorLikeType:
if len(tensors) == 0:
msg = "concatenate expects at least one tensor, but received zero!"
raise ValueError(msg)
for tensor in tensors:
assert isinstance(tensor, TensorLike)
utils.check_same_dtype(*tensors)
utils.check_same_device(*tensors, allow_cpu_scalar_tensors=False)
shape = tensors[0].shape
utils.validate_idx(tensors[0].ndim, dim)
# Verifies same shape (except in the concat dimension)
concat_length = 0
for tensor in tensors:
for idx, (common_length, length) in enumerate(zip(shape,
tensor.shape)):
if idx == dim:
concat_length = concat_length + length
else:
assert length == common_length
new_shape = list(tensors[0].shape).copy()
new_shape[dim] = concat_length
return TensorMeta(
tensors[0],
shape=new_shape,
strides=utils.make_contiguous_strides_for(new_shape),
)
def _concatenate_meta(tensors: Sequence[TensorLikeType],
dim: int) -> TensorLikeType:
assert len(tensors) > 0
for tensor in tensors:
assert isinstance(tensor, TensorLike)
utils.check_same_dtype(tensors)
utils.check_same_device(tensors, allow_scalars=False)
shape = tensors[0].shape
utils.validate_idx(shape, dim)
# Verifies same shape (except in the concat dimension)
concat_length = 0
for tensor in tensors:
for idx, (common_length, length) in enumerate(zip(shape,
tensor.shape)):
if idx == dim:
concat_length = concat_length + length
else:
assert length == common_length
new_shape = list(tensors[0].shape).copy()
new_shape[dim] = concat_length
return TensorMeta(
tensors[0],
shape=new_shape,
strides=utils.make_contiguous_strides_for(new_shape),
)
def _elementwise_meta(*args, type_promotion):
"""
Meta function for elementwise operations that produce outputs in the same dtype
as their inputs.
Stride logic is currently incorrect.
"""
assert len(args) > 0
utils.check_same_device(*args, allow_cpu_scalar_tensors=True)
utils.check_same_shape(*args, allow_cpu_scalar_tensors=True)
utils.check_same_dtype(*args)
strides = None
tensor = None
number = None
for arg in args:
if isinstance(arg, TensorLike):
if strides is None:
strides = arg.stride()
if tensor is None:
tensor = arg
if arg.stride() != strides:
return TensorMeta(arg,
strides=utils.make_contiguous_strides_for(
arg.shape))
elif isinstance(arg, Number):
if number is None:
number = arg
# TODO: fix strides
if tensor is not None:
if 0 in tensor.stride() and tensor.numel() > 0:
return TensorMeta(tensor,
strides=utils.make_contiguous_strides_for(
tensor.shape))
else:
return TensorMeta(tensor)
return TensorMeta(number)
def _reduction_meta(inp, dims, *, output_dtype=None):
"""
Meta function for single output reduction operations
Stride logic is incorrect
"""
assert isinstance(inp, TensorLike)
if output_dtype is None:
output_dtype = inp.dtype
output_shape = utils.compute_reduction_output_shape(inp.shape, dims)
return TensorMeta(
shape=output_shape,
strides=utils.make_contiguous_strides_for(output_shape),
dtype=output_dtype,
device=inp.device,
)
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 _resize_meta(a: TensorLikeType, shape: Union[torch.Size, List[int],
Tuple[int, ...]]):
return TensorMeta(a,
shape=shape,
strides=utils.make_contiguous_strides_for(shape))
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_
