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

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

17""" 

18 

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_masked_scatter_args(input_tensor, mask, source): 

25 return (input_tensor, mask, source), {} 

26 

27 

28class MaskedScatterDistributedOp(DistributedOp): 

29 """Distributed implementation for torch.Tensor.masked_scatter.""" 

30 

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

32 """ 

33 Preprocess arguments for MaskedScatter operator. 

34 

35 torch.masked_scatter(input, mask, source) takes three positional tensor inputs 

36 with no keyword-only parameters. 

37 

38 Args: 

39 args (tuple): Positional arguments (input, mask, source). 

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

41 

42 Returns: 

43 tuple: (local_args, local_kwargs, cache_values) 

44 """ 

45 args, kwargs = _normalize_masked_scatter_args(*args, **kwargs) 

46 input_tensor, mask, source = args[0], args[1], args[2] 

47 local_args = ( 

48 input_tensor.to_local(), 

49 mask.to_local(), 

50 source.to_local(), 

51 ) 

52 local_kwargs = {} 

53 cache_values = [ 

54 input_tensor.layout, 

55 mask.layout, 

56 source.layout, 

57 ] 

58 return local_args, local_kwargs, cache_values 

59 

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

61 """ 

62 Infer output layout for torch.Tensor.masked_scatter. 

63 

64 PyTorch semantics: 

65 masked_scatter_(mask, source) 

66 Copies elements from source into self tensor at positions where the mask is True. 

67 Elements from source are taken in order. 

68 

69 Distributed restrictions: 

70 Because `masked_scatter` consumes elements from `source` sequentially based on 

71 the flattened index of `True` values in `mask`, sharding the input or mask 

72 would require a global prefix sum (scan) to determine the correct offset 

73 in `source` for each rank. Without this communication overhead, correct 

74 behavior cannot be guaranteed on sharded tensors. 

75 

76 Therefore, this implementation enforces that all inputs (input, mask, source) 

77 must be fully Replicated (Unsharded). 

78 

79 Rules: 

80 1. Inputs must not have Partial status. 

81 2. All inputs (input, mask, source) must be fully Replicated — no dimension may be sharded. 

82 3. Output layout follows the input layout. 

83 

84 Args: 

85 cache_values (list): [input_layout, mask_layout, source_layout] 

86 

87 Returns: 

88 tuple: ((output_layout,), None) 

89 

90 Raises: 

91 ValueError: If any input tensor is sharded. 

92 """ 

93 input_layout = cache_values[0] 

94 mask_layout = cache_values[1] 

95 source_layout = cache_values[2] 

96 

97 layouts = [input_layout, mask_layout, source_layout] 

98 if not self._allow_partial_inputs: 

99 self._check_partial_inputs(layouts) 

100 

101 # Check strict replication for all involved distributed tensors 

102 for i, layout in enumerate(layouts): 

103 if layout is None: 

104 continue 

105 

106 # Use alias_tensor_map for sharding checks — it handles StridedShard tuple mappings. 

107 # "None" means Replicated; anything else indicates sharding. 

108 for dim_alias in layout.alias_tensor_map: 

109 if dim_alias != "None": 

110 raise ValueError( 

111 f"For {self.op_name}, input {i} is sharded (Layout: {layout}). " 

112 f"masked_scatter currently only supports fully Replicated (Unsharded) tensors " 

113 f"due to sequential dependency on source elements." 

114 ) 

115 

116 # Output layout follows input layout (which we verified is Replicated/None) 

117 return ((input_layout,), None)