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

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 InplaceScatterValue operator. 

17""" 

18 

19from copy import deepcopy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_inplace_scatter_value_args(input_tensor, dim, index, value): 

26 """Normalize InplaceScatterValue arguments to positional args + empty kwargs.""" 

27 return (input_tensor, dim, index, value), {} 

28 

29 

30class InplaceScatterValueDistributedOp(DistributedOp): 

31 """Distributed implementation for InplaceScatterValue operator.""" 

32 

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

34 """ 

35 Preprocess arguments for InplaceScatterValue operator. 

36 

37 Args: 

38 args (tuple): Input arguments (input_tensor, dim, index, value). 

39 kwargs (dict): Keyword arguments (none expected). 

40 

41 Returns: 

42 tuple: (local_args, local_kwargs, cache_values) 

43 """ 

44 args, kwargs = _normalize_inplace_scatter_value_args(*args, **kwargs) 

45 input_tensor, dim, index, value = args 

46 

47 local_args = ( 

48 input_tensor.to_local() if hasattr(input_tensor, '_layout') else input_tensor, 

49 dim, 

50 index.to_local() if hasattr(index, '_layout') else index, 

51 value, 

52 ) 

53 

54 cache_values = [ 

55 input_tensor.layout if hasattr(input_tensor, '_layout') else None, 

56 index.layout if hasattr(index, '_layout') else None, 

57 dim, 

58 ] 

59 local_kwargs = kwargs 

60 return local_args, local_kwargs, cache_values 

61 

62 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

63 """ 

64 Infer output layout for InplaceScatterValue operator. 

65 

66 Rules: 

67 1. Input and index must be DTensors (layout not None). 

68 2. Input must not have Partial status. 

69 3. dim must be an integer and within bounds after normalization. 

70 4. Input and index must have the same number of dimensions. 

71 5. Input and index must use the same sharding on all non-dim axes. 

72 6. The target dim axis must be replicated (unsharded). 

73 7. Output layout is identical to input layout (inplace operation). 

74 

75 Args: 

76 cache_values (list): [input_layout, index_layout, dim] 

77 

78 Returns: 

79 tuple: ((output_layout,), None) 

80 

81 Raises: 

82 ValueError: If any validation rule above is violated. 

83 """ 

84 input_layout, index_layout, dim = cache_values 

85 

86 if not self._allow_partial_inputs: 

87 self._check_partial_inputs([input_layout, index_layout]) 

88 

89 if input_layout is None or not hasattr(input_layout, "tensor_map"): 

90 raise ValueError( 

91 f"For {self.op_name}, input layout should not be None" 

92 ) 

93 

94 if index_layout is None or not hasattr(index_layout, "tensor_map"): 

95 raise ValueError( 

96 f"For {self.op_name}, index must be a DTensor when input is a DTensor" 

97 ) 

98 

99 input_map = input_layout.alias_tensor_map 

100 index_map = index_layout.alias_tensor_map 

101 

102 if not isinstance(dim, int): 

103 raise ValueError( 

104 f"For {self.op_name}, dim should be an integer, but got {type(dim).__name__}" 

105 ) 

106 

107 ndim = len(input_map) 

108 original_dim = dim 

109 if dim < 0: 

110 dim += ndim 

111 if dim < 0 or dim >= ndim: 

112 raise ValueError( 

113 f"For {self.op_name}, dim {original_dim} is out of bounds for tensor with {ndim} dims" 

114 ) 

115 

116 if len(input_map) != len(index_map): 

117 raise ValueError( 

118 f"For {self.op_name}, input and index must have the same number of dimensions, " 

119 f"but got input rank={len(input_map)}, index rank={len(index_map)}" 

120 ) 

121 

122 for axis, (input_axis_map, index_axis_map) in enumerate(zip(input_map, index_map)): 

123 if axis == dim: 

124 continue 

125 if input_axis_map != index_axis_map: 

126 raise ValueError( 

127 f"For {self.op_name}, input and index must use the same sharding " 

128 f"on non-dim axes, " 

129 f"but got mismatch at axis {axis}: " 

130 f"input='{input_axis_map}', index='{index_axis_map}'" 

131 ) 

132 

133 if input_map[dim] != "None": 

134 raise ValueError( 

135 f"For {self.op_name}, scatter along sharded dimension {dim} " 

136 f"is not supported, " 

137 f"dim {dim} is sharded on '{input_map[dim]}'" 

138 ) 

139 

140 return ((deepcopy(input_layout),), None)