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

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

17""" 

18from copy import deepcopy 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.layout import Layout 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_squeeze_args(x, axis=None): 

26 return (x, axis), {} 

27 

28 

29class SqueezeDistributedOp(DistributedOp): 

30 """Distributed implementation for Squeeze operator.""" 

31 

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

33 args, _ = _normalize_squeeze_args(*args, **kwargs) 

34 input_tensor, axis = args[0], args[1] 

35 

36 if axis is None: 

37 local_args = (input_tensor.to_local(),) 

38 else: 

39 local_args = (input_tensor.to_local(), axis) 

40 

41 input_shape = getattr(input_tensor, "shape", None) 

42 cache_values = [input_tensor.layout, axis, input_shape] 

43 return local_args, {}, cache_values 

44 

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

46 """ 

47 Infer output layout for Squeeze operator. 

48 

49 Rules: 

50 1. Input must not have Partial status. 

51 2. input_shape rank must match input layout rank. 

52 3. If axis is None, only size-1 replicated dimensions are removed. 

53 4. If axis is specified, every axis must be in range, have size 1, and 

54 must not be sharded. 

55 5. Output layout removes the squeezed dimensions and preserves the 

56 remaining input dimension mappings. 

57 

58 Args: 

59 cache_values (list): [input_layout, axis, input_shape] 

60 

61 Returns: 

62 tuple: ((output_layout,), None) 

63 

64 Raises: 

65 ValueError: If input has Partial status, input_shape is missing, axis 

66 is invalid, or a requested squeeze dimension is sharded. 

