Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / region_tensor.py: 90%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Symbolic checkpoint tensors, which track where every element of a converted tensor is read from. 

16 

17A weight conversion cuts, stacks, reorders and reshapes the tensors of a checkpoint into the tensors of 

18a model. Given :class:`RegionTensor` inputs instead of real tensors, the same conversion code computes 

19no values. For every element of each result it tracks which checkpoint tensor and index the element is 

20copied from, and which elementwise scalar operations are applied to it on the way. 

21:meth:`RegionTensor.remap_blocks` hands that over as :class:`RemapBlock` entries, which 

22:class:`RemapLoadPlanner` loads without ever assembling a whole checkpoint tensor or a whole result. 

23 

24What cannot be tracked that way raises :class:`UnsupportedRegionOp`: an operation that combines two 

25tensors elementwise, reads values, writes in place, or indexes with tensors. 

26""" 

27import math 

28import operator 

29from collections.abc import Callable, Iterator, Mapping, Sequence 

30from dataclasses import dataclass, replace 

31from functools import partial 

32from typing import Any 

33 

34import torch # pylint: disable=forbidden-backend-import 

35 

36from hyper_parallel.core.distributed_checkpoint.remap_planner import RemapBlock 

37 

38# A region cut into more blocks than this is not worth describing block by block. The conversion is 

39# run on the real tensors instead. 

40MAX_BLOCKS = 1 << 16 

41 

42 

43class UnsupportedRegionOp(TypeError): 

44 """Raised by an operation whose result cannot be described as regions of checkpoint tensors.""" 

45 

46 

47@dataclass(frozen=True) 

48class RegionBlock: 

49 """ 

50 A box of a symbolic tensor, and where in one checkpoint tensor each element of it comes from. 

51 

52 Attributes: 

53 offsets (tuple[int, ...]): Where the box starts in the symbolic tensor. 

54 lengths (tuple[int, ...]): Size of the box, positive in every dimension. 

55 source (str): Name of the checkpoint tensor. 

56 base (tuple[int, ...]): Checkpoint index of the first element of the box. 

57 steps (tuple[tuple[int, ...], ...]): One column per dimension of the symbolic tensor, giving how 

58 far along each checkpoint dimension one step along that dimension moves. 

59 dtype (torch.dtype): Dtype of the values once ``post`` has been applied. 

60 post (tuple[Callable[[Any], Any], ...]): Elementwise operations applied in order. Default (). 

61 """ 

62 

63 offsets: tuple[int, ...] 

64 lengths: tuple[int, ...] 

65 source: str 

66 base: tuple[int, ...] 

67 steps: tuple[tuple[int, ...], ...] 

68 dtype: torch.dtype 

69 post: tuple[Callable[[Any], Any], ...] = () 

70 

71 

72def _replaced(values: tuple, index: int, value: Any) -> tuple: 

73 """``values`` with the entry at ``index`` replaced by ``value``.""" 

74 return values[:index] + (value,) + values[index + 1:] 

75 

76 

77def _strides(sizes: Sequence[int]) -> tuple[int, ...]: 

78 """How many elements one step along each of ``sizes`` skips, in row-major order.""" 

79 strides, stride = [], 1 

80 for size in reversed(sizes): 

81 strides.append(stride) 

82 stride *= size 

83 return tuple(reversed(strides)) 

84 

85 

86def _check_count(count: int) -> None: 

87 """Refuse to go on describing a region cut into more than :data:`MAX_BLOCKS` blocks.""" 

88 if count > MAX_BLOCKS: 

89 raise UnsupportedRegionOp(f"the result is cut into more than {MAX_BLOCKS} blocks") 

90 

91 

92def _restrict(block: RegionBlock, axis: int, begin: int, length: int) -> RegionBlock: 

93 """The part of ``block`` from ``begin`` to ``begin + length`` along ``axis``, in the same coordinates.""" 

94 delta = begin - block.offsets[axis] 

95 return replace( 

96 block, 

97 offsets=_replaced(block.offsets, axis, begin), 

98 lengths=_replaced(block.lengths, axis, length), 

99 base=tuple(origin + step * delta for origin, step in zip(block.base, block.steps[axis])), 

100 ) 

101 

102 

103def _slice_axis(blocks: Sequence[RegionBlock], axis: int, start: int, stop: int, step: int) -> list[RegionBlock]: 

104 """ 

