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, Optional

24from hyper_parallel.platform import get_platform

25from hyper_parallel.core.dtensor.dtensor import DTensor

26from hyper_parallel.core.dtensor.layout import Layout

27from hyper_parallel.core.dtensor.placement_types import RaggedShard, Shard, StridedShard

28from .parallel_ops import DistributedOp

30platform = get_platform()

31Tensor = platform.Tensor

33_BASIC = "basic"

34_ADVANCED = "advanced"

35_BOOL_MASK = "bool_mask"

38def _normalize___getitem___args(self_t, key):

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

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

43 Args:

44 self_t: The DTensor self argument.

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

47 Returns:

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

49 """

50 return (self_t, key), {}

53def _is_long_tensor(obj) -> bool:

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

55 return isinstance(obj, Tensor) and not _is_bool_tensor(obj)

58def _is_bool_tensor(obj) -> bool:

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

60 if not isinstance(obj, Tensor):

61 return False

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

63 return False

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

67def _is_advanced_elem(k) -> bool:

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

69 if isinstance(k, list):

70 return True

71 if isinstance(k, Tensor):

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

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

74 return False

75 return True

76 return False

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

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

82 Args:

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

84 op_name: Operator name for error messages.

86 Returns:

87 tuple: A descriptor tuple for this element.

89 Raises:

90 ValueError: If the key element type is unsupported.

91 """

92 if k is None:

93 return ("none",)

94 if k is Ellipsis:

95 return ("ellipsis",)

96 if isinstance(k, int):

97 return ("int", k)

98 if isinstance(k, slice):

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

100 if isinstance(k, list):

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

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

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

104 if isinstance(k, Tensor):

105 shape = tuple(k.shape)

106 if isinstance(k, DTensor):

107 alias = tuple(k.layout.alias_tensor_map)

108 return ("idx_tensor", shape, alias)

109 return ("idx_tensor", shape)

110 raise ValueError(

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

112 )

113

114

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

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

117

118 The descriptor captures only what affects layout derivation:

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

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

121

122 Args:

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

124 op_name: Operator name for error messages.

125

126 Returns:

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

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

129

130 Raises:

131 ValueError: If key contains unsupported types.

132 """

133 if not isinstance(key, tuple):

134 key = (key,)

135

136 # First pass: detect BoolTensor anywhere in the key

137 for k in key:

138 if _is_bool_tensor(k):

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

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

141

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

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

144 kind = _ADVANCED if has_advanced else _BASIC

145

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

147 return tuple(desc), kind

148

149

150def _desc_action_to_expanded(action_type, d, input_dim):

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

152

153 Args:

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

155 d: Descriptor tuple for this element.

156 input_dim: Current input dimension index.

157

158 Returns:

159 tuple: (expanded_entries, input_dim_delta)

160 expanded_entries: list of action tuples to append.

161 input_dim_delta: how much to advance input_dim.

162 """

163 if action_type == "none":

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

165 if action_type == "int":

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

167 if action_type == "slice":

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

169 if action_type == "idx_list_len":

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

171 if action_type == "idx_tensor":

172 shape = d[1]

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

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

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

176 raise ValueError(

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

178 )

179

180

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

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

183

184 Args:

185 desc: Key descriptor tuple from _key_cache_descriptor.

186 ndim: Number of input tensor dimensions.

187 op_name: Operator name for error messages.

188

189 Returns:

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

191

192 Raises:

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

194 """

195 if desc is None:

196 return []

197

198 # Pass through bool_mask descriptor

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

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

201

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

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

204

205 if n_ellipsis > 1:

206 raise ValueError(

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

208 )

209

210 n_specified_dims = len(non_none) - n_ellipsis

211 n_fill = ndim - n_specified_dims

212

213 if n_fill < 0:

214 raise ValueError(

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

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

217 )

218

219 expanded = []

220 input_dim = 0

221

222 for d in desc:

223 action_type = d[0]

224

225 if action_type == "ellipsis":

226 for _ in range(n_fill):

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

228 input_dim += 1

229 else:

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

231 expanded.extend(entries)

232 input_dim += dim_delta

233

234 # Implicitly add full slices for unspecified trailing dimensions

235 while input_dim < ndim:

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

237 input_dim += 1

238

239 return expanded

240

241

242def _unwrap_key_for_local(key):

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

244

245 Args:

246 key: Raw key from __getitem__ call.

247

248 Returns:

249 Unwrapped key suitable for passing to local func.

250 """

