Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/shard/ops/parallel

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

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

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"""

16Distributed implementation for __getitem__ operator (PyTorch only).

18Supports basic indexing (int, slice, None, Ellipsis) which produces views,

19and advanced indexing (list, LongTensor) which produces copies.

20BoolTensor masks are not supported because the output shape is data-dependent.

21"""

22from typing import Callable, NamedTuple, Optional

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

25from hyper_parallel.core.dtensor.dtensor import DTensor

26from hyper_parallel.core.dtensor.layout import Layout

27from hyper_parallel.core.dtensor.placement_types import (

28 _StridedRaggedShard,

29 RaggedShard,

30 Replicate,

31 Shard,

32 StridedShard,

33)

34from .parallel_ops import DistributedOp

36Tensor = torch.Tensor

38_BASIC = "basic"

39_ADVANCED = "advanced"

40_BOOL_MASK = "bool_mask"

43class _StridedRaggedIndexInfo(NamedTuple):

44 """Owner metadata for dim-zero indexing across a strided shard chain."""

46 shard_axes: tuple[int, ...]

47 owner_coordinates: tuple[int, ...]

48 local_index: int

49 output_global_shape: tuple[int, ...]

52def _normalize___getitem___args(self_t, key):

53 """Normalize __getitem__ arguments to canonical positional form.

55 __getitem__ is always called as func(self, key) with empty kwargs.

57 Args:

58 self_t: The DTensor self argument.

59 key: The indexing key (int, slice, None, Ellipsis, tuple, list, Tensor).

61 Returns:

62 tuple: ((self_t, key), {})

63 """

64 return (self_t, key), {}

67def _is_long_tensor(obj) -> bool:

68 """Check if obj is a non-bool Tensor (LongTensor or similar)."""

69 return isinstance(obj, torch.Tensor) and not _is_bool_tensor(obj)

72def _is_bool_tensor(obj) -> bool:

73 """Check if obj is a BoolTensor."""

74 if not isinstance(obj, torch.Tensor):

75 return False

76 if not hasattr(obj, 'dtype'):

77 return False

78 return str(obj.dtype).rsplit('.', maxsplit=1)[-1] in ('bool', 'bool_')

81def _is_advanced_elem(k) -> bool:

82 """Check if key element triggers advanced indexing."""

83 if isinstance(k, list):

84 return True

85 if isinstance(k, torch.Tensor):

86 # 0-D long tensor is treated as basic (equivalent to int)

87 if k.ndim == 0 and not _is_bool_tensor(k):

88 return False

89 return True

90 return False

93def _build_key_element_descriptor(k, op_name="__getitem__"):

94 """Build a descriptor tuple for a single key element.

96 Args:

97 k: A single key element (int, slice, None, Ellipsis, list, Tensor).

98 op_name: Operator name for error messages.

100 Returns:

101 tuple: A descriptor tuple for this element.

102

103 Raises:

104 ValueError: If the key element type is unsupported.

105 """

106 if k is None:

107 return ("none",)

108 if k is Ellipsis:

109 return ("ellipsis",)

110 if isinstance(k, int):

111 return ("int", k)

112 if isinstance(k, slice):

113 return ("slice", k.start, k.stop, k.step)

114 if isinstance(k, list):

115 return ("idx_list_len", len(k))

116 if isinstance(k, torch.Tensor) and k.ndim == 0 and not _is_bool_tensor(k):

117 return ("int", int(k.item()))

118 if isinstance(k, torch.Tensor):

119 shape = tuple(k.shape)

120 if isinstance(k, DTensor):

121 alias = tuple(k.layout.alias_tensor_map)

122 return ("idx_tensor", shape, alias)

123 return ("idx_tensor", shape)

124 raise ValueError(

125 f"For {op_name}, unsupported index type: {type(k).__name__}."

126 )

127

128

129def _key_cache_descriptor(key, op_name="__getitem__"):

