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

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 

16""" 

17Distributed implementation for RepeatInterleave operator. 

18""" 

19import copy 

20from typing import Tuple 

21 

22from hyper_parallel.core.dtensor.layout import Layout 

23from .parallel_ops import DistributedOp 

24 

25 

26def _normalize_repeat_interleave_args(x, repeats, dim=None, *, output_size=None): 

27 """Normalize positional and keyword arguments into a canonical form. 

28 

29 Args: 

30 x: Input tensor. 

31 repeats: Number of repetitions (int or Tensor). 

32 dim: Dimension along which to repeat values. None means flatten first. 

33 output_size: Total output size (keyword-only in torch). 

34 

35 Returns: 

36 tuple: (positional_args_tuple, kwargs_dict) 

37 """ 

38 kwargs = {} 

39 if output_size is not None: 

40 kwargs['output_size'] = output_size 

41 return (x, repeats, dim), kwargs 

42 

43 

44class RepeatInterleaveDistributedOp(DistributedOp): 

45 """Distributed implementation for torch.repeat_interleave. 

46 

47 Sharding constraints: 

48 - When dim is specified: the repeat dimension must be replicated. 

49 - When dim is None (flatten mode): only the first dimension may be sharded. 

50 

51 Output layout: 

52 - When dim is specified: same as input layout. 

53 - When dim is None: 1-D layout preserving sharding on dim 0. 

54 """ 

55 

56 @staticmethod 

57 def _validate_input_layouts(input_layout, dim, op_name: str) -> None: 

58 """Validate sharding constraints for repeat_interleave. 

59 

60 Args: 

61 input_layout: Layout of the input tensor. 

62 dim: The repeat dimension, or None for flatten mode. 

63 op_name: Operator name used in error messages. 

64 

65 Raises: 

66 ValueError: If dim is out of range or the repeat dimension is sharded. 

67 """ 

68 in_tensor_map = input_layout.alias_tensor_map 

69 ndim = len(in_tensor_map) 

70 

71 if dim is not None: 

72 actual_dim = dim if dim >= 0 else ndim + dim 

73 if not 0 <= actual_dim < ndim: 

74 raise ValueError( 

75 f"For {op_name}, dimension should be in [0, {ndim}), but got {dim}." 

76 ) 

77 mapping = in_tensor_map[actual_dim] 

78 if isinstance(mapping, (list, tuple)): 

79 is_sharded = any(axis != "None" for axis in mapping) 

80 else: 

81 is_sharded = mapping != "None" 

82 if is_sharded: 

83 raise ValueError( 

84 f"For {op_name}, the repeat dimension should be replicated, " 

85 f"but got dim={dim} mapped to {mapping}." 

86 ) 

87 

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

89 """Preprocess arguments for RepeatInterleave operator. 

90 

91 Extracts local tensors from DTensor inputs and builds cache_values 

92 containing only the information needed for layout inference. 

93 

94 Args: 

95 args: Positional arguments (input, repeats, dim). 

96 kwargs: Keyword arguments (output_size). 

97 

98 Returns: 

99 tuple: (local_args, local_kwargs, cache_values) 

100 """ 

101 args, kwargs = _normalize_repeat_interleave_args(*args, **kwargs) 

102 input_tensor = args[0] 

103 repeats = args[1] 

104 dim = args[2] 

105 output_size = kwargs.get('output_size', None) 

106 

107 local_input = input_tensor.to_local() 

108 

109 # repeats can be int or Tensor; handle DTensor defensively 

110 if hasattr(repeats, 'to_local'): 

111 local_repeats = repeats.to_local() 

112 else: 

113 local_repeats = repeats 

114 

115 local_args = (local_input, local_repeats, dim) 

116 local_kwargs = {} 

117 if output_size is not None: 

118 local_kwargs['output_size'] = output_size 

119 

120 cache_values = [input_tensor.layout, dim] 

121 return local_args, local_kwargs, cache_values 

122 

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

124 """Infer output layout for RepeatInterleave operator. 

125 

126 Rules: 

127 1. Input must not have Partial status. 

128 2. When dim is specified: the repeat dimension must be replicated. 

129 Output layout is a deep copy of the input layout. 

130 3. When dim is None (flatten mode): only dim 0 may be sharded. 

131 Output is a 1-D layout preserving the mesh axis on dim 0. 

132 

133 Args: 

134 cache_values: [input_layout, dim] where dim is the repeat dimension 

135 or None for flatten mode. 

136 

137 Returns: 

138 tuple: ((output_layout,), None) 

139 

140 Raises: 

141 ValueError: If input has Partial status, dim is out of range, 

142 the repeat dimension is sharded, or flatten is attempted on 

143 a tensor sharded on a non-first dimension. 

144 """ 

145 input_layout = cache_values[0] 

146 dim = cache_values[1] 

147 

148 self._check_partial_inputs([input_layout]) 

149 self._validate_input_layouts(input_layout, dim, self.op_name) 

150 

151 in_tensor_map = input_layout.alias_tensor_map 

152 

153 if dim is None: 

154 # Flatten mode: output is 1-D. 

155 sharded_dims = [] 

156 for i, shard in enumerate(in_tensor_map): 

157 if isinstance(shard, (list, tuple)): 

158 is_sharded = any(axis != "None" for axis in shard) 

159 else: 

160 is_sharded = shard != "None" 

161 if is_sharded: 

162 sharded_dims.append(i) 

163 if not sharded_dims: 

164 output_tensor_map = ("None",) 

165 elif sharded_dims == [0]: 

166 output_tensor_map = (in_tensor_map[0],) 

167 else: 

168 raise ValueError( 

169 f"For {self.op_name}, sharded dims in flatten mode should be [] or [0], " 

170 f"but got {sharded_dims}." 

171 ) 

172 

173 output_layout = Layout( 

174 mesh_shape=input_layout.mesh_shape, 

175 alias_name=input_layout.alias_name, 

176 rank_list=input_layout.rank_list 

177 ) 

178 return (output_layout(*output_tensor_map),), None 

179 

180 # dim specified: output layout = deep copy of input layout 

181 return (copy.deepcopy(input_layout),), None