251 if isinstance(key, DTensor):

252 return key.to_local()

253 if isinstance(key, tuple):

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

255 return key

256

257

258def _copy_partial_state(src_layout, dst_layout):

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

260

261 Args:

262 src_layout: Source Layout to copy partials from.

263 dst_layout: Destination Layout to copy partials to.

264 """

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

266 if op is not None:

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

268

269

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

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

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

273 return False

274 input_dim = action[-1]

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

276 normalized_step = 1 if step is None else step

277 return (

278 normalized_step == 1

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

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

281 )

282

283

284class GetItemDistributedOp(DistributedOp):

285 """Distributed implementation for tensor.__getitem__.

286

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

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

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

290

291 Sharding constraints:

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

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

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

295 - Advanced index tensors must themselves be replicated.

296 - Input must not have Partial status.

297 - slice step != 1 is not supported.

298 """

299

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

301 """Preprocess arguments for __getitem__.

302

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

304 DTensors to local tensors for execution.

305

306 Args:

307 args: (self_tensor, key)

308 kwargs: Empty dict for __getitem__.

309

310 Returns:

311 tuple: (local_args, local_kwargs, cache_values)

312 """

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

314 self_t, key = norm_args

315

316 self_layout = self_t.layout

317 global_shape = tuple(self_t.shape)

318

319 # raw key -> hashable descriptor for cache

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

321

322 # Unwrap DTensor in key to local for local execution

323 local_key = _unwrap_key_for_local(key)

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

325 local_kwargs = {}

326

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

328 return local_args, local_kwargs, cache_values

329

330 @staticmethod

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

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

333 mask_shape = "unknown"

334 for action in expanded_actions:

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

336 mask_shape = action[1]

337 break

338 raise ValueError(

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

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

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

342 )

343

344 @staticmethod

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

346 """Validate an int indexing action."""

347 input_dim = action[-1]

348 idx = action[1]

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

350 raise ValueError(

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

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

353 )

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

355 raise ValueError(

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

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

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

359 )

360

361 @staticmethod

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

363 """Validate a slice indexing action."""

364 input_dim = action[-1]

365 step = action[3]

366 step = step if step is not None else 1

367

368 if step != 1:

369 raise ValueError(

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

371 f"but got {step}."

372 )

373

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

375 raise ValueError(

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

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

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

379 )

380

381 @staticmethod

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

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

384 input_dim = action[-1]

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

386 raise ValueError(

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

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

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

390 )

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

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

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

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

395 raise ValueError(

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

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

398 )

399

400 @staticmethod

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

402 op_name="__getitem__"):

403 """Validate sharding constraints for __getitem__.

404

405 Rules:

406 1. BoolTensor mask indexing is not supported.

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

408 must be replicated.

409 3. slice step must be None or 1.

410 4. Advanced index tensors must be replicated.

411 5. int indices must be in range.

412

413 Args:

414 self_layout: Layout of self tensor.

415 expanded_actions: Expanded key actions from _descriptor_to_expanded_actions.

416 global_shape: Global shape of self tensor.

417 kind: "basic" or "advanced".

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

419

420 Raises:

421 ValueError: If any constraint is violated.