130 """Convert raw key to hashable (desc, kind) for cache_values.

131

132 The descriptor captures only what affects layout derivation:

133 int value, slice bounds, list length, and tensor shape.

134 Exact list/tensor contents are not preserved — only shape matters for layout.

135

136 Args:

137 key: Raw key from __getitem__ call (int, slice, None, Ellipsis, tuple, list, Tensor).

138 op_name: Operator name for error messages.

139

140 Returns:

141 (desc, kind): desc is a tuple of hashable descriptor tuples,

142 kind is "basic", "advanced", or "bool_mask".

143

144 Raises:

145 ValueError: If key contains unsupported types.

146 """

147 if not isinstance(key, tuple):

148 key = (key,)

149

150 # First pass: detect BoolTensor anywhere in the key

151 for k in key:

152 if _is_bool_tensor(k):

153 mask_shape = tuple(k.shape) if hasattr(k, 'shape') else ()

154 return (("bool_mask", mask_shape),), _BOOL_MASK

155

156 # Check for advanced indexing elements (list or non-0D tensor)

157 has_advanced = any(_is_advanced_elem(k) for k in key)

158 kind = _ADVANCED if has_advanced else _BASIC

159

160 desc = [_build_key_element_descriptor(k, op_name) for k in key]

161 return tuple(desc), kind

162

163

164def _desc_action_to_expanded(action_type, d, input_dim):

165 """Convert a single descriptor element to expanded action.

166

167 Args:

168 action_type: Type string from descriptor (none, int, slice, etc.).

169 d: Descriptor tuple for this element.

170 input_dim: Current input dimension index.

171

172 Returns:

173 tuple: (expanded_entries, input_dim_delta)

174 expanded_entries: list of action tuples to append.

175 input_dim_delta: how much to advance input_dim.

176 """

177 if action_type == "none":

178 return [("newaxis",)], 0

179 if action_type == "int":

180 return [("int", d[1], input_dim)], 1

181 if action_type == "slice":

182 return [("slice", d[1], d[2], d[3], input_dim)], 1

183 if action_type == "idx_list_len":

184 return [("idx_list", (0,) * d[1], input_dim)], 1

185 if action_type == "idx_tensor":

186 shape = d[1]

187 alias = d[2] if len(d) > 2 else None

188 idx_info = {"shape": shape, "alias": alias}

189 return [("idx_tensor", idx_info, input_dim)], 1

190 raise ValueError(

191 f"Unsupported descriptor type: {action_type}."

192 )

193

194

195def _descriptor_to_expanded_actions(desc, ndim, op_name="__getitem__"):

196 """Reconstruct expanded_actions from key descriptor for layout derivation.

197

198 Args:

199 desc: Key descriptor tuple from _key_cache_descriptor.

200 ndim: Number of input tensor dimensions.

201 op_name: Operator name for error messages.

202

203 Returns:

204 list: Expanded actions, each a tuple (action_type, *args).

205

206 Raises:

207 ValueError: On invalid descriptor (multiple Ellipsis, too many indices).

208 """

209 if desc is None:

210 return []

211

212 # Pass through bool_mask descriptor

213 if len(desc) == 1 and desc[0][0] == "bool_mask":

214 return [("bool_mask", desc[0][1])]

215

216 non_none = [d for d in desc if d[0] != "none"]

217 n_ellipsis = sum(1 for d in non_none if d[0] == "ellipsis")

218

219 if n_ellipsis > 1:

220 raise ValueError(

221 f"For {op_name}, an index can only have a single Ellipsis ('...')."

222 )

223

224 n_specified_dims = len(non_none) - n_ellipsis

225 n_fill = ndim - n_specified_dims

226

227 if n_fill < 0:

228 raise ValueError(

229 f"For {op_name}, too many indices for tensor of dimension {ndim}, "

230 f"but got {len(non_none)} non-None indices."

231 )

232

233 expanded = []

234 input_dim = 0

235

236 for d in desc:

237 action_type = d[0]

238

239 if action_type == "ellipsis":

240 for _ in range(n_fill):

241 expanded.append(("slice", None, None, None, input_dim))

242 input_dim += 1

243 else:

244 entries, dim_delta = _desc_action_to_expanded(action_type, d, input_dim)

245 expanded.extend(entries)

246 input_dim += dim_delta

247

248 # Implicitly add full slices for unspecified trailing dimensions

