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.components.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 _non_singleton_shape(shape: tuple[int, ...]) -> list[int]: 

314 """Return dimensions retained by the reshape run matching algorithm.""" 

315 return [size for size in shape if size != 1] 

316 

317 

318def _reshape( 

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

320) -> list[RegionBlock]: 

321 """ 

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

323 

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

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

326 

327 Args: 

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

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

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

331 

332 Returns: 

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

334 """ 

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

336 return [] 

337 blocks = list(blocks) 

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

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

340 old = _non_singleton_shape(old_shape) 

341 new = _non_singleton_shape(new_shape) 

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

343 if i1 - i0 > 1: 

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

345 if j1 - j0 > 1: 

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

347 _check_count(len(blocks)) 

348 for axis, size in enumerate(new_shape): 

349 if size == 1: 

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

351 return blocks 

352 

353 

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

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

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

357 return tuple(shape[0]) 

358 return tuple(shape) 

359 

360 

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

362 """ 

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

364 

365 Args: 

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

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

368 

369 Returns: 

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

371 

372 Raises: 

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

374 """ 

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

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

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

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

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

380 expanded = [] 

381 for item in items: 

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

383 return expanded 

384 

385 

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

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

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

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

390 

391 

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

393 """``tensor`` converted to ``dtype``.""" 

394 return tensor.to(dtype) 

395 

396 

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

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

399 return function(tensor, other) 

400 

401 

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

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

404 return function(other, tensor) 

405 

406 

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

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

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

410 

411 

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

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

414 

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

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

417 del args, kwargs 

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

419 

420 refuse.__name__ = name 

421 return refuse 

422 

423 

424class RegionTensor: 

425 """ 

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

427 

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

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

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

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

432 

433 Example:: 

434 

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

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

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

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

439 """ 

440 

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

442 

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

444 """ 

445 Args: 

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

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

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

449 """ 

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

451 self._dtype = dtype 

452 self._blocks = tuple(blocks) 

453 _check_count(len(self._blocks)) 

454 

455 @classmethod 

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

457 """ 

458 The whole of one checkpoint tensor. 

459 

460 Args: 

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

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

463 dtype (torch.dtype): Its dtype. 

464 

465 Returns: 

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

467 """ 

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

469 ndim = len(shape) 

470 blocks = () 

471 if math.prod(shape): 

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

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

474 return cls(shape, dtype, blocks) 

475 

476 @classmethod 

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

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

479 del types 

480 kwargs = dict(kwargs or {}) 

481 name = _FUNCTION_NAMES.get(func) 

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

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

484 kwargs.pop("out", None) 

485 if name in _SEQUENCE_FUNCTIONS: 

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

487 args = list(args) 

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

489 if isinstance(tensor, RegionTensor): 

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

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

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

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

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

495 

496 @property 

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

498 """Shape of the tensor.""" 

499 return torch.Size(self._shape) 

500 

501 @property 

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

503 """Dtype of the tensor.""" 

504 return self._dtype 

505 

506 @property 

507 def ndim(self) -> int: 

508 """Number of dimensions.""" 

509 return len(self._shape) 

510 

511 @property 

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

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

514 return torch.device("meta") 

515 

516 @property 

517 def requires_grad(self) -> bool: 

518 """Always False.""" 

519 return False 

520 

521 @property 

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

523 """Blocks tiling the tensor.""" 

524 return self._blocks 

525 

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

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

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

529 

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

531 """ 

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

533 

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

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

536 

537 Returns: 

538 list[RemapBlock]: One entry per block. 

539 """ 

540 remapped = [] 

541 for block in self._blocks: 

542 post = block.post 

543 if block.dtype != self._dtype: 

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

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

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

547 return remapped 

548 

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

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

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

552 

553 def dim(self) -> int: 

554 """Number of dimensions.""" 

555 return len(self._shape) 

556 

557 def numel(self) -> int: 

558 """Number of elements.""" 

559 return math.prod(self._shape) 

560 

561 def nelement(self) -> int: 

562 """Number of elements.""" 

563 return self.numel() 

564 

565 def element_size(self) -> int: 

566 """Bytes one element takes.""" 

567 return self._meta().element_size() 

568 

569 def is_floating_point(self) -> bool: 

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

571 return self._dtype.is_floating_point 

572 

573 def is_complex(self) -> bool: 

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

575 return self._dtype.is_complex 

576 

577 def __len__(self) -> int: 

578 """Size of the first dimension.""" 

579 if not self._shape: 

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

581 return self._shape[0] 

582 

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

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

585 if not self._shape: 

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

587 return iter(self.unbind(0)) 

588 

589 def __repr__(self) -> str: 

590 """Shape, dtype and blocks.""" 

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

592 

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

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

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

596 

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

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

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

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

601 return self.reshape(*shape) 

602 

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

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

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

606 

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

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

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

610 

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

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

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

614 

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

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

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

618 

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

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

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

622 if not self._shape: 

623 return self 

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

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

626 order[first], order[second] = order[second], order[first] 

627 return self._permuted(order) 

628 

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

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

631 return self.transpose(axis0, axis1) 

632 

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

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

635 return self.transpose(dim0, dim1) 

636 

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

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

639 self._meta().t() 

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

641 

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

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

644 dims = _shape_argument(dims) 

645 self._meta().permute(dims) 

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

647 

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

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

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

651 if not self._shape: 

652 return self 

653 if dim is None: 

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

655 else: 

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

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

658 blocks = list(self._blocks) 

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

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

661 return self._checked(expected, blocks) 

662 

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

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

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

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

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

668 

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

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

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

672 

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

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

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

676 

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

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

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

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

681 

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

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

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

685 axis = dim % len(self._shape) 

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

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

688 

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

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

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

692 axis = dim % len(self._shape) 

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

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

695 return self._checked(expected, blocks) 

696 

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

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

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