422 """

423 alias_map = self_layout.alias_tensor_map

424

425 if kind == _BOOL_MASK:

426 GetItemDistributedOp._reject_bool_mask(expanded_actions, op_name)

427

428 for action in expanded_actions:

429 action_type = action[0]

430

431 if action_type == "newaxis":

432 continue

433 if action_type == "int":

434 GetItemDistributedOp._validate_int_action(

435 action, alias_map, global_shape, op_name

436 )

437 elif action_type == "slice":

438 GetItemDistributedOp._validate_slice_action(

439 action, alias_map, global_shape, op_name

440 )

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

442 GetItemDistributedOp._validate_advanced_action(

443 action, alias_map, op_name

444 )

445

446 @staticmethod

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

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

449 placements = tuple(self_layout.placements)

450 placement = placements[0] if len(placements) == 1 else None

451 if (

452 kind != _BASIC

453 or len(global_shape) < 2

454 or len(self_layout.mesh_shape) != 1

455 or not expanded_actions

456 or expanded_actions[0][0] != "int"

457 or expanded_actions[0][-1] != 0

458 or not all(

459 _is_full_slice_action(action, global_shape)

460 for action in expanded_actions[1:]

461 )

462 or not isinstance(placement, Shard)

463 or isinstance(placement, StridedShard)

464 or not placement.is_shard(0)

465 ):

466 return None

467

468 index = expanded_actions[0][1]

469 global_dim0 = global_shape[0]

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

471 return None

472

473 mesh_size = self_layout.mesh.size(0)

474 if global_dim0 % mesh_size != 0:

475 return None

476

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

478 rows_per_rank = global_dim0 // mesh_size

479 owner_rank = normalized_index // rows_per_rank

480 local_index = normalized_index % rows_per_rank

481 output_global_shape = tuple(global_shape[1:])

482 local_units = tuple(1 if rank == owner_rank else 0 for rank in range(mesh_size))

483

484 output_layout = Layout.from_device_mesh(self_layout.mesh)

485 output_layout.set_placements((RaggedShard(tuple(range(len(output_global_shape))), local_units),))

486 output_layout.placement_to_tensor_map(len(output_global_shape))

487 return output_layout, (owner_rank, local_index, output_global_shape)

488

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

490 """Infer output layout for __getitem__.

491

492 Rules:

493 1. Input must not have Partial status.

494 2. BoolTensor mask indexing is not supported.

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

496 view on a 1-D mesh.

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

498 must be replicated.

499 5. slice step must be None or 1.

500 6. Advanced index tensors must be replicated.

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

502 dims, inserting Replicate for newaxis, and preserving sharding

503 for full-slice dims.

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

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

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

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

508

509 Args:

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

511

512 Returns:

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

514

515 Raises:

516 ValueError: If any constraint is violated.

517 """

518 self_layout = cache_values[0]

519 key_desc = cache_values[1]

520 global_shape = cache_values[2]

521 kind = cache_values[3]

522

523 # key_desc -> expanded_actions

524 expanded_actions = _descriptor_to_expanded_actions(

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

526 )

527

528 if not self._allow_partial_inputs:

529 self._check_partial_inputs([self_layout])

530

531 ragged_result = self._infer_shard_dim0_int(

532 self_layout, expanded_actions, global_shape, kind

533 )

534 if ragged_result is not None:

535 output_layout, info = ragged_result

536 return ((output_layout,), info)

537

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

539

540 if kind == _BASIC:

541 out_layout = self._infer_basic_output_layout(self_layout, expanded_actions)

542 elif kind == _ADVANCED:

543 out_layout = self._infer_advanced_output_layout(

544 self_layout, expanded_actions, op_name=self.op_name

545 )

546 else:

547 # _BOOL_MASK is rejected by _validate_input_layouts above,

548 # so we should never reach this branch.

549 raise ValueError(

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

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

552 )

553

554 return ((out_layout,), None)

555

556 def get_expand_impl(

557 self,

558 func: Callable,

559 infer_result: tuple,

560 cache_values: list,

561 ) -> Optional[Callable]:

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

563 info = infer_result[1]

564 if info is None:

565 return None

566

567 input_layout = cache_values[0]

568 output_layout = infer_result[0][0]

569 owner_rank, local_index, output_global_shape = info

570

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

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

573 del local_key

574 if not local_input.is_contiguous():

575 raise ValueError(

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

577 "requires a contiguous local tensor."

578 )

579

580 local_rank = input_layout.mesh.get_local_rank(0)

581 if local_rank == owner_rank:

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

583 else:

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

585 return DTensor.from_local_with_layout(

586 local_view,

587 output_layout,

588 shape=output_global_shape,

589 )

590

591 return ragged_getitem_impl

592

593 @staticmethod

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

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

596 if isinstance(py_output, DTensor):

597 return py_output

598 return DistributedOp.wrap_output(py_output, output_layouts)

599

600 @staticmethod

601 def _infer_basic_output_layout(self_layout, expanded_actions):

602 """Derive output layout for basic indexing.