249 while input_dim < ndim:

250 expanded.append(("slice", None, None, None, input_dim))

251 input_dim += 1

252

253 return expanded

254

255

256def _unwrap_key_for_local(key):

257 """Convert key for local execution: DTensors to local.

258

259 Args:

260 key: Raw key from __getitem__ call.

261

262 Returns:

263 Unwrapped key suitable for passing to local func.

264 """

265 if isinstance(key, DTensor):

266 return key.to_local()

267 if isinstance(key, tuple):

268 return tuple(_unwrap_key_for_local(k) for k in key)

269 return key

270

271

272def _copy_partial_state(src_layout, dst_layout):

273 """Copy partial state from src_layout to dst_layout using public API.

274

275 Args:

276 src_layout: Source Layout to copy partials from.

277 dst_layout: Destination Layout to copy partials to.

278 """

279 for dev_idx, op in enumerate(src_layout.partial):

280 if op is not None:

281 dst_layout.set_partial_by_dev_axis(dst_layout.alias_name[dev_idx], op)

282

283

284def _is_full_slice_action(action, global_shape) -> bool:

285 """Return whether an expanded action keeps one complete input dimension."""

286 if action[0] != "slice":

287 return False

288 input_dim = action[-1]

289 start, stop, step = action[1], action[2], action[3]

290 normalized_step = 1 if step is None else step

291 return (

292 normalized_step == 1

293 and (start is None or start == 0)

294 and (stop is None or stop >= global_shape[input_dim])

295 )

296

297

298def _is_dim0_int_full_slice(expanded_actions, global_shape, kind) -> bool:

299 """Return whether an index selects one row and keeps all trailing dims."""

300 return (

301 kind == _BASIC

302 and len(global_shape) >= 2

303 and bool(expanded_actions)

304 and expanded_actions[0][0] == "int"

305 and expanded_actions[0][-1] == 0

306 and all(

307 _is_full_slice_action(action, global_shape)

308 for action in expanded_actions[1:]

309 )

310 )

311

312

313def _dim0_shard_axes(placements) -> tuple[int, ...]:

314 """Return mesh axes that shard tensor dimension zero."""

315 return tuple(

316 mesh_dim

317 for mesh_dim, placement in enumerate(placements)

318 if isinstance(placement, Shard) and placement.is_shard(0)

319 )

320

321

322def _other_axes_are_replicated(placements, shard_axes) -> bool:

323 """Return whether every non-dim-zero-shard mesh axis is replicated."""

324 shard_axis_set = set(shard_axes)

325 return all(

326 mesh_dim in shard_axis_set or placement.is_replicate()

327 for mesh_dim, placement in enumerate(placements)

328 )

329

330

331def _build_regular_ragged_result(self_layout, output_global_shape, global_dim0, index):

332 """Build an owner-only RaggedShard result for a regular dim-zero shard."""

333 placements = tuple(self_layout.placements)

334 shard_axes = _dim0_shard_axes(placements)

335 if len(shard_axes) != 1 or not _other_axes_are_replicated(placements, shard_axes):

336 return None

337

338 shard_axis = shard_axes[0]

339 mesh_size = self_layout.mesh_shape[shard_axis]

340 if global_dim0 % mesh_size != 0:

341 return None

342

343 rows_per_rank = global_dim0 // mesh_size

344 owner_rank = index // rows_per_rank

345 local_index = index % rows_per_rank

346 local_units = tuple(

347 1 if rank == owner_rank else 0 for rank in range(mesh_size)

348 )

349 output_layout = Layout.from_device_mesh(self_layout.mesh)

350 output_placements = [Replicate() for _ in placements]

351 output_placements[shard_axis] = RaggedShard(

352 tuple(range(len(output_global_shape))), local_units

353 )

354 output_layout.set_placements(tuple(output_placements))

355 output_layout.placement_to_tensor_map(len(output_global_shape))

356 if len(placements) == 1:

357 # Preserve the legacy metadata shape used by existing one-dimensional tests.

358 info = (owner_rank, local_index, output_global_shape)

359 else:

360 info = _StridedRaggedIndexInfo(

361 (shard_axis,), (owner_rank,), local_index, output_global_shape

362 )