105 The blocks of ``tensor[..., start:stop:step, ...]`` along ``axis``, in the coordinates of the slice. 

106 

107 Args: 

108 blocks (Sequence[RegionBlock]): Blocks of the tensor being sliced. 

109 axis (int): Dimension to slice. 

110 start (int): First index kept, already clamped into the dimension. 

111 stop (int): Index the slice stops before, already clamped into the dimension. 

112 step (int): Positive distance between the indices kept. 

113 

114 Returns: 

115 list[RegionBlock]: The part of every block the slice keeps. 

116 """ 

117 sliced = [] 

118 for block in blocks: 

119 low = max(block.offsets[axis], start) 

120 high = min(block.offsets[axis] + block.lengths[axis], stop) 

121 if low >= high: 

122 continue 

123 first = -(-(low - start) // step) 

124 last = (high - 1 - start) // step 

125 if first > last: 

126 continue 

127 delta = start + step * first - block.offsets[axis] 

128 column = block.steps[axis] 

129 sliced.append(replace( 

130 block, 

131 offsets=_replaced(block.offsets, axis, first), 

132 lengths=_replaced(block.lengths, axis, last - first + 1), 

133 base=tuple(origin + c * delta for origin, c in zip(block.base, column)), 

134 steps=_replaced(block.steps, axis, tuple(c * step for c in column)), 

135 )) 

136 return sliced 

137 

138 

139def _drop_axis(block: RegionBlock, axis: int) -> RegionBlock: 

140 """``block`` without the dimension ``axis``, along which it is one element long.""" 

141 return replace( 

142 block, 

143 offsets=block.offsets[:axis] + block.offsets[axis + 1:], 

144 lengths=block.lengths[:axis] + block.lengths[axis + 1:], 

145 steps=block.steps[:axis] + block.steps[axis + 1:], 

146 ) 

147 

148 

149def _insert_axis(block: RegionBlock, axis: int) -> RegionBlock: 

150 """``block`` with a new dimension of size one at ``axis``.""" 

151 return replace( 

152 block, 

153 offsets=block.offsets[:axis] + (0,) + block.offsets[axis:], 

154 lengths=block.lengths[:axis] + (1,) + block.lengths[axis:], 

155 steps=block.steps[:axis] + ((0,) * len(block.base),) + block.steps[axis:], 

156 ) 

157 

158 

159def _shift(block: RegionBlock, axis: int, distance: int) -> RegionBlock: 

160 """``block`` moved ``distance`` elements along ``axis``.""" 

161 return replace(block, offsets=_replaced(block.offsets, axis, block.offsets[axis] + distance)) 

162 

163 

164def _permute_block(block: RegionBlock, order: Sequence[int]) -> RegionBlock: 

165 """``block`` with its dimensions in ``order``.""" 

166 return replace( 

167 block, 

168 offsets=tuple(block.offsets[axis] for axis in order), 

169 lengths=tuple(block.lengths[axis] for axis in order), 

170 steps=tuple(block.steps[axis] for axis in order), 

171 ) 

172 

173 

174def _merge_axes(blocks: Sequence[RegionBlock], start: int, sizes: Sequence[int]) -> list[RegionBlock]: 

175 """ 

176 The blocks of a tensor whose dimensions ``start`` to ``start + len(sizes)`` are flattened into one. 

177 

178 A block stays whole when it covers consecutive flat indices and walks the checkpoint the same way 

179 along them. That holds when every dimension after its first varying one is covered whole and moves 

180 through the checkpoint as a multiple of the last. Any other block is cut into pieces one element 

181 long along its first varying dimension, and those pieces are merged in turn. 

182 

183 Args: 

184 blocks (Sequence[RegionBlock]): Blocks of the tensor. 