67 """ 

68 if not cache_values: 

69 raise ValueError( 

70 f"For {self.op_name}, cache_values should contain input layout, " 

71 f"but got empty cache_values." 

72 ) 

73 

74 x_layout = cache_values[0] 

75 if not self._allow_partial_inputs: 

76 self._check_partial_inputs([x_layout]) 

77 

78 if x_layout.mesh_shape is None: 

79 raise ValueError( 

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

81 f"but got None." 

82 ) 

83 

84 axis = cache_values[1] if len(cache_values) > 1 else None 

85 input_shape = cache_values[2] if len(cache_values) > 2 else None 

86 if input_shape is None: 

87 raise ValueError( 

88 f"For {self.op_name}, input_shape should be provided in cache_values, " 

89 f"but got None." 

90 ) 

91 if not isinstance(input_shape, (list, tuple)): 

92 raise ValueError( 

93 f"For {self.op_name}, input_shape should be list or tuple, " 

94 f"but got {type(input_shape)}." 

95 ) 

96 

97 output_layout = self._compute_squeeze_layout(x_layout, axis, input_shape) 

98 return ((output_layout,), None) 

99 

100 def _compute_squeeze_layout(self, x_layout, axis, input_shape): 

101 """Compute the squeezed layout.""" 

102 # Handle scalar case 

103 if not input_shape: 

104 return self._handle_scalar_case(x_layout, axis) 

105 

106 # Validate input_shape matches layout rank 

107 self._validate_input_shape(x_layout, input_shape) 

108 

109 # Find dimensions to squeeze 

110 dims_to_squeeze = self._get_dims_to_squeeze(x_layout, axis, input_shape) 

111 

112 # Create output layout 

113 return self._create_output_layout(x_layout, dims_to_squeeze) 

114 

115 def _handle_scalar_case(self, x_layout, axis): 

116 """Handle scalar input case.""" 

117 if axis is not None and axis != [] and axis != (): 

118 raise ValueError( 

119 f"For {self.op_name}, axis should be None for scalar input, " 

120 f"but got {axis}." 

121 ) 

122 

123 return deepcopy(x_layout) 

124 

125 def _validate_input_shape(self, x_layout, input_shape): 

126 """Validate that input shape matches layout rank.""" 

127 x_map = list(x_layout.alias_tensor_map) 

128 in_rank = len(x_map) 

129 

130 if len(input_shape) != in_rank: 

131 raise ValueError( 

132 f"For {self.op_name}, input shape rank should match layout rank, " 

133 f"but got {len(input_shape)} and {in_rank}." 

134 ) 

135 

136 def _get_dims_to_squeeze(self, x_layout, axis, input_shape): 

137 """Get list of dimensions to squeeze.""" 

138 x_map = list(x_layout.alias_tensor_map) 

139 in_rank = len(x_map) 

140 

141 if axis is None: 

142 return self._get_all_squeezable_dims(x_map, input_shape) 

143 return self._get_specified_dims_to_squeeze(x_map, axis, input_shape, in_rank) 

144 

145 def _get_all_squeezable_dims(self, x_map, input_shape): 

146 """Get all squeezable dimensions when axis is None.""" 

147 dims_to_squeeze = [] 

148 for i, shape in enumerate(input_shape): 

149 if shape == 1 and x_map[i] == "None": 

150 dims_to_squeeze.append(i) 

151 return dims_to_squeeze 

152 

153 def _get_specified_dims_to_squeeze(self, x_map, axis, input_shape, in_rank): 

154 """Get dimensions to squeeze when axis is specified.""" 

155 # Convert axis to list if it's a single integer 

156 if isinstance(axis, int): 

157 axis = [axis] 

158 elif isinstance(axis, tuple): 

159 axis = list(axis) 

160 elif not isinstance(axis, list): 

161 raise ValueError( 

162 f"For {self.op_name}, axis should be int, list or tuple, " 

163 f"but got {type(axis)}." 

164 ) 

165 

166 if not all(isinstance(ax, int) for ax in axis): 

167 raise ValueError( 

168 f"For {self.op_name}, every axis value should be int, " 

169 f"but got {axis}." 

170 ) 

171 

172 # Convert negative indices to positive 

173 axis = [ax if ax >= 0 else ax + in_rank for ax in axis] 

174 

175 # Validate axis range 

176 self._validate_axis_range(axis, in_rank) 

177 

178 # Check all specified axes 

179 for ax in axis: 

180 self._validate_axis_for_squeeze(x_map, input_shape, ax) 

181 

182 # Return sorted unique axes 

183 return sorted(set(axis)) 

184 

185 def _validate_axis_range(self, axis, in_rank): 

186 """Validate axis values are within range.""" 

187 for ax in axis: 

188 if ax < 0 or ax >= in_rank: 

189 raise ValueError( 

190 f"For {self.op_name}, axis should be in range [{-in_rank}, {in_rank-1}], " 

191 f"but got {ax}." 

192 ) 

193 

194 def _validate_axis_for_squeeze(self, x_map, input_shape, ax): 

195 """Validate a specific axis can be squeezed.""" 

196 # Check shape == 1 

197 if input_shape[ax] != 1: 

198 raise ValueError( 

199 f"For {self.op_name}, dimension should have size 1, " 

200 f"but got shape {input_shape[ax]} at dimension {ax}." 

201 ) 

202 

203 # Check mapping is "None" (not distributed) 

204 if x_map[ax] != "None": 

205 raise ValueError( 

206 f"For {self.op_name}, dimension should not be distributed, " 

207 f"but got dimension {ax} mapped to device axis {x_map[ax]}." 

208 ) 

209 

210 def _create_output_layout(self, x_layout, dims_to_squeeze): 

211 """Create output layout after squeezing dimensions.""" 

212 if not dims_to_squeeze: 

213 return deepcopy(x_layout) 

214 

215 # Get current alias tensor map 

216 x_map = list(x_layout.alias_tensor_map) 

217 

218 # Sort in descending order for safe removal 

219 dims_to_squeeze = sorted(set(dims_to_squeeze), reverse=True) 

220 

221 # Remove specified dimensions 

222 for dim in dims_to_squeeze: 

223 del x_map[dim] 

224 

225 new_map = x_map 

226 

227 # Create output layout with new mapping 

228 output_layout = Layout( 

229 mesh_shape=x_layout.mesh_shape, 

230 alias_name=x_layout.alias_name, 

231 rank_list=x_layout.rank_list 

232 ) 

233 

234 if new_map: 

235 output_layout = output_layout(*new_map) 

236 else: 

237 # For scalar result 

238 output_layout = output_layout() 

239 

240 # Copy partial operations from input layout 

241 self._copy_partial_operations(x_layout, output_layout, new_map) 

242 

243 return output_layout 

244 

245 def _copy_partial_operations(self, x_layout, output_layout, new_map): 

246 """Copy partial operations from input to output layout.""" 

247 for i, partial_op in enumerate(x_layout.partial): 

248 if partial_op is not None: 

249 dev_axis_name = x_layout.alias_name[i] 

250 # Check if this device axis is still used in the output 

251 if dev_axis_name in new_map: 

252 output_layout.set_partial_by_dev_axis(dev_axis_name, partial_op)