603

604 Walk through expanded_actions:

605 - int: remove the dimension from alias_tensor_map

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

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

608 - newaxis: insert Replicate dimension

609

610 Args:

611 self_layout: Layout of self tensor.

612 expanded_actions: Expanded key actions.

613

614 Returns:

615 Layout: Output layout.

616 """

617 alias_map = self_layout.alias_tensor_map

618 mesh = self_layout.mesh

619

620 out_alias = []

621 for action in expanded_actions:

622 action_type = action[0]

623 if action_type == "newaxis":

624 out_alias.append("None")

625 elif action_type == "slice":

626 input_dim = action[4]

627 out_alias.append(alias_map[input_dim])

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

629

630 out_layout = Layout.from_device_mesh(mesh)

631 out_layout = out_layout(*out_alias)

632

633 # Copy partial state from input

634 _copy_partial_state(self_layout, out_layout)

635 out_layout.tensor_map_to_placement()

636 out_layout.update_compact_str()

637 return out_layout

638

639 @staticmethod

640 def _compute_advanced_index_info(expanded_actions, op_name):

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

642

643 Returns:

644 (advanced_positions, advanced_pos_set, are_consecutive, b_ndim)

645 """

646 advanced_actions = []

647 for i, action in enumerate(expanded_actions):

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

649 advanced_actions.append((i, action))

650

651 index_shapes = []

652 for _, action in advanced_actions:

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

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

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

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

657

658 bcast_shape = _broadcast_shapes(index_shapes, op_name=op_name)

659 b_ndim = len(bcast_shape) if bcast_shape else 0

660

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

662 advanced_pos_set = set(advanced_positions)

663 are_consecutive = (

664 len(advanced_positions) > 0

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

666 advanced_positions[-1] + 1))

667 )

668 return advanced_positions, advanced_pos_set, are_consecutive, b_ndim

669

670 @staticmethod

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

672 """Derive output layout for advanced indexing.

673

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

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

676 index tensors is inserted at position p:

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

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

679

680 Args:

681 self_layout: Layout of self tensor.

682 expanded_actions: Expanded key actions.

683 op_name: Operator name for error messages.

684

685 Returns:

686 Layout: Output layout.

687 """

688 alias_map = self_layout.alias_tensor_map

689 mesh = self_layout.mesh

690

691 advanced_positions, advanced_pos_set, are_consecutive, b_ndim = \

692 GetItemDistributedOp._compute_advanced_index_info(expanded_actions, op_name)

693

694 def _append_non_advanced(action):

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

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

697 out_alias.append("None")

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

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

700 # int: dimension removed, skip

701

702 out_alias = []

703

704 if are_consecutive and advanced_positions:

705 first_adv_pos = advanced_positions[0]

706 last_adv_pos = advanced_positions[-1]

707

708 for i in range(first_adv_pos):

709 _append_non_advanced(expanded_actions[i])

710 for _ in range(b_ndim):

711 out_alias.append("None")

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

713 _append_non_advanced(expanded_actions[i])

714 else:

715 for _ in range(b_ndim):

716 out_alias.append("None")

717 for i, action in enumerate(expanded_actions):

718 if i in advanced_pos_set:

719 continue

720 _append_non_advanced(action)

721

722 out_layout = Layout.from_device_mesh(mesh)

723 out_layout = out_layout(*out_alias)

724

725 _copy_partial_state(self_layout, out_layout)

726 out_layout.tensor_map_to_placement()

727 out_layout.update_compact_str()

728 return out_layout

729

730

731

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

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

734

735 Args:

736 shapes: List of shape tuples.

737 op_name: Operator name for error messages.

738

739 Returns:

740 tuple: Broadcast shape.

741 """

742 if not shapes:

743 return ()

744 result = list(shapes[0])

745 for shape in shapes[1:]:

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

747 if ndim_diff < 0:

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

749 ndim_diff = 0

750 for i, d2 in enumerate(shape):

751 d1 = result[ndim_diff + i]

752 if d1 == 1:

753 result[ndim_diff + i] = d2

754 elif d2 not in (1, d1):

755 raise ValueError(

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

757 f"cannot be broadcast together."

758 )

759 return tuple(result)

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

329 statements