File size: 7,011 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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | from __future__ import annotations
from typing import Any, TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Sequence
import torch # noqa: TC002
from ._dim_entry import _match_levels, DimEntry, ndim_of_levels
def _wrap_dim(arg: Any, orig_ndim: int, allow_none: bool = True) -> DimEntry:
"""
Convert various dimension representations to DimEntry.
Args:
arg: The argument to convert (Dim, int, or other)
orig_ndim: Original number of dimensions
allow_none: Whether to allow None values
Returns:
DimEntry representation of the dimension
"""
from . import Dim
if arg is None and allow_none:
return DimEntry() # None entry
elif isinstance(arg, Dim):
return DimEntry(arg)
elif isinstance(arg, int):
if arg < 0:
pos = arg
else:
pos = arg - orig_ndim
return DimEntry(pos)
else:
return DimEntry()
def order(
tensor_or_dim: torch.Tensor | Any, *dims: Any | Sequence[Any]
) -> torch.Tensor:
"""
Reorder the dimensions of a tensor or create a tensor from a dimension.
It allows reordering tensor dimensions using first-class dimensions and
positional indices.
Args:
tensor_or_dim: Input tensor with first-class dimensions, or a Dim object
*dims: Dimensions or sequences of dimensions specifying the new order
Returns:
Tensor with reordered dimensions
Examples:
>>> import torch
>>> from functorch.dim import dims
>>> batch, channel, height, width = dims(4)
>>> x = torch.randn(2, 3, 4, 5)[batch, channel, height, width]
>>> # Reorder to [height, width, batch, channel]
>>> y = order(x, height, width, batch, channel)
"""
from . import Dim, DimList, Tensor
# Handle first argument - tensor or dimension
if isinstance(tensor_or_dim, Tensor):
# First-class tensor
orig_levels = tensor_or_dim._levels[:]
data = tensor_or_dim._tensor
has_device = tensor_or_dim._has_device
elif isinstance(tensor_or_dim, Dim):
# Single dimension - create range tensor
orig_levels = [DimEntry(tensor_or_dim)]
data = tensor_or_dim._get_range()
has_device = False
else:
raise ValueError("First argument must be a Tensor or Dim object")
flat_positional_dims = []
to_flatten = [] # List of (start_index, length) pairs for flattening
levels = orig_levels[:]
orig_ndim = ndim_of_levels(levels)
def append_dim(d: DimEntry) -> None:
"""Add a dimension to the reordering, removing it from available levels."""
try:
idx = levels.index(d)
except ValueError:
idx = None
if idx is None:
if d.is_positional():
raise ValueError(
f"tensor has {orig_ndim} positional dimensions, but {d.position() + orig_ndim} specified, "
f"or it was specified twice"
)
else:
raise ValueError(
f"tensor does not contain dim {d.dim()} or it was specified twice"
)
levels[idx] = DimEntry()
flat_positional_dims.append(d)
n_new_positional = 0
# Process each dimension argument
for arg in dims:
entry = _wrap_dim(arg, orig_ndim, False)
if not entry.is_none():
append_dim(entry)
n_new_positional += 1
elif isinstance(arg, DimList):
# Handle DimList
for dim in arg._dims:
append_dim(DimEntry(dim))
n_new_positional += 1
else:
# Handle sequences of dimensions for flattening
n_new_positional += 1
if not hasattr(arg, "__iter__"):
raise ValueError("expected a Dim, List[Dim], or Sequence[Dim]")
# Convert to list to get length
seq = list(arg)
to_flatten.append((len(flat_positional_dims), len(seq)))
for item in seq:
entry = _wrap_dim(item, orig_ndim, False)
if entry.is_none():
raise ValueError("expected a Dim or int")
append_dim(entry)
# Build new level ordering
insert_point = -1
new_levels: list[DimEntry] = []
# Add remaining (non-reordered) levels, finding insertion point for new dimensions
for level in levels:
if level.is_none():
continue
if level.is_positional():
if insert_point == -1:
insert_point = len(new_levels)
new_levels.extend(flat_positional_dims)
new_levels.append(level)
# If no positional dimensions found, append new dims at the end
if insert_point == -1:
insert_point = len(new_levels)
new_levels.extend(flat_positional_dims)
# Match tensor to new level structure
if data is None:
raise AssertionError("Cannot reorder None tensor")
ndata = _match_levels(data, orig_levels, new_levels)
# Handle dimension flattening if requested
if to_flatten:
# Now build the reshape target
view_shape = []
sizes = ndata.size()
# Add dimensions before the reordered ones
for i in range(insert_point):
view_shape.append(sizes[i])
# Process flattening groups
i = 0
for start_idx, length in to_flatten:
# Add individual dims before this flattening group
while i < start_idx:
view_shape.append(sizes[insert_point + i])
i += 1
# Flatten the group
new_size = 1
for j in range(length):
new_size *= sizes[insert_point + i + j]
view_shape.append(new_size)
i += length
# Add remaining individual dims
while i < len(flat_positional_dims):
view_shape.append(sizes[insert_point + i])
i += 1
# Add dimensions after the reordered ones
for i in range(insert_point + len(flat_positional_dims), len(levels)):
view_shape.append(sizes[i])
# Update levels by removing flattened dimensions
n_to_remove = len(flat_positional_dims) - n_new_positional
if n_to_remove > 0:
# Remove flattened levels
new_levels = (
new_levels[:insert_point] + new_levels[insert_point + n_to_remove :]
)
ndata = ndata.reshape(view_shape)
# Renumber positional dimensions (negative indexing from the right)
seen = 0
for i in range(len(new_levels) - 1, -1, -1):
if new_levels[i].is_positional() or (
i >= insert_point and i < insert_point + n_new_positional
):
seen -= 1
new_levels[i] = DimEntry(seen)
result = Tensor.from_positional(ndata, new_levels, has_device)
return result # type: ignore[return-value]
|