363 return output_layout, info

364

365

366def _strided_owner_info(self_layout, placements, shard_axes, global_dim0, index):

367 """Return local row index and owner coordinate for a strided shard chain."""

368 total_shard_count = 1

369 for mesh_dim in shard_axes:

370 total_shard_count *= self_layout.mesh_shape[mesh_dim]

371 if global_dim0 % total_shard_count != 0:

372 raise ValueError(

373 "For __getitem__, StridedShard dim 0 size should be divisible "

374 "by the total shard count, but got "

375 f"dim0={global_dim0}, total_shard_count={total_shard_count}."

376 )

377

378 rows_per_shard = global_dim0 // total_shard_count

379 combined_shard_rank = index // rows_per_shard

380 local_index = index % rows_per_shard

381 shard_order = sorted(

382 shard_axes,

383 key=lambda mesh_dim: (

384 placements[mesh_dim].split_factor

385 if isinstance(placements[mesh_dim], StridedShard)

386 else 1

387 ),

388 )

389 owner_by_axis = {}

390 for mesh_dim in reversed(shard_order):

391 owner_by_axis[mesh_dim] = combined_shard_rank % self_layout.mesh_shape[mesh_dim]

392 combined_shard_rank //= self_layout.mesh_shape[mesh_dim]

393 owner_coordinate = tuple(owner_by_axis[mesh_dim] for mesh_dim in shard_axes)

394 return owner_coordinate, local_index

395

396

397def _build_strided_ragged_result(

398 self_layout,

399 placements,

400 shard_axes,

401 output_global_shape,

402 global_dim0,

403 index,

404):

405 """Build the multi-axis strided ragged result and owner metadata."""

406 owner_coordinate, local_index = _strided_owner_info(

407 self_layout, placements, shard_axes, global_dim0, index

408 )

409 output_dims = tuple(range(len(output_global_shape)))

410 output_placements = [Replicate() for _ in placements]

411 for mesh_dim, owner_axis_coordinate in zip(shard_axes, owner_coordinate):

412 local_units = tuple(

413 int(coordinate == owner_axis_coordinate)

414 for coordinate in range(self_layout.mesh_shape[mesh_dim])

415 )

416 input_placement = placements[mesh_dim]

417 split_factor = (

418 input_placement.split_factor

419 if isinstance(input_placement, StridedShard)

420 else 1

421 )

422 output_placements[mesh_dim] = _StridedRaggedShard(

423 output_dims,

424 local_units,

425 split_factor,

426 )

427

428 output_layout = Layout.from_device_mesh(self_layout.mesh)

429 output_layout.set_placements(tuple(output_placements))

430 output_layout.placement_to_tensor_map(len(output_global_shape))

431 info = _StridedRaggedIndexInfo(

432 shard_axes,

433 owner_coordinate,

434 local_index,

435 output_global_shape,

436 )

437 return output_layout, info

438

439

440class GetItemDistributedOp(DistributedOp):

441 """Distributed implementation for tensor.__getitem__.

442

443 Supports basic indexing (int, slice, None, Ellipsis) which produces views,

444 and advanced indexing (list, LongTensor) which produces copies.

445 BoolTensor masks are rejected because they produce data-dependent shapes.

446

447 Sharding constraints:

448 - Integer indexing on Shard(0) is supported for the owner-only RaggedShard

449 view case on a 1-D mesh; other indexed dimensions must be replicated.

450 - Any dimension indexed by non-full slice or advanced index must be replicated.

451 - Advanced index tensors must themselves be replicated.

452 - Input must not have Partial status.

453 - slice step != 1 is not supported.

454 """

455

456 def preprocess(self, args: tuple, kwargs: dict) -> tuple:

457 """Preprocess arguments for __getitem__.

458

459 Converts raw key to hashable descriptor for cache, and unwraps

460 DTensors to local tensors for execution.

461

462 Args:

463 args: (self_tensor, key)

464 kwargs: Empty dict for __getitem__.

465

466 Returns:

467 tuple: (local_args, local_kwargs, cache_values)

468 """

469 norm_args, _ = _normalize___getitem___args(*args, **kwargs)