185 start (int): First dimension to flatten. 

186 sizes (Sequence[int]): Sizes of the dimensions to flatten, all larger than one. 

187 

188 Returns: 

189 list[RegionBlock]: Blocks of the flattened tensor. 

190 """ 

191 stop = start + len(sizes) 

192 strides = _strides(sizes) 

193 merged, pending = [], list(blocks) 

194 while pending: 

195 block = pending.pop() 

196 varying = [axis for axis in range(len(sizes)) if block.lengths[start + axis] > 1] 

197 column = block.steps[stop - 1] if varying else (0,) * len(block.base) 

198 length = 1 

199 if varying: 

200 first = varying[0] 

201 whole = all( 

202 block.offsets[start + axis] == 0 and block.lengths[start + axis] == sizes[axis] 

203 for axis in range(first + 1, len(sizes)) 

204 ) 

205 in_step = all( 

206 block.steps[start + axis] == tuple(c * strides[axis] for c in column) 

207 for axis in range(first, len(sizes)) 

208 ) 

209 if not (whole and in_step): 

210 origin = block.offsets[start + first] 

211 pending.extend( 

212 _restrict(block, start + first, origin + index, 1) 

213 for index in range(block.lengths[start + first]) 

214 ) 

215 _check_count(len(merged) + len(pending)) 

216 continue 

217 length = block.lengths[start + first] * strides[first] 

218 offset = sum(index * stride for index, stride in zip(block.offsets[start:stop], strides)) 

219 merged.append(replace( 

220 block, 

221 offsets=block.offsets[:start] + (offset,) + block.offsets[stop:], 

222 lengths=block.lengths[:start] + (length,) + block.lengths[stop:], 

223 steps=block.steps[:start] + (column,) + block.steps[stop:], 

224 )) 

225 return merged 

226 

227 

228def _decompose(start: int, length: int, sizes: Sequence[int]) -> list[tuple[tuple[int, ...], tuple[int, ...]]]: 

229 """ 

230 Cut a run of flat indices over dimensions of ``sizes`` into boxes. 

231 

232 Args: 

233 start (int): First flat index of the run. 

234 length (int): Number of flat indices in the run. 

235 sizes (Sequence[int]): Sizes of the dimensions the flat index is spread over. 

236 

237 Returns: 

238 list[tuple[tuple[int, ...], tuple[int, ...]]]: ``(offsets, lengths)`` of every box: a partial 

239 first row, the whole rows, then a partial last row. 

240 """ 

241 if len(sizes) == 1: 

242 return [((start,), (length,))] 

243 inner = tuple(sizes[1:]) 

244 stride = math.prod(inner) 

245 boxes = [] 

246 row, column = divmod(start, stride) 

247 if column: 

248 head = min(stride - column, length) 

249 boxes.extend(((row,) + offsets, (1,) + lengths) for offsets, lengths in _decompose(column, head, inner)) 

250 row, length = row + 1, length - head 

251 rows, tail = divmod(length, stride) 

252 if rows: 

253 boxes.append(((row,) + (0,) * len(inner), (rows,) + inner)) 

254 if tail: 

255 boxes.extend(((row + rows,) + offsets, (1,) + lengths) for offsets, lengths in _decompose(0, tail, inner)) 

256 return boxes 

257 

258 

259def _split_axis(blocks: Sequence[RegionBlock], axis: int, sizes: Sequence[int]) -> list[RegionBlock]: 

260 """ 

261 The blocks of a tensor whose dimension ``axis`` is unflattened into dimensions of ``sizes``. 

262 

263 Args: 

264 blocks (Sequence[RegionBlock]): Blocks of the tensor. 

265 axis (int): Dimension to unflatten. 

266 sizes (Sequence[int]): Sizes it is unflattened into, all larger than one. 

267 

268 Returns: 

269 list[RegionBlock]: Blocks of the unflattened tensor. 

