File size: 3,884 Bytes
a40a8f5 | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Sequence
from . import Dim
import torch # noqa: TC002
# NB: The old code represented dimension was from as negative number, so we
# follow this convention even though it shouldn't be necessary now
class DimEntry:
# The dimension this is from the rhs, or a FCD
data: Dim | int
def __init__(self, data: Dim | int | None = None) -> None:
from . import Dim
if type(data) is int:
if data >= 0:
raise AssertionError(f"Expected negative int, got {data}")
elif data is None:
data = 0
else:
if not isinstance(data, Dim):
raise AssertionError(f"Expected Dim, got {type(data)}")
self.data = data
def __eq__(self, other: object) -> bool:
if not isinstance(other, DimEntry):
return False
# Use 'is' for Dim objects to avoid triggering __torch_function__
# Use '==' only for positional (int) comparisons
if self.is_positional() and other.is_positional():
# Both are positional (ints)
return self.data == other.data
elif not self.is_positional() and not other.is_positional():
# Both are Dim objects - use 'is' to avoid __eq__
return self.data is other.data
else:
# One is positional, one is Dim - they can't be equal
return False
def is_positional(self) -> bool:
return type(self.data) is int and self.data < 0
def is_none(self) -> bool:
# Use isinstance to check for Dim objects, avoid triggering __torch_function__
from . import Dim
if isinstance(self.data, Dim):
# This is a Dim object, it can't be "none" (which is represented by 0)
return False
else:
# This is an int or other type
return self.data == 0
def position(self) -> int:
if not isinstance(self.data, int):
raise AssertionError(f"Expected int, got {type(self.data)}")
return self.data
def dim(self) -> Dim:
if isinstance(self.data, int):
raise AssertionError("Expected Dim, got int")
return self.data
def __repr__(self) -> str:
return repr(self.data)
def ndim_of_levels(levels: Sequence[DimEntry]) -> int:
r = 0
for l in levels:
if l.is_positional():
r += 1
return r
def _match_levels(
tensor: torch.Tensor,
from_levels: list[DimEntry],
to_levels: list[DimEntry],
drop_levels: bool = False,
) -> torch.Tensor:
"""
Reshape a tensor to match target levels using as_strided.
Args:
tensor: Input tensor to reshape
from_levels: Current levels of the tensor
to_levels: Target levels to match
drop_levels: If True, missing dimensions are assumed to have stride 0
Returns:
Reshaped tensor
"""
if from_levels == to_levels:
return tensor
sizes = tensor.size()
strides = tensor.stride()
if not drop_levels:
if len(from_levels) > len(to_levels):
raise AssertionError("Cannot expand dimensions without drop_levels")
new_sizes = []
new_strides = []
for level in to_levels:
# Find index of this level in from_levels
try:
idx = from_levels.index(level)
except ValueError:
# Level not found in from_levels
if level.is_positional():
new_sizes.append(1)
else:
new_sizes.append(level.dim().size)
new_strides.append(0)
else:
new_sizes.append(sizes[idx])
new_strides.append(strides[idx])
return tensor.as_strided(new_sizes, new_strides, tensor.storage_offset())
|