470 self_t, key = norm_args

471

472 self_layout = self_t.layout

473 global_shape = tuple(self_t.shape)

474

475 # raw key -> hashable descriptor for cache

476 key_desc, kind = _key_cache_descriptor(key, op_name=self.op_name)

477

478 # Unwrap DTensor in key to local for local execution

479 local_key = _unwrap_key_for_local(key)

480 local_args = (self_t.to_local(), local_key)

481 local_kwargs = {}

482

483 cache_values = [self_layout, key_desc, global_shape, kind]

484 return local_args, local_kwargs, cache_values

485

486 @staticmethod

487 def _reject_bool_mask(expanded_actions, op_name="__getitem__"):

488 """Raise ValueError for bool_mask key kind."""

489 mask_shape = "unknown"

490 for action in expanded_actions:

491 if action[0] == "bool_mask":

492 mask_shape = action[1]

493 break

494 raise ValueError(

495 f"For {op_name}, boolean-mask indexing has data-dependent "

496 f"output shape and is not supported in DTensor. "

497 f"Got mask of shape {mask_shape}."

498 )

499

500 @staticmethod

501 def _validate_int_action(action, alias_map, global_shape, op_name="__getitem__"):

502 """Validate an int indexing action."""

503 input_dim = action[-1]

504 idx = action[1]

505 if idx < -global_shape[input_dim] or idx >= global_shape[input_dim]:

506 raise ValueError(

507 f"For {op_name}, index {idx} is out of range "

508 f"for dimension {input_dim} with size {global_shape[input_dim]}."

509 )

510 if alias_map[input_dim] != "None":

511 raise ValueError(

512 f"For {op_name}, indexing with int on non-replicate "

513 f"dim {input_dim} is not supported, "

514 f"but got sharding {alias_map[input_dim]} on dim {input_dim}."

515 )

516

517 @staticmethod

518 def _validate_slice_action(action, alias_map, global_shape, op_name="__getitem__"):

519 """Validate a slice indexing action."""

520 input_dim = action[-1]

521 step = action[3]

522 step = step if step is not None else 1

523

524 if step != 1:

525 raise ValueError(

526 f"For {op_name}, slice step should be 1 or None, "

527 f"but got {step}."

528 )

529

530 if not _is_full_slice_action(action, global_shape) and alias_map[input_dim] != "None":

531 raise ValueError(

532 f"For {op_name}, non-full slice on non-replicate "

533 f"dim {input_dim} is not supported, "

534 f"but got sharding {alias_map[input_dim]} on dim {input_dim}."

535 )

536

537 @staticmethod

538 def _validate_advanced_action(action, alias_map, op_name="__getitem__"):

539 """Validate an advanced indexing action (idx_list or idx_tensor)."""

540 input_dim = action[-1]

541 if alias_map[input_dim] != "None":

542 raise ValueError(

543 f"For {op_name}, advanced indexing on non-replicate "

544 f"dim {input_dim} is not supported, "

545 f"but got sharding {alias_map[input_dim]} on dim {input_dim}."

546 )

547 # For idx_tensor, the index tensor itself must be replicated

548 if action[0] == "idx_tensor":

549 idx_alias = action[1].get("alias")

550 if idx_alias and any(x != "None" for x in idx_alias):

551 raise ValueError(

552 f"For {op_name}, advanced index tensor must be "

553 f"replicated, but got layout with sharding {idx_alias}."

554 )

555

556 @staticmethod

557 def _validate_input_layouts(self_layout, expanded_actions, global_shape, kind,

558 op_name="__getitem__"):

559 """Validate sharding constraints for __getitem__.

560

561 Rules:

562 1. BoolTensor mask indexing is not supported.

563 2. Any dimension indexed by int, non-full slice, or advanced index

564 must be replicated.

565 3. slice step must be None or 1.

566 4. Advanced index tensors must be replicated.

567 5. int indices must be in range.

568

569 Args:

570 self_layout: Layout of self tensor.

571 expanded_actions: Expanded key actions from _descriptor_to_expanded_actions.

572 global_shape: Global shape of self tensor.

573 kind: "basic" or "advanced".

574 op_name: Operator name for error messages (default "__getitem__").

575

576 Raises:

577 ValueError: If any constraint is violated.

578 """

