Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_elementwise.py: 83%
308 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""
16Distributed implementation for Element-wise operator.
17"""
19import copy
20from typing import Callable, Optional, Tuple
22from .parallel_ops import DistributedOp
25_INPLACE_ELEMENTWISE_OPS = frozenset({
26 "add_", "sub_", "InplaceAddExt", "InplaceSubExt",
27})
30def _partial_signature(layout, mesh_ndim: int) -> tuple:
31 """Return one Partial entry per output mesh dimension."""
32 if layout is None:
33 return (None,) * mesh_ndim
34 partial = tuple(layout.partial)
35 if len(partial) != mesh_ndim:
36 raise ValueError(
37 f"Input and output mesh dimensions must match, but got "
38 f"input={len(partial)} and output={mesh_ndim}."
39 )
40 return partial
43def _contributes_to_partial_output(input_layout, output_layout) -> bool:
44 """Return whether this rank contributes an input to a Partial output.
46 An input that is not Partial on one of the output's Partial axes is
47 replicated along that axis. Only coordinate zero may contribute that
48 replicated value, otherwise the eventual reduction would count it once
49 per rank on the added axis.
50 """
51 output_partial = tuple(output_layout.partial)
52 input_partial = _partial_signature(input_layout, len(output_partial))
53 for mesh_dim, output_partial_type in enumerate(output_partial):
54 if (
55 output_partial_type is not None
56 and input_partial[mesh_dim] is None
57 and output_layout.mesh.get_local_rank(mesh_dim) != 0
58 ):
59 return False
60 return True
63def _zero_contribution(value):
64 """Create a strict zero while retaining a floating tensor's grad edge."""
65 is_complex = getattr(value, "is_complex", None)
66 if callable(is_complex) and is_complex():
67 return value * 0
68 if hasattr(value, "clamp"):
69 # ``value * 0`` turns +/-Inf into NaN. Clamp first so the forward value
70 # is finite, then multiply by zero to make the derivative zero even at
71 # the clamp boundary.
72 return value.clamp(0, 0) * 0
73 if isinstance(value, (bool, int, float, complex)):
74 return type(value)(0)
75 return value * 0
78def _unwrap_local_value(value):
79 """Convert DTensor-like values to local tensors while preserving containers."""
80 if hasattr(value, "_layout"):
81 return value.to_local()
82 if isinstance(value, tuple):
83 return tuple(_unwrap_local_value(item) for item in value)
84 if isinstance(value, list):
85 return [_unwrap_local_value(item) for item in value]
86 return value
89def _collect_layout_and_shape(value):
90 """Collect layout and shape from one argument for layout inference cache."""
91 layout = value.layout if hasattr(value, "_layout") else None
92 shape = value.shape if hasattr(value, "shape") else None
93 return layout, shape
96def _build_elementwise_cache_values(args, kwargs):
97 """Build cache_values from the real element-wise input arguments."""
98 input_layouts = []
99 input_shapes = []
100 for value in args:
101 layout, shape = _collect_layout_and_shape(value)
102 input_layouts.append(layout)
103 input_shapes.append(shape)
104 for value in kwargs.values():
105 layout, shape = _collect_layout_and_shape(value)
106 input_layouts.append(layout)
107 input_shapes.append(shape)
108 return [*input_layouts, input_shapes]
111class ElementWiseDistributedOp(DistributedOp):
112 """
113 Base class for distributed element-wise operators.
115 Supports broadcasting following broadcasting rules and handles
116 distributed tensor layouts with proper sharding strategy inference.
118 Args:
119 op_name (str): Name of the operator to register.
120 """
122 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
123 """
124 Preprocess arguments for element-wise operators.
126 Args:
127 args (tuple): Positional arguments passed to the operator.
128 kwargs (dict): Keyword arguments passed to the operator.
130 Returns:
131 tuple: (local_args, local_kwargs, cache_values)
132 """
133 local_kwargs = {key: _unwrap_local_value(value) for key, value in kwargs.items()}
135 local_args = tuple(_unwrap_local_value(arg) for arg in args)
136 cache_values = _build_elementwise_cache_values(args, kwargs)
137 return local_args, local_kwargs, cache_values
139 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
140 """
141 Infer output layouts for element-wise operations with broadcasting support.
143 Rules:
144 1. Inputs must not have Partial status unless the operator explicitly allows it.
145 2. Input shapes must be broadcast-compatible when shape information is available.
146 3. Broadcasting dimensions cannot be sharded.
147 4. Non-broadcast sharding patterns must be compatible across inputs.
148 5. Output layout uses the merged sharding strategy and merged Partial status.
150 Args:
151 cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes].
153 Returns:
154 tuple: ((output_layout,), None)
156 Raises:
157 ValueError: If input layouts are not compatible for broadcasting.
158 """
159 layouts = tuple(cache_values[:-1])
160 input_shapes = cache_values[-1] if cache_values else None
162 if not layouts:
163 return None
165 valid_layouts = [layout for layout in layouts if layout is not None]
167 if not valid_layouts:
168 return None
170 if not self._allow_partial_inputs:
171 self._check_partial_inputs(layouts)
173 if len(valid_layouts) == 1:
174 return ((copy.deepcopy(valid_layouts[0]),), None)
176 if not input_shapes:
177 return ((self._handle_no_input_shapes(valid_layouts),), None)
179 aligned_layouts, aligned_shapes = self._align_layouts_and_shapes(layouts, input_shapes)
181 if len(aligned_layouts) <= 1 or len(aligned_layouts) != len(aligned_shapes):
182 return ((copy.deepcopy(valid_layouts[0]),), None)
184 output_shape = self._compute_output_shape(aligned_shapes)
185 merged_tensor_map, merged_partial = self._merge_all_layouts(
186 aligned_layouts,
187 aligned_shapes,
188 output_shape,
189 layouts
190 )
192 self._check_all_inputs_broadcasts_and_partial(aligned_layouts, aligned_shapes, output_shape)
194 output_layout = self._create_output_layout(aligned_layouts[0], merged_tensor_map, merged_partial)
195 return ((output_layout,), None)
197 def _handle_no_input_shapes(self, valid_layouts):
198 """
199 Handle the case when input shapes are not available.
200 """
201 first_layout = valid_layouts[0]
202 first_alias_map = first_layout.alias_tensor_map
203 for layout in valid_layouts[1:]:
204 if layout.alias_tensor_map != first_alias_map:
205 raise ValueError(
206 f"For {self.op_name}, cannot infer layout without shapes: "
207 f"mismatched alias_tensor_map {first_alias_map} vs {layout.alias_tensor_map}."
208 )
209 return copy.deepcopy(first_layout)
211 def _align_layouts_and_shapes(self, layouts, input_shapes):
212 """
213 Align layouts with shapes by position, skipping None layouts.
214 """
215 aligned_layouts = []
216 aligned_shapes = []
217 for layout, shape in zip(layouts, input_shapes):
218 if layout is None:
219 continue
220 aligned_layouts.append(layout)
221 aligned_shapes.append(shape)
222 return aligned_layouts, aligned_shapes
224 def _compute_output_shape(self, aligned_shapes):
225 """
226 Compute broadcasted output shape from all input shapes.
227 """
228 output_shape = aligned_shapes[0]
229 for shape in aligned_shapes[1:]:
230 output_shape = self._broadcast_shapes(output_shape, shape)
231 return output_shape
233 def _merge_all_layouts(self, aligned_layouts, aligned_shapes, output_shape, layouts):
234 """
235 Merge all input layouts sequentially to get final tensor_map and partial status.
236 """
237 base_layout = aligned_layouts[0]
239 merged_tensor_map = self._merge_tensor_maps_for_broadcast(
240 aligned_layouts[0],
241 aligned_layouts[1],
242 aligned_shapes[0],
243 aligned_shapes[1],
244 output_shape
245 )
247 merged_partial = self._merge_partial_status(
248 base_layout.partial,
249 aligned_layouts[1].partial,
250 merged_tensor_map,
251 aligned_layouts[0].tensor_map if aligned_layouts[0].tensor_map else tuple(),
252 aligned_layouts[1].tensor_map if aligned_layouts[1].tensor_map else tuple(),
253 layouts
254 )
256 for i in range(2, len(aligned_layouts)):
257 temp_layout = self._create_output_layout(base_layout, merged_tensor_map, merged_partial)
258 merged_tensor_map = self._merge_tensor_maps_for_broadcast(
259 temp_layout,
260 aligned_layouts[i],
261 output_shape,
262 aligned_shapes[i],
263 output_shape
264 )
265 merged_partial = self._merge_partial_status(
266 merged_partial,
267 aligned_layouts[i].partial,
268 merged_tensor_map,
269 temp_layout.tensor_map if temp_layout.tensor_map else tuple(),
270 aligned_layouts[i].tensor_map if aligned_layouts[i].tensor_map else tuple(),
271 layouts
272 )
274 return merged_tensor_map, merged_partial
276 def _merge_partial_status(self, partial1, partial2, merged_tensor_map, tensor_map1, tensor_map2, layouts):
277 """
278 Merge partial status from two inputs.
280 Rules:
281 1. Both None → None
282 2. One None → Use the other
283 3. Both not None and same → Use it
284 4. Both not None and different → Error
285 5. Check Shard + Partial conflicts for each input
287 Args:
288 partial1: Partial status list from first input
289 partial2: Partial status list from second input
290 merged_tensor_map: Merged tensor map for output
291 tensor_map1: Tensor map of first input
292 tensor_map2: Tensor map of second input
294 Returns:
295 List: Merged partial status
297 Raises:
298 ValueError: If partial operations conflict or Shard+Partial conflict found
299 """
300 # Check Shard + Partial conflicts for input1
301 self._check_shard_partial_conflict(tensor_map1, partial1, layouts)
303 # Check Shard + Partial conflicts for input2
304 self._check_shard_partial_conflict(tensor_map2, partial2, layouts)
306 # Determine mesh dimension from partial lists
307 mesh_dim = max(len(partial1) if partial1 else 0, len(partial2) if partial2 else 0)
309 merged_partial = [None] * mesh_dim
311 for i in range(mesh_dim):
312 op1 = partial1[i] if partial1 and i < len(partial1) else None
313 op2 = partial2[i] if partial2 and i < len(partial2) else None
315 # Both have partial status with different operations
316 if op1 is not None and op2 is not None and op1 != op2:
317 raise ValueError(
318 f"For {self.op_name}, partial operations should be same for device axis {i}, "
319 f"but got {op1} and {op2}"
320 )
322 # Merge: prefer non-None, or either if both same
323 if op1 is not None:
324 merged_partial[i] = op1
325 elif op2 is not None:
326 merged_partial[i] = op2
328 # Check final output for Shard + Partial conflicts
329 self._check_shard_partial_conflict(merged_tensor_map, merged_partial, layouts)
331 return merged_partial
333 def _check_shard_partial_conflict(self, tensor_map, partial_list, layouts):
334 """
335 Check for conflicts between Shard and Partial on same device axis.
337 Args:
338 tensor_map: Tensor map to check
339 partial_list: Partial status list
341 Raises:
342 ValueError: If Shard and Partial conflict found
343 """
344 if not partial_list:
345 return
347 mesh_dim = len(partial_list)
349 # Collect all device axis used for sharding
350 sharded_axis = set()
351 if tensor_map:
352 for map_val in tensor_map:
353 if isinstance(map_val, tuple):
354 for sub_val in map_val:
355 if sub_val != -1:
356 # Convert to device axis index
357 axis_idx = mesh_dim - 1 - sub_val
358 sharded_axis.add(axis_idx)
359 elif map_val != -1:
360 axis_idx = mesh_dim - 1 - map_val
361 sharded_axis.add(axis_idx)
363 # Check if any sharded axis has partial status
364 for axis_idx in sharded_axis:
365 if 0 <= axis_idx < len(partial_list) and partial_list[axis_idx] is not None:
366 raise ValueError(
367 f"For {self.op_name}, Shard and Partial should not coexist on same device axis "
368 f"{axis_idx}, but got Partial({partial_list[axis_idx]}). "
369 f"Please check layouts: {layouts}."
370 )
372 def _check_all_inputs_broadcasts_and_partial(self, layouts, input_shapes, output_shape):
373 """
374 Check if any input broadcasts and has Partial status.
375 """
376 for i, (layout, input_shape) in enumerate(zip(layouts, input_shapes)):
377 if layout is None:
378 continue
380 input_name = f"input{i+1}"
382 input_len = len(input_shape)
383 output_len = len(output_shape)
385 if input_len < output_len:
386 aligned_input_shape = (1,) * (output_len - input_len) + tuple(input_shape)
387 else:
388 aligned_input_shape = input_shape
390 broadcasts = False
391 for in_dim, out_dim in zip(aligned_input_shape, output_shape):
392 if in_dim == 1 and out_dim > 1:
393 broadcasts = True
394 break
396 if broadcasts and layout.is_partial():
397 raise ValueError(
398 f"For {self.op_name}, {input_name} has Partial status and broadcasts. "
399 f"Should be without Partial status for broadcasting without communication"
400 )
402 def _broadcast_shapes(self, shape1, shape2):
403 """
404 Calculate the broadcasted shape of two shapes according to broadcasting rules.
406 Broadcasting rules:
407 1. If two arrays have different numbers of dimensions, pad the shape of the
408 lower-dimensional array with 1s on the left until both shapes have the same length.
409 2. If two arrays have the same number of dimensions but different lengths in some
410 dimensions, dimensions with length 1 will be expanded to match the other array's
411 dimension length.
412 3. If two arrays have the same number of dimensions but any dimension has different
413 lengths and neither is 1, raise an error.
415 Args:
416 shape1 (tuple): Shape of the first tensor, e.g., (3, 1, 5)
417 shape2 (tuple): Shape of the second tensor, e.g., (4, 5)
419 Returns:
420 tuple: Broadcasted shape, e.g., (3, 4, 5)
422 Raises:
423 ValueError: If shapes cannot be broadcast together.
424 """
425 # Rule 1: Right-align, pad with 1s on the left to make dimensions equal
426 len1, len2 = len(shape1), len(shape2)
427 max_len = max(len1, len2)
429 padded_shape1 = (1,) * (max_len - len1) + tuple(shape1)
430 padded_shape2 = (1,) * (max_len - len2) + tuple(shape2)
432 # Rules 2 and 3: Check if each dimension can be broadcast
433 result_shape = []
434 for dim1, dim2 in zip(padded_shape1, padded_shape2):
435 if dim1 == dim2:
436 # Dimensions are the same, use directly
437 result_shape.append(dim1)
438 elif dim1 == 1:
439 # First shape has 1 in this dimension, expand to dim2
440 result_shape.append(dim2)
441 elif dim2 == 1:
442 # Second shape has 1 in this dimension, expand to dim1
443 result_shape.append(dim1)
444 else:
445 # Rule 3: Dimensions are different and neither is 1, cannot broadcast
446 raise ValueError(
447 f"For {self.op_name}, shapes {shape1} and {shape2} cannot be broadcast together. "
448 f"Dimension mismatch: {dim1} vs {dim2}"
449 )
451 return tuple(result_shape)
453 def _align_tensor_maps_for_broadcast(self, layout1, layout2, shape1, shape2):
454 """
455 Align tensor_maps of two layouts to support broadcasting.
457 When two tensors have different dimensions, the tensor_map of the
458 lower-dimensional tensor is padded with -1 (indicating no sharding) at the front.
460 Args:
461 layout1: Layout of the first tensor
462 layout2: Layout of the second tensor
463 shape1 (tuple): Global shape of the first tensor
464 shape2 (tuple): Global shape of the second tensor
466 Returns:
467 tuple: (aligned_map1, aligned_map2) - Aligned tensor_maps
468 """
469 len1, len2 = len(shape1), len(shape2)
470 max_len = max(len1, len2)
472 map1 = layout1.tensor_map if layout1.tensor_map else tuple([-1] * len1)
473 map2 = layout2.tensor_map if layout2.tensor_map else tuple([-1] * len2)
475 aligned_map1 = (-1,) * (max_len - len1) + map1
476 aligned_map2 = (-1,) * (max_len - len2) + map2
478 return aligned_map1, aligned_map2
480 def _normalize_tensor_map_element(self, map_element):
481 """
482 Normalize a tensor_map element to a tuple of device axis for unified processing.
484 Args:
485 map_element: Element from tensor_map, can be:
486 - int: -1 (no sharding) or device axis index
487 - tuple: multiple device axis
489 Returns:
490 tuple: Tuple of device axis (empty tuple if not sharded)
491 """
492 if map_element == -1:
493 return ()
494 if isinstance(map_element, int):
495 return (map_element,)
496 if isinstance(map_element, tuple):
497 return tuple(dim for dim in map_element if dim != -1)
498 return ()
500 def _denormalize_tensor_map_element(self, device_axis_tuple):
501 """
502 Convert a tuple of device axis back to tensor_map element format.
504 Args:
505 device_axis_tuple (tuple): Tuple of device axis
507 Returns:
508 int or tuple: -1 if empty, single int if one element, tuple if multiple elements
509 """
510 if not device_axis_tuple:
511 return -1
512 if len(device_axis_tuple) == 1:
513 return device_axis_tuple[0]
514 return device_axis_tuple
516 def _merge_tensor_maps_for_broadcast(self, layout1, layout2, shape1, shape2, output_shape):
517 """
518 Merge tensor_maps of two inputs to generate output tensor_map.
520 This method handles both simple int-type and complex tuple-type tensor_map elements,
521 ensuring correct sharding strategy for the broadcasted output.
523 Args:
524 layout1: Layout of the first input
525 layout2: Layout of the second input
526 shape1 (tuple): Global shape of the first input
527 shape2 (tuple): Global shape of the second input
528 output_shape (tuple): Global shape of the output
530 Returns:
531 tuple: Merged tensor_map for the output
533 Raises:
534 ValueError: If sharding strategies conflict or broadcasting dimension is sharded
535 """
536 map1, map2 = self._align_tensor_maps_for_broadcast(layout1, layout2, shape1, shape2)
538 len1, len2 = len(shape1), len(shape2)
539 max_len = len(output_shape)
540 padded_shape1 = (1,) * (max_len - len1) + tuple(shape1)
541 padded_shape2 = (1,) * (max_len - len2) + tuple(shape2)
543 merged_map = []
544 for i, (dim1, dim2, out_dim) in enumerate(zip(padded_shape1, padded_shape2, output_shape)):
545 m1, m2 = map1[i], map2[i]
547 m1_axis = self._normalize_tensor_map_element(m1)
548 m2_axis = self._normalize_tensor_map_element(m2)
550 m1_axis_for_compare = frozenset(m1_axis)
551 m2_axis_for_compare = frozenset(m2_axis)
553 m1_is_sharded = bool(m1_axis)
554 m2_is_sharded = bool(m2_axis)
556 if not m1_is_sharded and not m2_is_sharded:
557 merged_map.append(-1)
559 elif not m1_is_sharded:
560 if dim2 == 1 and out_dim > 1:
561 raise ValueError(
562 f"For {self.op_name}, dimension {i} of second input has size 1 "
563 f"but is sharded on device axis {m2_axis}. "
564 f"Broadcasting dimension cannot be sharded."
565 )
566 merged_map.append(self._denormalize_tensor_map_element(m2_axis))
568 elif not m2_is_sharded:
569 if dim1 == 1 and out_dim > 1:
570 raise ValueError(
571 f"For {self.op_name}, dimension {i} of first input has size 1 "
572 f"but is sharded on device axis {m1_axis}. "
573 f"Broadcasting dimension cannot be sharded."
574 )
575 merged_map.append(self._denormalize_tensor_map_element(m1_axis))
577 else:
578 if m1_axis_for_compare != m2_axis_for_compare:
579 raise ValueError(
580 f"For {self.op_name}, inputs should have same sharding pattern, "
581 f"but got confilcting sharding at dimension {i}, "
582 f"input1 shaded on {m1_axis} and input2 shaded on {m2_axis}."
583 )
585 if (dim1 == 1 or dim2 == 1) and dim1 != dim2:
586 raise ValueError(
587 f"For {self.op_name}, dimension {i} is broadcast from size 1 "
588 f"to {out_dim} but is sharded on device axis {m1_axis}. "
589 f"Broadcasting dimension cannot be sharded."
590 )
592 merged_map.append(self._denormalize_tensor_map_element(m1_axis))
594 return tuple(merged_map)
596 def _create_output_layout(self, base_layout, output_tensor_map, partial_list=None):
597 """
598 Create output layout based on input layout.
600 Args:
601 base_layout: Base layout (usually from the first input)
602 output_tensor_map (tuple): Tensor_map for the output
603 partial_list (list): Partial status list for the output
605 Returns:
606 Layout: New Layout object with updated tensor_map and alias_tensor_map
607 """
608 new_layout = copy.deepcopy(base_layout)
609 new_layout.set_tensor_map(output_tensor_map)
611 alias_tensor_map = []
612 for tensor_dim in output_tensor_map:
613 if tensor_dim == -1:
614 alias_tensor_map.append("None")
615 elif isinstance(tensor_dim, tuple):
616 alias_tuple = tuple(
617 base_layout.alias_name[len(base_layout.alias_name) - 1 - dim]
618 for dim in tensor_dim
619 if dim != -1
620 )
621 alias_tensor_map.append(alias_tuple if alias_tuple else "None")
622 else:
623 alias_tensor_map.append(
624 base_layout.alias_name[len(base_layout.alias_name) - 1 - tensor_dim]
625 )
627 new_layout.set_alias_tensor_map(tuple(alias_tensor_map))
628 new_layout.tensor_map_to_placement()
630 # Set partial status if provided
631 if partial_list:
632 for i, partial_op in enumerate(partial_list):
633 if partial_op is not None and i < len(new_layout.alias_name):
634 new_layout.set_partial_by_dev_axis(new_layout.alias_name[i], partial_op)
636 return new_layout
639class ElementWiseWithPartialDistributedOp(ElementWiseDistributedOp):
640 """
641 Base class for elementwise operations that support partial status propagation.
642 """
643 def __init__(self, op_name):
644 super().__init__(op_name)
645 self._allow_partial_inputs = True
648class AddDistributedOp(ElementWiseWithPartialDistributedOp):
649 """
650 Distributed implementation for Add operator.
652 This operator supports partial status propagation from inputs to output,
653 which is useful for operations like gradient accumulation where partial
654 results need to be preserved through the computation graph.
655 """
657 @staticmethod
658 def _canonicalize_binary_operands(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
659 """Move semantic ``input`` and ``other`` operands before option values."""
660 local_kwargs = dict(kwargs)
661 if args:
662 if "input" in local_kwargs:
663 raise ValueError("Binary elementwise input was provided both positionally and by keyword.")
664 lhs = args[0]
665 trailing_args = args[1:]
666 elif "input" in local_kwargs:
667 lhs = local_kwargs.pop("input")
668 trailing_args = ()
669 else:
670 return args, local_kwargs
672 if trailing_args:
673 if "other" in local_kwargs:
674 raise ValueError("Binary elementwise other was provided both positionally and by keyword.")
675 rhs = trailing_args[0]
676 trailing_args = trailing_args[1:]
677 elif "other" in local_kwargs:
678 rhs = local_kwargs.pop("other")
679 else:
680 return args, local_kwargs
681 return (lhs, rhs, *trailing_args), local_kwargs
683 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
684 """Canonicalize binary operands before local unwrapping and layout caching.
686 Args:
687 args: Positional backend arguments.
688 kwargs: Keyword backend arguments, potentially including
689 ``input``, ``other``, and options such as ``alpha``.
691 Returns:
692 Local arguments, local keyword arguments, and layout cache values
693 with the semantic left and right operands in the first two slots.
694 """
695 canonical_args, canonical_kwargs = self._canonicalize_binary_operands(args, kwargs)
696 return super().preprocess(canonical_args, canonical_kwargs)
698 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
699 """
700 Infer output layout for Add operator.
702 Rules:
703 1. Follow element-wise broadcasting and sharding merge rules.
704 2. Partial status is allowed and propagated.
705 3. Propagated Partial status must be "sum" or None.
707 Args:
708 cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes].
710 Returns:
711 tuple: ((output_layout,), None)
713 Raises:
714 ValueError: If propagated Partial status is not "sum" or None.
715 """
716 layouts = tuple(cache_values[:-1])
717 if self.op_name in _INPLACE_ELEMENTWISE_OPS and len(layouts) >= 2:
718 mesh_ndim = max(
719 (len(layout.partial) for layout in layouts[:2] if layout is not None),
720 default=0,
721 )
722 partial_signatures = tuple(
723 _partial_signature(layout, mesh_ndim)
724 for layout in layouts[:2]
725 )
726 if partial_signatures[0] != partial_signatures[1]:
727 raise ValueError(
728 f"For {self.op_name}, input Partial placements should be identical "
729 f"for in-place execution, but got {partial_signatures}."
730 )
732 infer_result = super().infer_layout(cache_values)
733 if infer_result is None:
734 return infer_result
736 output_layout = infer_result[0][0]
737 for i, partial_type in enumerate(output_layout.partial):
738 if partial_type is not None and partial_type != "sum":
739 raise ValueError(
740 f"For {self.op_name}, inputs partial status should be 'sum' or None, "
741 f"but got {partial_type} at index {i}."
742 )
744 return infer_result
746 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
747 cache_values: list) -> Optional[Callable]:
748 """Build a rank-local implementation that avoids repeated values.
750 Args:
751 func: Local backend add or subtract implementation.
752 infer_result: Inferred output layout and auxiliary result.
753 cache_values: Input layouts followed by input shapes.
755 Returns:
756 A contribution-aware local implementation, or ``None`` when both
757 inputs already have identical Partial placements.
758 """
759 layouts = tuple(cache_values[:-1])
760 x1_layout = layouts[0]
761 x2_layout = layouts[1]
762 output_layout = infer_result[0][0]
763 output_mesh_ndim = len(output_layout.partial)
764 x1_partial = _partial_signature(x1_layout, output_mesh_ndim)
765 x2_partial = _partial_signature(x2_layout, output_mesh_ndim)
767 if x1_partial == x2_partial:
768 return None
770 x1_contributes = _contributes_to_partial_output(x1_layout, output_layout)
771 x2_contributes = _contributes_to_partial_output(x2_layout, output_layout)
773 # ``*args``/``**kwargs`` forward add-specific extras such as ``alpha``.
774 # A strict zero retains a valid zero-gradient autograd edge.
775 def _expand_impl(x1, x2, *args, **kwargs):
776 local_x1 = x1 if x1_contributes else _zero_contribution(x1)
777 local_x2 = x2 if x2_contributes else _zero_contribution(x2)
778 return func(local_x1, local_x2, *args, **kwargs)
780 return _expand_impl
783class SubDistributedOp(AddDistributedOp):
784 """Distributed subtraction with correct replicated-minus-partial signs.
786 When the first operand is replicated and the second is Partial(sum), only
787 the first Partial rank may contribute the replicated value. Every other
788 rank must contribute ``0 - x2_local`` so reducing the output reconstructs
789 ``x1 - sum(x2_local)``.
790 """