270 """ 

271 strides = _strides(sizes) 

272 pieces = [] 

273 for block in blocks: 

274 column = block.steps[axis] 

275 scaled = tuple(tuple(c * stride for c in column) for stride in strides) 

276 for offsets, lengths in _decompose(block.offsets[axis], block.lengths[axis], sizes): 

277 delta = sum(index * stride for index, stride in zip(offsets, strides)) - block.offsets[axis] 

278 pieces.append(replace( 

279 block, 

280 offsets=block.offsets[:axis] + offsets + block.offsets[axis + 1:], 

281 lengths=block.lengths[:axis] + lengths + block.lengths[axis + 1:], 

282 base=tuple(origin + c * delta for origin, c in zip(block.base, column)), 

283 steps=block.steps[:axis] + scaled + block.steps[axis + 1:], 

284 )) 

285 return pieces 

286 

287 

288def _runs(old: Sequence[int], new: Sequence[int]) -> list[tuple[int, int, int, int]]: 

289 """ 

290 Match the dimensions of two shapes up into runs holding as many elements on either side. 

291 

292 Args: 

293 old (Sequence[int]): Sizes of one shape, none of them one. 

294 new (Sequence[int]): Sizes of the other, none of them one, holding as many elements in all. 

295 

296 Returns: 

297 list[tuple[int, int, int, int]]: ``(old start, old stop, new start, new stop)`` of every run. 

298 """ 

299 runs, i, j = [], 0, 0 

300 while i < len(old): 

301 i0, j0 = i, j 

302 old_size, new_size = old[i], new[j] 

303 i, j = i + 1, j + 1 

304 while old_size != new_size: 

305 if old_size < new_size: 

306 old_size, i = old_size * old[i], i + 1 

307 else: 

308 new_size, j = new_size * new[j], j + 1 

309 runs.append((i0, i, j0, j)) 

310 return runs 

311 

312 

313def _reshape( 

314 blocks: Sequence[RegionBlock], old_shape: tuple[int, ...], new_shape: tuple[int, ...] 

315) -> list[RegionBlock]: 

316 """ 

317 The blocks of a tensor of ``old_shape`` reshaped to ``new_shape``, which holds as many elements. 

318 

319 Dimensions of size one are dropped first and put back last. The rest are matched up into runs 

320 holding as many elements on either side, and each run is flattened and then unflattened. 

321 

322 Args: 

323 blocks (Sequence[RegionBlock]): Blocks of the tensor. 

324 old_shape (tuple[int, ...]): Its shape. 

325 new_shape (tuple[int, ...]): The shape it is reshaped to. 

326 

327 Returns: 

328 list[RegionBlock]: Blocks of the reshaped tensor. 

329 """ 

330 if math.prod(old_shape) == 0: 

331 return [] 

332 blocks = list(blocks) 

333 for axis in reversed([axis for axis, size in enumerate(old_shape) if size == 1]): 

334 blocks = [_drop_axis(block, axis) for block in blocks] 

335 old = [size for size in old_shape if size != 1] 

336 new = [size for size in new_shape if size != 1] 

337 for i0, i1, j0, j1 in reversed(_runs(old, new)): 

338 if i1 - i0 > 1: 

339 blocks = _merge_axes(blocks, i0, old[i0:i1]) 

340 if j1 - j0 > 1: 

341 blocks = _split_axis(blocks, i0, new[j0:j1]) 

342 _check_count(len(blocks)) 

343 for axis, size in enumerate(new_shape): 

344 if size == 1: 

345 blocks = [_insert_axis(block, axis) for block in blocks] 

346 return blocks 

347 

348 

349def _shape_argument(shape: tuple) -> tuple: 

350 """A shape passed either as one sequence or as separate sizes.""" 

351 if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): 

352 return tuple(shape[0]) 

353 return tuple(shape) 

354 

355 

356def _expand_index(index: Any, ndim: int) -> list: 

357 """ 

358 The items of an index into a region, with Ellipsis spelled out as whole slices. 

359 

360 Args: 

361 index (Any): What ``__getitem__`` was given. 

362 ndim (int): Number of dimensions of the region. 

363 

364 Returns: 