579 alias_map = self_layout.alias_tensor_map

580

581 if kind == _BOOL_MASK:

582 GetItemDistributedOp._reject_bool_mask(expanded_actions, op_name)

583

584 for action in expanded_actions:

585 action_type = action[0]

586

587 if action_type == "newaxis":

588 continue

589 if action_type == "int":

590 GetItemDistributedOp._validate_int_action(

591 action, alias_map, global_shape, op_name

592 )

593 elif action_type == "slice":

594 GetItemDistributedOp._validate_slice_action(

595 action, alias_map, global_shape, op_name

596 )

597 elif action_type in ("idx_list", "idx_tensor"):

598 GetItemDistributedOp._validate_advanced_action(

599 action, alias_map, op_name

600 )

601

602 @staticmethod

603 def _infer_shard_dim0_int(self_layout, expanded_actions, global_shape, kind):

604 """Return the owner-only RaggedShard layout and local index, if supported."""

605 placements = tuple(self_layout.placements)

606 if not _is_dim0_int_full_slice(expanded_actions, global_shape, kind):

607 return None

608

609 shard_axes = _dim0_shard_axes(placements)

610 if not shard_axes or not _other_axes_are_replicated(placements, shard_axes):

611 return None

612

613 index = expanded_actions[0][1]

614 global_dim0 = global_shape[0]

615 if index < -global_dim0 or index >= global_dim0:

616 return None

617

618 normalized_index = index if index >= 0 else index + global_dim0

619 output_global_shape = tuple(global_shape[1:])

620 if not any(isinstance(placements[mesh_dim], StridedShard) for mesh_dim in shard_axes):

621 return _build_regular_ragged_result(

622 self_layout, output_global_shape, global_dim0, normalized_index

623 )

624 return _build_strided_ragged_result(

625 self_layout,

626 placements,

627 shard_axes,

628 output_global_shape,

629 global_dim0,

630 normalized_index,

631 )

632

633 def infer_layout(self, cache_values: list) -> tuple: # pylint: disable=W0221

634 """Infer output layout for __getitem__.

635

636 Rules:

637 1. Input must not have Partial status.

638 2. BoolTensor mask indexing is not supported.

639 3. Integer indexing on Shard(0) may produce an owner-only RaggedShard

640 view on a 1-D mesh.

641 4. Other dimensions indexed by int, non-full slice, or advanced index

642 must be replicated.

643 5. slice step must be None or 1.

644 6. Advanced index tensors must be replicated.

645 7. BASIC: Output alias_tensor_map is derived by removing int-indexed

646 dims, inserting Replicate for newaxis, and preserving sharding

647 for full-slice dims.

648 8. ADVANCED: Advanced indices broadcast shape B is inserted at

649 position p (consecutive: p = first advanced dim; non-consecutive:

650 p = 0). B dims are all Replicate. Other dims preserve sharding.

651 9. Partial state is copied to the output layout for preserved dims.

652

653 Args:

654 cache_values: [self_layout, key_desc, global_shape, kind]

655

656 Returns:

657 tuple: Output layouts and optional owner-only RaggedShard metadata.

658

659 Raises:

660 ValueError: If any constraint is violated.

661 """

662 self_layout = cache_values[0]

663 key_desc = cache_values[1]

664 global_shape = cache_values[2]

665 kind = cache_values[3]

666

667 # key_desc -> expanded_actions

668 expanded_actions = _descriptor_to_expanded_actions(

669 key_desc, len(global_shape), op_name=self.op_name

670 )

671

672 if not self._allow_partial_inputs:

673 self._check_partial_inputs([self_layout])

674

675 ragged_result = self._infer_shard_dim0_int(

676 self_layout, expanded_actions, global_shape, kind

677 )

678 if ragged_result is not None:

679 output_layout, info = ragged_result

680 return ((output_layout,), info)

681

682 self._validate_input_layouts(self_layout, expanded_actions, global_shape, kind)

683

684 if kind == _BASIC:

685 out_layout = self._infer_basic_output_layout(self_layout, expanded_actions)

686 elif kind == _ADVANCED:

687 out_layout = self._infer_advanced_output_layout(

688 self_layout, expanded_actions, op_name=self.op_name

689 )

690 else:

691 # _BOOL_MASK is rejected by _validate_input_layouts above,

692 # so we should never reach this branch.

693 raise ValueError(

694 f"For __getitem__, unexpected kind: {kind}. "

695 f"Expected 'basic' or 'advanced'."

696 )

697

698 return ((out_layout,), None)

699

700 def get_expand_impl(

701 self,

702 func: Callable,

703 infer_result: tuple,

704 cache_values: list,

705 ) -> Optional[Callable]:

706 """Return a local view implementation for owner-only RaggedShard indexing."""

707 info = infer_result[1]

708 if info is None:

709 return None

710

711 input_layout = cache_values[0]

712 output_layout = infer_result[0][0]

713 if isinstance(info, _StridedRaggedIndexInfo):

714 shard_axes = info.shard_axes

715 owner_coordinates = info.owner_coordinates

716 local_index = info.local_index

717 output_global_shape = info.output_global_shape

718 else:

719 owner_rank, local_index, output_global_shape = info

720 shard_axes = (0,)

721 owner_coordinates = (owner_rank,)

722

723 def ragged_getitem_impl(local_input: Tensor, local_key: object) -> DTensor:

724 """Return the selected owner view or an empty non-owner view."""

725 del local_key

726 if not local_input.is_contiguous():

727 raise ValueError(

728 f"For {self.op_name}, Shard(0) integer indexing to RaggedShard "

729 "requires a contiguous local tensor."

730 )

731

732 is_owner = all(

733 input_layout.mesh.get_local_rank(mesh_dim) == owner_coordinate

734 for mesh_dim, owner_coordinate in zip(shard_axes, owner_coordinates)

735 )

736 if is_owner:

737 local_view = func(local_input, local_index).view(-1)

738 else:

739 local_view = local_input.view(-1)[:0]

740 return DTensor.from_local_with_layout(

741 local_view,

742 output_layout,

743 shape=output_global_shape,

744 )

745

746 return ragged_getitem_impl

747

748 @staticmethod

749 def wrap_output(py_output: object, output_layouts: object) -> object:

750 """Pass through the RaggedShard DTensor built by the local implementation."""

751 if isinstance(py_output, DTensor):

752 return py_output

753 return DistributedOp.wrap_output(py_output, output_layouts)

754

755 @staticmethod

756 def _infer_basic_output_layout(self_layout, expanded_actions):

757 """Derive output layout for basic indexing.

758

759 Walk through expanded_actions:

760 - int: remove the dimension from alias_tensor_map

761 - full slice (None, None, 1): keep the dimension with same sharding

762 - non-full slice: keep the dimension, must be replicated (validated earlier)

763 - newaxis: insert Replicate dimension

764

765 Args:

766 self_layout: Layout of self tensor.

767 expanded_actions: Expanded key actions.

768

769 Returns:

770 Layout: Output layout.

771 """

772 alias_map = self_layout.alias_tensor_map

773 mesh = self_layout.mesh

774

775 out_alias = []

776 for action in expanded_actions:

777 action_type = action[0]

778 if action_type == "newaxis":

779 out_alias.append("None")

780 elif action_type == "slice":

781 input_dim = action[4]

782 out_alias.append(alias_map[input_dim])

783 # "int" actions are skipped (dimension removed)

784

785 out_layout = Layout.from_device_mesh(mesh)

786 out_layout = out_layout(*out_alias)

787

788 # Copy partial state from input

789 _copy_partial_state(self_layout, out_layout)

790 out_layout.tensor_map_to_placement()

791 out_layout.update_compact_str()

792 return out_layout

793

794 @staticmethod

795 def _compute_advanced_index_info(expanded_actions, op_name):

796 """Analyze advanced indexing actions for output layout derivation.

797

798 Returns:

799 (advanced_positions, advanced_pos_set, are_consecutive, b_ndim)

800 """

801 advanced_actions = []

802 for i, action in enumerate(expanded_actions):

