File size: 869 Bytes
c4cbbbc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | from typing import Tuple, Optional
import torch
def make_tensor_strided_like(
q_tensor: torch.Tensor,
o_shape: Tuple[int, ...],
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
):
"""
Create an empty tensor with the given shape that mimics the layout/strides of
the provided `q_tensor` as closely as possible.
"""
q_strides = q_tensor.stride()
rank_out = len(o_shape)
order = tuple(sorted(range(min(len(q_strides), rank_out)), key=lambda i: q_strides[i]))
strides = [0] * rank_out
current = 1
for dim in order:
strides[dim] = current
current *= o_shape[dim]
return torch.empty_strided(
o_shape,
tuple(strides),
dtype=dtype if dtype is not None else q_tensor.dtype,
device=device if device is not None else q_tensor.device,
)
|