Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_elementwise.py: 78%
307 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +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 NOTE: aclop packed-args normalization (for MindSpore aclop operators
127 like Mod, StopGradient that pack args as ``(prim, name, (real_args...))``)
128 is handled upstream in ``OpDispatcher._dispatch_layout_infer`` via
129 ``_normalize_aclop_args``. This method receives clean unpacked args.
131 Args:
132 args (tuple): Positional arguments passed to the operator.
133 kwargs (dict): Keyword arguments passed to the operator.
135 Returns:
136 tuple: (local_args, local_kwargs, cache_values)
137 """
138 local_kwargs = {key: _unwrap_local_value(value) for key, value in kwargs.items()}
140 local_args = tuple(_unwrap_local_value(arg) for arg in args)
141 cache_values = _build_elementwise_cache_values(args, kwargs)
142 return local_args, local_kwargs, cache_values
144 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
145 """
146 Infer output layouts for element-wise operations with broadcasting support.
148 Rules:
149 1. Inputs must not have Partial status unless the operator explicitly allows it.
150 2. Input shapes must be broadcast-compatible when shape information is available.
151 3. Broadcasting dimensions cannot be sharded.
152 4. Non-broadcast sharding patterns must be compatible across inputs.
153 5. Output layout uses the merged sharding strategy and merged Partial status.
155 Args:
156 cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes].
158 Returns:
159 tuple: ((output_layout,), None)
161 Raises:
162 ValueError: If input layouts are not compatible for broadcasting.
163 """
164 layouts = tuple(cache_values[:-1])
165 input_shapes = cache_values[-1] if cache_values else None
167 if not layouts:
168 return None
170 valid_layouts = [layout for layout in layouts if layout is not None]
172 if not valid_layouts:
173 return None
175 if not self._allow_partial_inputs:
176 self._check_partial_inputs(layouts)
178 if len(valid_layouts) == 1:
179 return ((copy.deepcopy(valid_layouts[0]),), None)
181 if not input_shapes:
182 return ((self._handle_no_input_shapes(valid_layouts),), None)
184 aligned_layouts, aligned_shapes = self._align_layouts_and_shapes(layouts, input_shapes)
186 if len(aligned_layouts) <= 1 or len(aligned_layouts) != len(aligned_shapes):
187 return ((copy.deepcopy(valid_layouts[0]),), None)
189 output_shape = self._compute_output_shape(aligned_shapes)
190 merged_tensor_map, merged_partial = self._merge_all_layouts(
191 aligned_layouts,
192 aligned_shapes,
193 output_shape,
194 layouts
195 )
197 self._check_all_inputs_broadcasts_and_partial(aligned_layouts, aligned_shapes, output_shape)
199 output_layout = self._create_output_layout(aligned_layouts[0], merged_tensor_map, merged_partial)
200 return ((output_layout,), None)
202 def _handle_no_input_shapes(self, valid_layouts):
203 """
204 Handle the case when input shapes are not available.
205 """
206 first_layout = valid_layouts[0]
207 first_alias_map = first_layout.alias_tensor_map
208 for layout in valid_layouts[1:]:
209 if layout.alias_tensor_map != first_alias_map:
210 raise ValueError(
211 f"For {self.op_name}, cannot infer layout without shapes: "
212 f"mismatched alias_tensor_map {first_alias_map} vs {layout.alias_tensor_map}."
213 )
214 return copy.deepcopy(first_layout)
216 def _align_layouts_and_shapes(self, layouts, input_shapes):
217 """
218 Align layouts with shapes by position, skipping None layouts.
219 """
220 aligned_layouts = []
221 aligned_shapes = []
222 for layout, shape in zip(layouts, input_shapes):
223 if layout is None:
224 continue
225 aligned_layouts.append(layout)
226 aligned_shapes.append(shape)
227 return aligned_layouts, aligned_shapes
229 def _compute_output_shape(self, aligned_shapes):
230 """
231 Compute broadcasted output shape from all input shapes.
232 """
233 output_shape = aligned_shapes[0]
234 for shape in aligned_shapes[1:]:
235 output_shape = self._broadcast_shapes(output_shape, shape)
236 return output_shape
238 def _merge_all_layouts(self, aligned_layouts, aligned_shapes, output_shape, layouts):
239 """
240 Merge all input layouts sequentially to get final tensor_map and partial status.
241 """
242 base_layout = aligned_layouts[0]
244 merged_tensor_map = self._merge_tensor_maps_for_broadcast(
245 aligned_layouts[0],
246 aligned_layouts[1],
247 aligned_shapes[0],
248 aligned_shapes[1],
249 output_shape
250 )
252 merged_partial = self._merge_partial_status(
253 base_layout.partial,
254 aligned_layouts[1].partial,
255 merged_tensor_map,
256 aligned_layouts[0].tensor_map if aligned_layouts[0].tensor_map else tuple(),
257 aligned_layouts[1].tensor_map if aligned_layouts[1].tensor_map else tuple(),
258 layouts
259 )
261 for i in range(2, len(aligned_layouts)):
262 temp_layout = self._create_output_layout(base_layout, merged_tensor_map, merged_partial)
263 merged_tensor_map = self._merge_tensor_maps_for_broadcast(
264 temp_layout,
265 aligned_layouts[i],
266 output_shape,
267 aligned_shapes[i],
268 output_shape
269 )
270 merged_partial = self._merge_partial_status(
271 merged_partial,
272 aligned_layouts[i].partial,
273 merged_tensor_map,
274 temp_layout.tensor_map if temp_layout.tensor_map else tuple(),
275 aligned_layouts[i].tensor_map if aligned_layouts[i].tensor_map else tuple(),
276 layouts
277 )
279 return merged_tensor_map, merged_partial
281 def _merge_partial_status(self, partial1, partial2, merged_tensor_map, tensor_map1, tensor_map2, layouts):
282 """
283 Merge partial status from two inputs.
285 Rules:
286 1. Both None → None
287 2. One None → Use the other
288 3. Both not None and same → Use it
289 4. Both not None and different → Error
290 5. Check Shard + Partial conflicts for each input
292 Args:
293 partial1: Partial status list from first input
294 partial2: Partial status list from second input
295 merged_tensor_map: Merged tensor map for output
296 tensor_map1: Tensor map of first input
297 tensor_map2: Tensor map of second input
299 Returns:
300 List: Merged partial status
302 Raises:
303 ValueError: If partial operations conflict or Shard+Partial conflict found
304 """
305 # Check Shard + Partial conflicts for input1
306 self._check_shard_partial_conflict(tensor_map1, partial1, layouts)
308 # Check Shard + Partial conflicts for input2
309 self._check_shard_partial_conflict(tensor_map2, partial2, layouts)
311 # Determine mesh dimension from partial lists
312 mesh_dim = max(len(partial1) if partial1 else 0, len(partial2) if partial2 else 0)
314 merged_partial = [None] * mesh_dim
316 for i in range(mesh_dim):
317 op1 = partial1[i] if partial1 and i < len(partial1) else None
318 op2 = partial2[i] if partial2 and i < len(partial2) else None
320 # Both have partial status with different operations
321 if op1 is not None and op2 is not None and op1 != op2:
322 raise ValueError(
323 f"For {self.op_name}, partial operations should be same for device axis {i}, "
324 f"but got {op1} and {op2}"
325 )
327 # Merge: prefer non-None, or either if both same
328 if op1 is not None:
329 merged_partial[i] = op1
330 elif op2 is not None:
331 merged_partial[i] = op2
333 # Check final output for Shard + Partial conflicts
334 self._check_shard_partial_conflict(merged_tensor_map, merged_partial, layouts)
336 return merged_partial
338 def _check_shard_partial_conflict(self, tensor_map, partial_list, layouts):
339 """
340 Check for conflicts between Shard and Partial on same device axis.
342 Args:
343 tensor_map: Tensor map to check
344 partial_list: Partial status list
346 Raises:
347 ValueError: If Shard and Partial conflict found
348 """
349 if not partial_list:
350 return
352 mesh_dim = len(partial_list)
354 # Collect all device axis used for sharding
355 sharded_axis = set()
356 if tensor_map:
357 for map_val in tensor_map:
358 if isinstance(map_val, tuple):
359 for sub_val in map_val:
360 if sub_val != -1:
361 # Convert to device axis index
362 axis_idx = mesh_dim - 1 - sub_val
363 sharded_axis.add(axis_idx)
364 elif map_val != -1:
365 axis_idx = mesh_dim - 1 - map_val
366 sharded_axis.add(axis_idx)
368 # Check if any sharded axis has partial status
369 for axis_idx in sharded_axis:
370 if 0 <= axis_idx < len(partial_list) and partial_list[axis_idx] is not None:
371 raise ValueError(
372 f"For {self.op_name}, Shard and Partial should not coexist on same device axis "
373 f"{axis_idx}, but got Partial({partial_list[axis_idx]}). "
374 f"Please check layouts: {layouts}."
375 )
377 def _check_all_inputs_broadcasts_and_partial(self, layouts, input_shapes, output_shape):
378 """
379 Check if any input broadcasts and has Partial status.
380 """
381 for i, (layout, input_shape) in enumerate(zip(layouts, input_shapes)):
382 if layout is None:
383 continue
385 input_name = f"input{i+1}"
387 input_len = len(input_shape)
388 output_len = len(output_shape)
390 if input_len < output_len:
391 aligned_input_shape = (1,) * (output_len - input_len) + tuple(input_shape)
392 else:
393 aligned_input_shape = input_shape
395 broadcasts = False
396 for in_dim, out_dim in zip(aligned_input_shape, output_shape):
397 if in_dim == 1 and out_dim > 1:
398 broadcasts = True
399 break
401 if broadcasts and layout.is_partial():
402 raise ValueError(
403 f"For {self.op_name}, {input_name} has Partial status and broadcasts. "
404 f"Should be without Partial status for broadcasting without communication"
405 )
407 def _broadcast_shapes(self, shape1, shape2):
408 """
409 Calculate the broadcasted shape of two shapes according to broadcasting rules.
411 Broadcasting rules:
412 1. If two arrays have different numbers of dimensions, pad the shape of the
413 lower-dimensional array with 1s on the left until both shapes have the same length.
414 2. If two arrays have the same number of dimensions but different lengths in some
415 dimensions, dimensions with length 1 will be expanded to match the other array's
416 dimension length.
417 3. If two arrays have the same number of dimensions but any dimension has different
418 lengths and neither is 1, raise an error.
420 Args:
421 shape1 (tuple): Shape of the first tensor, e.g., (3, 1, 5)
422 shape2 (tuple): Shape of the second tensor, e.g., (4, 5)
424 Returns:
425 tuple: Broadcasted shape, e.g., (3, 4, 5)
427 Raises:
428 ValueError: If shapes cannot be broadcast together.
429 """
430 # Rule 1: Right-align, pad with 1s on the left to make dimensions equal
431 len1, len2 = len(shape1), len(shape2)
432 max_len = max(len1, len2)
434 padded_shape1 = (1,) * (max_len - len1) + tuple(shape1)
435 padded_shape2 = (1,) * (max_len - len2) + tuple(shape2)
437 # Rules 2 and 3: Check if each dimension can be broadcast
438 result_shape = []
439 for dim1, dim2 in zip(padded_shape1, padded_shape2):
440 if dim1 == dim2:
441 # Dimensions are the same, use directly
442 result_shape.append(dim1)
443 elif dim1 == 1:
444 # First shape has 1 in this dimension, expand to dim2
445 result_shape.append(dim2)
446 elif dim2 == 1:
447 # Second shape has 1 in this dimension, expand to dim1
448 result_shape.append(dim1)
449 else:
450 # Rule 3: Dimensions are different and neither is 1, cannot broadcast
451 raise ValueError(
452 f"For {self.op_name}, shapes {shape1} and {shape2} cannot be broadcast together. "
453 f"Dimension mismatch: {dim1} vs {dim2}"
454 )
456 return tuple(result_shape)
458 def _align_tensor_maps_for_broadcast(self, layout1, layout2, shape1, shape2):
459 """
460 Align tensor_maps of two layouts to support broadcasting.
462 When two tensors have different dimensions, the tensor_map of the
463 lower-dimensional tensor is padded with -1 (indicating no sharding) at the front.
465 Args:
466 layout1: Layout of the first tensor
467 layout2: Layout of the second tensor
468 shape1 (tuple): Global shape of the first tensor
469 shape2 (tuple): Global shape of the second tensor
471 Returns:
472 tuple: (aligned_map1, aligned_map2) - Aligned tensor_maps
473 """
474 len1, len2 = len(shape1), len(shape2)
475 max_len = max(len1, len2)
477 map1 = layout1.tensor_map if layout1.tensor_map else tuple([-1] * len1)
478 map2 = layout2.tensor_map if layout2.tensor_map else tuple([-1] * len2)
480 aligned_map1 = (-1,) * (max_len - len1) + map1
481 aligned_map2 = (-1,) * (max_len - len2) + map2
483 return aligned_map1, aligned_map2
485 def _normalize_tensor_map_element(self, map_element):
486 """
487 Normalize a tensor_map element to a tuple of device axis for unified processing.
489 Args:
490 map_element: Element from tensor_map, can be:
491 - int: -1 (no sharding) or device axis index
492 - tuple: multiple device axis
494 Returns:
495 tuple: Tuple of device axis (empty tuple if not sharded)
496 """
497 if map_element == -1:
498 return ()
499 if isinstance(map_element, int):
500 return (map_element,)
501 if isinstance(map_element, tuple):
502 return tuple(dim for dim in map_element if dim != -1)
503 return ()
505 def _denormalize_tensor_map_element(self, device_axis_tuple):
506 """
507 Convert a tuple of device axis back to tensor_map element format.
509 Args:
510 device_axis_tuple (tuple): Tuple of device axis
512 Returns:
513 int or tuple: -1 if empty, single int if one element, tuple if multiple elements
514 """
515 if not device_axis_tuple:
516 return -1
517 if len(device_axis_tuple) == 1:
518 return device_axis_tuple[0]
519 return device_axis_tuple
521 def _merge_tensor_maps_for_broadcast(self, layout1, layout2, shape1, shape2, output_shape):
522 """
523 Merge tensor_maps of two inputs to generate output tensor_map.
525 This method handles both simple int-type and complex tuple-type tensor_map elements,
526 ensuring correct sharding strategy for the broadcasted output.
528 Args:
529 layout1: Layout of the first input
530 layout2: Layout of the second input
531 shape1 (tuple): Global shape of the first input
532 shape2 (tuple): Global shape of the second input
533 output_shape (tuple): Global shape of the output
535 Returns:
536 tuple: Merged tensor_map for the output
538 Raises:
539 ValueError: If sharding strategies conflict or broadcasting dimension is sharded
540 """
541 map1, map2 = self._align_tensor_maps_for_broadcast(layout1, layout2, shape1, shape2)
543 len1, len2 = len(shape1), len(shape2)
544 max_len = len(output_shape)
545 padded_shape1 = (1,) * (max_len - len1) + tuple(shape1)
546 padded_shape2 = (1,) * (max_len - len2) + tuple(shape2)
548 merged_map = []
549 for i, (dim1, dim2, out_dim) in enumerate(zip(padded_shape1, padded_shape2, output_shape)):
550 m1, m2 = map1[i], map2[i]
552 m1_axis = self._normalize_tensor_map_element(m1)
553 m2_axis = self._normalize_tensor_map_element(m2)
555 m1_axis_for_compare = frozenset(m1_axis)
556 m2_axis_for_compare = frozenset(m2_axis)
558 m1_is_sharded = bool(m1_axis)
559 m2_is_sharded = bool(m2_axis)
561 if not m1_is_sharded and not m2_is_sharded:
562 merged_map.append(-1)
564 elif not m1_is_sharded:
565 if dim2 == 1 and out_dim > 1:
566 raise ValueError(
567 f"For {self.op_name}, dimension {i} of second input has size 1 "
568 f"but is sharded on device axis {m2_axis}. "
569 f"Broadcasting dimension cannot be sharded."
570 )
571 merged_map.append(self._denormalize_tensor_map_element(m2_axis))
573 elif not m2_is_sharded:
574 if dim1 == 1 and out_dim > 1:
575 raise ValueError(
576 f"For {self.op_name}, dimension {i} of first input has size 1 "
577 f"but is sharded on device axis {m1_axis}. "
578 f"Broadcasting dimension cannot be sharded."
579 )
580 merged_map.append(self._denormalize_tensor_map_element(m1_axis))
582 else:
583 if m1_axis_for_compare != m2_axis_for_compare:
584 raise ValueError(
585 f"For {self.op_name}, inputs should have same sharding pattern, "
586 f"but got confilcting sharding at dimension {i}, "
587 f"input1 shaded on {m1_axis} and input2 shaded on {m2_axis}."
588 )
590 if (dim1 == 1 or dim2 == 1) and dim1 != dim2:
591 raise ValueError(
592 f"For {self.op_name}, dimension {i} is broadcast from size 1 "
593 f"to {out_dim} but is sharded on device axis {m1_axis}. "
594 f"Broadcasting dimension cannot be sharded."
595 )
597 merged_map.append(self._denormalize_tensor_map_element(m1_axis))
599 return tuple(merged_map)
601 def _create_output_layout(self, base_layout, output_tensor_map, partial_list=None):
602 """
603 Create output layout based on input layout.
605 Args:
606 base_layout: Base layout (usually from the first input)
607 output_tensor_map (tuple): Tensor_map for the output
608 partial_list (list): Partial status list for the output
610 Returns:
611 Layout: New Layout object with updated tensor_map and alias_tensor_map
612 """
613 new_layout = copy.deepcopy(base_layout)
614 new_layout.set_tensor_map(output_tensor_map)
616 alias_tensor_map = []
617 for tensor_dim in output_tensor_map:
618 if tensor_dim == -1:
619 alias_tensor_map.append("None")
620 elif isinstance(tensor_dim, tuple):
621 alias_tuple = tuple(
622 base_layout.alias_name[len(base_layout.alias_name) - 1 - dim]
623 for dim in tensor_dim
624 if dim != -1
625 )
626 alias_tensor_map.append(alias_tuple if alias_tuple else "None")
627 else:
628 alias_tensor_map.append(
629 base_layout.alias_name[len(base_layout.alias_name) - 1 - tensor_dim]
630 )
632 new_layout.set_alias_tensor_map(tuple(alias_tensor_map))
634 # Set partial status if provided
635 if partial_list:
636 for i, partial_op in enumerate(partial_list):
637 if partial_op is not None and i < len(new_layout.alias_name):
638 new_layout.set_partial_by_dev_axis(new_layout.alias_name[i], partial_op)
640 return new_layout
643class ElementWiseWithPartialDistributedOp(ElementWiseDistributedOp):
644 """
645 Base class for elementwise operations that support partial status propagation.
646 """
647 def __init__(self, op_name):
648 super().__init__(op_name)
649 self._allow_partial_inputs = True
652class AddDistributedOp(ElementWiseWithPartialDistributedOp):
653 """
654 Distributed implementation for Add operator.
656 This operator supports partial status propagation from inputs to output,
657 which is useful for operations like gradient accumulation where partial
658 results need to be preserved through the computation graph.
659 """
661 @staticmethod
662 def _canonicalize_binary_operands(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
663 """Move semantic ``input`` and ``other`` operands before option values."""
664 local_kwargs = dict(kwargs)
665 if args:
666 if "input" in local_kwargs:
667 raise ValueError("Binary elementwise input was provided both positionally and by keyword.")
668 lhs = args[0]
669 trailing_args = args[1:]
670 elif "input" in local_kwargs:
671 lhs = local_kwargs.pop("input")
672 trailing_args = ()
673 else:
674 return args, local_kwargs
676 if trailing_args:
677 if "other" in local_kwargs:
678 raise ValueError("Binary elementwise other was provided both positionally and by keyword.")
679 rhs = trailing_args[0]
680 trailing_args = trailing_args[1:]
681 elif "other" in local_kwargs:
682 rhs = local_kwargs.pop("other")
683 else:
684 return args, local_kwargs
685 return (lhs, rhs, *trailing_args), local_kwargs
687 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
688 """Canonicalize binary operands before local unwrapping and layout caching.
690 Args:
691 args: Positional backend arguments.
692 kwargs: Keyword backend arguments, potentially including
693 ``input``, ``other``, and options such as ``alpha``.
695 Returns:
696 Local arguments, local keyword arguments, and layout cache values
697 with the semantic left and right operands in the first two slots.
698 """
699 canonical_args, canonical_kwargs = self._canonicalize_binary_operands(args, kwargs)
700 return super().preprocess(canonical_args, canonical_kwargs)
702 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
703 """
704 Infer output layout for Add operator.
706 Rules:
707 1. Follow element-wise broadcasting and sharding merge rules.
708 2. Partial status is allowed and propagated.
709 3. Propagated Partial status must be "sum" or None.
711 Args:
712 cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes].
714 Returns:
715 tuple: ((output_layout,), None)
717 Raises:
718 ValueError: If propagated Partial status is not "sum" or None.
719 """
720 layouts = tuple(cache_values[:-1])
721 if self.op_name in _INPLACE_ELEMENTWISE_OPS and len(layouts) >= 2:
722 mesh_ndim = max(
723 (len(layout.partial) for layout in layouts[:2] if layout is not None),
724 default=0,
725 )
726 partial_signatures = tuple(
727 _partial_signature(layout, mesh_ndim)
728 for layout in layouts[:2]
729 )
730 if partial_signatures[0] != partial_signatures[1]:
731 raise ValueError(
732 f"For {self.op_name}, input Partial placements should be identical "
733 f"for in-place execution, but got {partial_signatures}."
734 )
736 infer_result = super().infer_layout(cache_values)
737 if infer_result is None:
738 return infer_result
740 output_layout = infer_result[0][0]
741 for i, partial_type in enumerate(output_layout.partial):
742 if partial_type is not None and partial_type != "sum":
743 raise ValueError(
744 f"For {self.op_name}, inputs partial status should be 'sum' or None, "
745 f"but got {partial_type} at index {i}."
746 )
748 return infer_result
750 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
751 cache_values: list) -> Optional[Callable]:
752 """Build a rank-local implementation that avoids repeated values.
754 Args:
755 func: Local backend add or subtract implementation.
756 infer_result: Inferred output layout and auxiliary result.
757 cache_values: Input layouts followed by input shapes.
759 Returns:
760 A contribution-aware local implementation, or ``None`` when both
761 inputs already have identical Partial placements.
762 """
763 layouts = tuple(cache_values[:-1])
764 x1_layout = layouts[0]
765 x2_layout = layouts[1]
766 output_layout = infer_result[0][0]
767 output_mesh_ndim = len(output_layout.partial)
768 x1_partial = _partial_signature(x1_layout, output_mesh_ndim)
769 x2_partial = _partial_signature(x2_layout, output_mesh_ndim)
771 if x1_partial == x2_partial:
772 return None
774 x1_contributes = _contributes_to_partial_output(x1_layout, output_layout)
775 x2_contributes = _contributes_to_partial_output(x2_layout, output_layout)
777 # ``*args``/``**kwargs`` forward add-specific extras such as ``alpha``.
778 # A strict zero retains a valid zero-gradient autograd edge.
779 def _expand_impl(x1, x2, *args, **kwargs):
780 local_x1 = x1 if x1_contributes else _zero_contribution(x1)
781 local_x2 = x2 if x2_contributes else _zero_contribution(x2)
782 return func(local_x1, local_x2, *args, **kwargs)
784 return _expand_impl
787class SubDistributedOp(AddDistributedOp):
788 """Distributed subtraction with correct replicated-minus-partial signs.
790 When the first operand is replicated and the second is Partial(sum), only
791 the first Partial rank may contribute the replicated value. Every other
792 rank must contribute ``0 - x2_local`` so reducing the output reconstructs
793 ``x1 - sum(x2_local)``.
794 """