803 if action[0] in ("idx_list", "idx_tensor"):

804 advanced_actions.append((i, action))

805

806 index_shapes = []

807 for _, action in advanced_actions:

808 if action[0] == "idx_list":

809 index_shapes.append((len(action[1]),))

810 elif action[0] == "idx_tensor":

811 index_shapes.append(tuple(action[1]["shape"]))

812

813 bcast_shape = _broadcast_shapes(index_shapes, op_name=op_name)

814 b_ndim = len(bcast_shape) if bcast_shape else 0

815

816 advanced_positions = [pos for pos, _ in advanced_actions]

817 advanced_pos_set = set(advanced_positions)

818 are_consecutive = (

819 len(advanced_positions) > 0

820 and advanced_positions == list(range(advanced_positions[0],

821 advanced_positions[-1] + 1))

822 )

823 return advanced_positions, advanced_pos_set, are_consecutive, b_ndim

824

825 @staticmethod

826 def _infer_advanced_output_layout(self_layout, expanded_actions, op_name="__getitem__"):

827 """Derive output layout for advanced indexing.

828

829 Advanced indices are on specific input dims L. Other input dims K

830 (including newaxis) are preserved. The broadcast shape B of the

831 index tensors is inserted at position p:

832 - consecutive advanced indices: p = position of first advanced action

833 - non-consecutive: p = 0 (B precedes all other dims)

834

835 Args:

836 self_layout: Layout of self tensor.

837 expanded_actions: Expanded key actions.

838 op_name: Operator name for error messages.

839

840 Returns:

841 Layout: Output layout.

842 """

843 alias_map = self_layout.alias_tensor_map

844 mesh = self_layout.mesh

845

846 advanced_positions, advanced_pos_set, are_consecutive, b_ndim = \

847 GetItemDistributedOp._compute_advanced_index_info(expanded_actions, op_name)

848

849 def _append_non_advanced(action):

850 """Append output alias for a non-advanced action (newaxis, slice, int)."""

851 if action[0] == "newaxis":

852 out_alias.append("None")

853 elif action[0] == "slice":

854 out_alias.append(alias_map[action[4]])

855 # int: dimension removed, skip

856

857 out_alias = []

858

859 if are_consecutive and advanced_positions:

860 first_adv_pos = advanced_positions[0]

861 last_adv_pos = advanced_positions[-1]

862

863 for i in range(first_adv_pos):

864 _append_non_advanced(expanded_actions[i])

865 for _ in range(b_ndim):

866 out_alias.append("None")

867 for i in range(last_adv_pos + 1, len(expanded_actions)):

868 _append_non_advanced(expanded_actions[i])

869 else:

870 for _ in range(b_ndim):

871 out_alias.append("None")

872 for i, action in enumerate(expanded_actions):

873 if i in advanced_pos_set:

874 continue

875 _append_non_advanced(action)

876

877 out_layout = Layout.from_device_mesh(mesh)

878 out_layout = out_layout(*out_alias)

879

880 _copy_partial_state(self_layout, out_layout)

881 out_layout.tensor_map_to_placement()

882 out_layout.update_compact_str()

883 return out_layout

884

885

886

887def _broadcast_shapes(shapes, op_name="__getitem__"):

888 """Compute broadcast shape for a list of shapes.

889

890 Args:

891 shapes: List of shape tuples.

892 op_name: Operator name for error messages.

893

894 Returns:

895 tuple: Broadcast shape.

896 """

897 if not shapes:

898 return ()

899 result = list(shapes[0])

900 for shape in shapes[1:]:

901 ndim_diff = len(result) - len(shape)

902 if ndim_diff < 0:

903 result = [1] * (-ndim_diff) + result

904 ndim_diff = 0

905 for i, d2 in enumerate(shape):

906 d1 = result[ndim_diff + i]

907 if d1 == 1:

908 result[ndim_diff + i] = d2

909 elif d2 not in (1, d1):

910 raise ValueError(

911 f"For {op_name}, advanced index shapes {shapes} "

912 f"cannot be broadcast together."

913 )

914 return tuple(result)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_getitem.py: 92%

393 statements