365 list: Integers, slices and None, one integer or slice per dimension indexed. 

366 

367 Raises: 

368 UnsupportedRegionOp: If the index holds anything but integers, slices, None and Ellipsis. 

369 """ 

370 items = index if isinstance(index, tuple) else (index,) 

371 supported = (int, slice, type(None), type(Ellipsis)) 

372 if any(isinstance(item, bool) or not isinstance(item, supported) for item in items): 

373 raise UnsupportedRegionOp(f"indexing a region with {index!r} is not supported") 

374 consumed = sum(1 for item in items if isinstance(item, (int, slice))) 

375 expanded = [] 

376 for item in items: 

377 expanded.extend([slice(None)] * (ndim - consumed) if item is Ellipsis else [item]) 

378 return expanded 

379 

380 

381def _check_scalar(other: Any) -> None: 

382 """Refuse an operand that is not a Python number.""" 

383 if not isinstance(other, (int, float)): 

384 raise UnsupportedRegionOp(f"only Python numbers can be combined with a region, got {type(other).__name__}") 

385 

386 

387def _to_dtype(tensor: Any, dtype: torch.dtype) -> Any: 

388 """``tensor`` converted to ``dtype``.""" 

389 return tensor.to(dtype) 

390 

391 

392def _forward(function: Callable[[Any, Any], Any], other: Any, tensor: Any) -> Any: 

393 """``function(tensor, other)``.""" 

394 return function(tensor, other) 

395 

396 

397def _reflected(function: Callable[[Any, Any], Any], other: Any, tensor: Any) -> Any: 

398 """``function(other, tensor)``.""" 

399 return function(other, tensor) 

400 

401 

402def _divide(other: Any, rounding_mode: str, tensor: Any) -> Any: 

403 """``tensor / other``, rounded as ``rounding_mode`` says.""" 

404 return torch.div(tensor, other, rounding_mode=rounding_mode) 

405 

406 

407def _unsupported(name: str) -> Callable[..., Any]: 

408 """A method that refuses the operation ``name``.""" 

409 

410 def refuse(self: "RegionTensor", *args: Any, **kwargs: Any) -> Any: 

411 """Refuse an operation a region cannot describe.""" 

412 del args, kwargs 

413 raise UnsupportedRegionOp(f"{name} is not supported on {self!r}") 

414 

415 refuse.__name__ = name 

416 return refuse 

417 

418 

419class RegionTensor: 

420 """ 

421 A tensor that knows, for each of its elements, which checkpoint tensor and index it is read from. 

422 

423 Supports what weight conversions do to lay tensors out differently: concatenating, stacking, 

424 splitting, slicing, reordering and reshaping, plus elementwise arithmetic with Python scalars and 

425 dtype conversions, as methods and through the matching ``torch`` functions. Shapes and dtypes follow 

426 torch exactly, since every operation is first run on a meta tensor of the same shape and dtype. 

427 

428 Example:: 

429 

430 q = RegionTensor.leaf("q_proj.weight", (8, 16), torch.bfloat16) 

431 k = RegionTensor.leaf("k_proj.weight", (4, 16), torch.bfloat16) 

432 fused = torch.cat([q, k]) 

433 fused.remap_blocks() # two blocks: rows 0-7 read from q_proj, rows 8-11 from k_proj 

434 """ 

435 

436 __slots__ = ("_shape", "_dtype", "_blocks") 

437 

438 def __init__(self, shape: Sequence[int], dtype: torch.dtype, blocks: Sequence[RegionBlock] = ()) -> None: 

439 """ 

440 Args: 

441 shape (Sequence[int]): Shape of the tensor. 

442 dtype (torch.dtype): Dtype of the tensor. 

443 blocks (Sequence[RegionBlock]): Blocks tiling the tensor. Default (), for an empty tensor. 

444 """ 

445 self._shape = tuple(int(size) for size in shape) 

446 self._dtype = dtype 

447 self._blocks = tuple(blocks) 

448 _check_count(len(self._blocks)) 

449 

450 @classmethod 

451 def leaf(cls, source: str, shape: Sequence[int], dtype: torch.dtype) -> "RegionTensor": 

452 """ 

