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

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

17""" 

18 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.layout import Layout 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_repeat_args(x, *sizes): 

26 return (x,) + sizes, {} 

27 

28 

29class RepeatDistributedOp(DistributedOp): 

30 """Distributed implementation for torch.Tensor.repeat.""" 

31 

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

33 """ 

34 Preprocess arguments for torch.Tensor.repeat. 

35 

36 Args: 

37 args (tuple): Input tensor followed by repeat sizes. 

38 kwargs (dict): Keyword arguments. 

39 

40 Returns: 

41 tuple: (local_args, local_kwargs, cache_values) 

42 """ 

43 args, _ = _normalize_repeat_args(*args, **kwargs) 

44 input_tensor = args[0] 

45 sizes = args[1:] 

46 

47 local_args = (input_tensor.to_local(),) + sizes 

48 local_kwargs = {} 

49 cache_values = [input_tensor.layout, sizes] 

50 return local_args, local_kwargs, cache_values 

51 

52 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221 

53 """ 

54 Infer output layout for torch.Tensor.repeat. 

55 

56 PyTorch semantics: 

57 - Repeats this tensor along the specified dimensions. 

58 - If the number of repeat dimensions is larger than the tensor dimensions, 

59 the tensor is implicitly unsqueezed at the front. 

60 - The number of repeat dimensions cannot be smaller than the tensor dimensions. 

61 - Dimensions being repeated (>1 or 0) MUST be unsharded. 

62 

63 Rules: 

64 1. Input must not have Partial status. 

65 2. Repeat sizes must be provided. 

66 3. New prepended dimensions are unsharded. 

67 4. Existing dimensions with repeat size 1 preserve the input sharding. 

68 5. Existing dimensions with repeat size other than 1 must not be sharded. 

69 

70 Args: 

71 cache_values (list): [input_layout, repeat_sizes]. 

72 

73 Returns: 

74 tuple: ((output_layout,), None) 

75 

76 Raises: 

77 ValueError: If input layout is missing, input has Partial status, repeat sizes 

78 are missing, or a sharded dimension is repeated. 

79 """ 

80 if not cache_values or cache_values[0] is None: 

81 raise ValueError( 

82 f"For {self.op_name}, repeat requires a valid input tensor layout." 

83 ) 

84 

85 input_layout = cache_values[0] 

86 self._check_partial_inputs([input_layout]) 

87 

88 in_tensor_map = input_layout.tensor_map 

89 in_alias_map = input_layout.alias_tensor_map 

90 input_ndim = len(in_tensor_map) 

91 

92 if len(cache_values) < 2 or not cache_values[1]: 

93 raise ValueError( 

94 f"For {self.op_name}, repeat requires repeat sizes in cache_values." 

95 ) 

96 

97 # Robustly handle sizes unpacking (e.g., if args are packed as a single tuple) 

98 repeat_sizes = cache_values[1] 

99 if len(repeat_sizes) == 1 and isinstance(repeat_sizes[0], (tuple, list)): 

100 flat_args = repeat_sizes[0] 

101 else: 

102 flat_args = repeat_sizes 

103 

104 # Normalize repeat sizes to tuple of ints 

105 repeats = [] 

106 for arg in flat_args: 

107 if not isinstance(arg, int): 

108 arg = int(arg) 

109 repeats.append(arg) 

110 repeats = tuple(repeats) 

111 output_ndim = len(repeats) 

112 

113 num_new_dims = output_ndim - input_ndim 

114 output_map = [] 

115 

116 # Rule 1: New prepended dimensions are always unsharded 

117 for _ in range(num_new_dims): 

118 output_map.append("None") 

119 

120 # Rule 2: Process existing dimensions 

121 for i in range(input_ndim): 

122 repeat_idx = num_new_dims + i 

123 repeat_times = repeats[repeat_idx] 

124 

125 if repeat_times == 1: 

126 # If the dimension is not repeated, keep the original sharding 

127 output_map.append(in_alias_map[i]) 

128 else: 

129 # If the dimension is repeated (or zeroed), it cannot be currently sharded 

130 mapping = in_alias_map[i] 

131 if mapping != "None": 

132 raise ValueError( 

133 f"For {self.op_name}, Cannot repeat dimension {i} which is sharded. " 

134 f"Please redistribute (unshard) the tensor along this dimension first." 

135 ) 

136 # Repeated dimension remains unsharded in output 

137 output_map.append("None") 

138 

139 # Construct output layout mapping 

140 mesh_shape = input_layout.mesh_shape 

141 alias_name = input_layout.alias_name 

142 rank_list = input_layout.rank_list 

143 

144 # Instantiate new layout 

145 output_layout = Layout( 

146 mesh_shape=mesh_shape, 

147 alias_name=alias_name, 

148 rank_list=rank_list 

149 ) 

150 output_layout = output_layout(*output_map) 

151 

152 return ((output_layout,), None)