453 The whole of one checkpoint tensor. 

454 

455 Args: 

456 source (str): Name of the checkpoint tensor. 

457 shape (Sequence[int]): Its shape. 

458 dtype (torch.dtype): Its dtype. 

459 

460 Returns: 

461 RegionTensor: A tensor of one block reading the checkpoint tensor in place. 

462 """ 

463 shape = tuple(int(size) for size in shape) 

464 ndim = len(shape) 

465 blocks = () 

466 if math.prod(shape): 

467 identity = tuple(tuple(int(row == column) for row in range(ndim)) for column in range(ndim)) 

468 blocks = (RegionBlock((0,) * ndim, shape, source, (0,) * ndim, identity, dtype),) 

469 return cls(shape, dtype, blocks) 

470 

471 @classmethod 

472 def __torch_function__(cls, func: Any, types: Any, args: tuple = (), kwargs: Any = None) -> Any: 

473 """Route the ``torch`` functions a region supports to its methods, and refuse the rest.""" 

474 del types 

475 kwargs = dict(kwargs or {}) 

476 name = _FUNCTION_NAMES.get(func) 

477 if name is None or kwargs.get("out") is not None: 

478 raise UnsupportedRegionOp(f"torch.{getattr(func, '__name__', func)} is not supported on regions") 

479 kwargs.pop("out", None) 

480 if name in _SEQUENCE_FUNCTIONS: 

481 return _SEQUENCE_FUNCTIONS[name](*args, **kwargs) 

482 args = list(args) 

483 tensor = kwargs.pop("input") if "input" in kwargs else args.pop(0) 

484 if isinstance(tensor, RegionTensor): 

485 return getattr(tensor, name)(*args, **kwargs) 

486 other = kwargs.pop("other") if "other" in kwargs else (args.pop(0) if args else None) 

487 if name in _REFLECTED and isinstance(other, RegionTensor) and not args and not kwargs: 

488 return getattr(other, _REFLECTED[name])(tensor) 

489 raise UnsupportedRegionOp(f"torch.{name} is not supported with {type(tensor).__name__} as its tensor") 

490 

491 @property 

492 def shape(self) -> torch.Size: 

493 """Shape of the tensor.""" 

494 return torch.Size(self._shape) 

495 

496 @property 

497 def dtype(self) -> torch.dtype: 

498 """Dtype of the tensor.""" 

499 return self._dtype 

500 

501 @property 

502 def ndim(self) -> int: 

503 """Number of dimensions.""" 

504 return len(self._shape) 

505 

506 @property 

507 def device(self) -> torch.device: 

508 """The meta device: a region holds no values anywhere.""" 

509 return torch.device("meta") 

510 

511 @property 

512 def requires_grad(self) -> bool: 

513 """Always False.""" 

514 return False 

515 

516 @property 

517 def blocks(self) -> tuple[RegionBlock, ...]: 

518 """Blocks tiling the tensor.""" 

519 return self._blocks 

520 

521 def sources(self) -> frozenset[str]: 

522 """Names of the checkpoint tensors the elements are read from.""" 

523 return frozenset(block.source for block in self._blocks) 

524 

525 def remap_blocks(self) -> list[RemapBlock]: 

526 """ 

527 The blocks as :class:`RemapLoadPlanner` loads them. 

528 

529 A block kept in another dtype than the tensor, which a concatenation of tensors of different 

530 dtypes leaves behind, is converted to the dtype of the tensor last, as the conversion would. 

531 

532 Returns: 

533 list[RemapBlock]: One entry per block. 

534 """ 

535 remapped = [] 

536 for block in self._blocks: 

537 post = block.post 

538 if block.dtype != self._dtype: 

539 post += (partial(_to_dtype, dtype=self._dtype),) 

540 coeff = tuple(tuple(column[dim] for column in block.steps) for dim in range(len(block.base))) 

541 remapped.append(RemapBlock(block.offsets, block.lengths, block.source, block.base, coeff, post)) 

542 return remapped 

543 

544 def size(self, dim: Any = None) -> Any: 

545 """The shape, or the size of dimension ``dim``.""" 

546 return self.shape if dim is None else self._meta().size(dim) 

547 

548 def dim(self) -> int: 

549 """Number of dimensions.""" 

550 return len(self._shape) 

551 

552 def numel(self) -> int: 

553 """Number of elements.""" 

554 return math.prod(self._shape) 

555 

556 def nelement(self) -> int: 

557 """Number of elements.""" 

558 return self.numel() 

559 

560 def element_size(self) -> int: 

561 """Bytes one element takes.""" 

562 return self._meta().element_size() 

563 

564 def is_floating_point(self) -> bool: 

565 """Whether the dtype is a floating point one.""" 

566 return self._dtype.is_floating_point 

567 

568 def is_complex(self) -> bool: 

569 """Whether the dtype is a complex one.""" 

570 return self._dtype.is_complex 

571 

572 def __len__(self) -> int: 

573 """Size of the first dimension.""" 

574 if not self._shape: 

575 raise TypeError("len() of a 0-d tensor") 

576 return self._shape[0] 

577 

578 def __iter__(self) -> Iterator["RegionTensor"]: 

579 """The tensor one index of its first dimension at a time.""" 

580 if not self._shape: 

581 raise TypeError("iteration over a 0-d tensor") 

582 return iter(self.unbind(0)) 

583 

584 def __repr__(self) -> str: 

585 """Shape, dtype and blocks.""" 

586 return f"RegionTensor(shape={self._shape}, dtype={self._dtype}, blocks={len(self._blocks)})" 

587 

588 def reshape(self, *shape: Any) -> "RegionTensor": 

589 """The tensor reshaped, as ``torch.reshape`` does it.""" 

590 return self._reshaped(tuple(self._meta().reshape(_shape_argument(shape)).shape)) 

591 

592 def view(self, *shape: Any) -> "RegionTensor": 

593 """The tensor reshaped. Viewing it as another dtype is not supported.""" 

594 if len(shape) == 1 and isinstance(shape[0], torch.dtype): 

595 raise UnsupportedRegionOp("viewing a region as another dtype reinterprets its bytes") 

596 return self.reshape(*shape) 

597 

598 def reshape_as(self, other: Any) -> "RegionTensor": 

599 """The tensor reshaped to the shape of ``other``.""" 

600 return self.reshape(tuple(other.shape)) 

601 

602 def view_as(self, other: Any) -> "RegionTensor": 

603 """The tensor reshaped to the shape of ``other``.""" 

604 return self.reshape(tuple(other.shape)) 

605 

606 def flatten(self, start_dim: int = 0, end_dim: int = -1) -> "RegionTensor": 

607 """The tensor with dimensions ``start_dim`` to ``end_dim`` flattened into one.""" 

608 return self._reshaped(tuple(self._meta().flatten(start_dim, end_dim).shape)) 

609 

610 def unflatten(self, dim: int, sizes: Sequence[int]) -> "RegionTensor": 

611 """The tensor with dimension ``dim`` unflattened into ``sizes``.""" 

612 return self._reshaped(tuple(self._meta().unflatten(dim, sizes).shape)) 

613 

614 def transpose(self, dim0: int, dim1: int) -> "RegionTensor": 

615 """The tensor with dimensions ``dim0`` and ``dim1`` swapped.""" 

616 self._meta().transpose(dim0, dim1) 

617 if not self._shape: 

618 return self 

619 order = list(range(len(self._shape))) 

620 first, second = dim0 % len(order), dim1 % len(order) 

621 order[first], order[second] = order[second], order[first] 

622 return self._permuted(order) 

623 

624 def swapaxes(self, axis0: int, axis1: int) -> "RegionTensor": 

625 """The tensor with dimensions ``axis0`` and ``axis1`` swapped.""" 

626 return self.transpose(axis0, axis1) 

627 

628 def swapdims(self, dim0: int, dim1: int) -> "RegionTensor": 

629 """The tensor with dimensions ``dim0`` and ``dim1`` swapped.""" 

630 return self.transpose(dim0, dim1) 

631 

632 def t(self) -> "RegionTensor": 

633 """The tensor with its two dimensions swapped, or itself when it has fewer.""" 

634 self._meta().t() 

635 return self.transpose(0, 1) if len(self._shape) == 2 else self 

636 

637 def permute(self, *dims: Any) -> "RegionTensor": 

638 """The tensor with its dimensions reordered.""" 

639 dims = _shape_argument(dims) 

640 self._meta().permute(dims) 

641 return self._permuted([dim % len(self._shape) for dim in dims]) if self._shape else self 

642 

643 def squeeze(self, dim: Any = None) -> "RegionTensor": 

644 """The tensor without the dimensions of size one ``dim`` names, or without all of them.""" 

645 expected = tuple((self._meta().squeeze() if dim is None else self._meta().squeeze(dim)).shape) 

646 if not self._shape: 

647 return self 

648 if dim is None: 

649 axes = {axis for axis, size in enumerate(self._shape) if size == 1} 

650 else: 

651 dims = dim if isinstance(dim, (tuple, list)) else (dim,) 

652 axes = {d % len(self._shape) for d in dims if self._shape[d % len(self._shape)] == 1} 

653 blocks = list(self._blocks) 

654 for axis in sorted(axes, reverse=True): 

655 blocks = [_drop_axis(block, axis) for block in blocks] 

656 return self._checked(expected, blocks) 

657 

658 def unsqueeze(self, dim: int) -> "RegionTensor": 

659 """The tensor with a new dimension of size one at ``dim``.""" 

660 expected = tuple(self._meta().unsqueeze(dim).shape) 

661 axis = dim % (len(self._shape) + 1) 

662 return self._checked(expected, [_insert_axis(block, axis) for block in self._blocks]) 

663 

664 def split(self, split_size_or_sections: Any, dim: int = 0) -> tuple["RegionTensor", ...]: 

665 """The tensor cut along ``dim``, as ``torch.split`` cuts it.""" 

666 return self._pieces(self._meta().split(split_size_or_sections, dim), dim) 

667 

668 def chunk(self, chunks: int, dim: int = 0) -> tuple["RegionTensor", ...]: 

669 """The tensor cut along ``dim``, as ``torch.chunk`` cuts it.""" 

670 return self._pieces(self._meta().chunk(chunks, dim), dim) 

671 

672 def unbind(self, dim: int = 0) -> tuple["RegionTensor", ...]: 

673 """Every index of dimension ``dim``, without that dimension.""" 

674 size = self._meta().size(dim) 

675 return tuple(self.select(dim, index) for index in range(size)) 

676 

677 def narrow(self, dim: int, start: int, length: int) -> "RegionTensor": 

678 """``length`` indices of dimension ``dim`` from ``start``.""" 

679 expected = tuple(self._meta().narrow(dim, start, length).shape) 

680 axis = dim % len(self._shape) 

681 start = start + self._shape[axis] if start < 0 else start 

682 return self._checked(expected, _slice_axis(self._blocks, axis, start, start + length, 1)) 

683 

684 def select(self, dim: int, index: int) -> "RegionTensor": 

685 """Index ``index`` of dimension ``dim``, without that dimension.""" 

686 expected = tuple(self._meta().select(dim, index).shape) 

687 axis = dim % len(self._shape) 

688 index = index + self._shape[axis] if index < 0 else index 

689 blocks = [_drop_axis(block, axis) for block in _slice_axis(self._blocks, axis, index, index + 1, 1)] 

690 return self._checked(expected, blocks) 

691 

692 def __getitem__(self, index: Any) -> "RegionTensor": 

693 """The part of the tensor integers, slices with a positive step, None and Ellipsis pick out.""" 

694 expanded = _expand_index(index, len(self._shape)) 

695 expected = tuple(self._meta()[index].shape) 

696 shape, blocks, axis = list(self._shape), list(self._blocks), 0