Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_expand.py: 77%
106 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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 Expand operator.
17"""
19from typing import Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from .parallel_ops import DistributedOp
25def _normalize_expand_args(input_tensor, *sizes):
26 return (input_tensor, *sizes), {}
29def _normalize_expand_as_args(input_tensor, target_tensor):
30 return (input_tensor, target_tensor), {}
33class ExpandDistributedOp(DistributedOp):
34 """Distributed implementation for torch.Tensor.expand."""
36 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
37 """
38 Preprocess arguments for Expand operator.
40 Args:
41 args (tuple): Input arguments (input_tensor, *sizes).
42 kwargs (dict): Keyword arguments (none for expand).
44 Returns:
45 tuple: (local_args, local_kwargs, cache_values)
46 """
47 args, kwargs = _normalize_expand_args(*args, **kwargs)
48 input_tensor = args[0]
49 sizes = tuple(args[1:])
50 local_args = (input_tensor.to_local(), *sizes)
51 cache_values = [input_tensor.layout, input_tensor.shape, sizes]
52 return local_args, {}, cache_values
54 @staticmethod
55 def _validate_input_layouts(
56 cache_values: list,
57 op_name: str,
58 ) -> Tuple[Layout, tuple, tuple, int]:
59 """Validate all inputs for expand layout inference.
61 Performs type checks, shape validation, sizes validation,
62 and dimension compatibility checks.
64 Rules:
65 1. input_shape and sizes must be tuples of positive ints.
66 2. Cannot reduce dimensions (output_ndim < input_ndim).
67 3. -1 cannot be used for new (prepended) dimensions.
69 Args:
70 cache_values: [input_layout, input_shape, sizes]
71 op_name: Operator name for error messages.
73 Returns:
74 tuple: (input_layout, input_shape, sizes, num_new_dims)
76 Raises:
77 ValueError: If any validation rule is violated.
78 """
79 if not cache_values:
80 raise ValueError(
81 f"For {op_name}, cache_values should contain input layout, "
82 f"but got empty cache_values."
83 )
84 input_layout = cache_values[0]
86 if not isinstance(input_layout, Layout):
87 raise ValueError(
88 f"For {op_name}, input layout should be a Layout, "
89 f"but got {type(input_layout)}."
90 )
92 input_shape = cache_values[1] if len(cache_values) > 1 else None
93 if not isinstance(input_shape, tuple):
94 raise ValueError(
95 f"For {op_name}, input_shape should be a tuple, "
96 f"but got {type(input_shape)}."
97 )
99 sizes = cache_values[2] if len(cache_values) > 2 else None
100 if sizes is None or len(sizes) < 1:
101 raise ValueError(
102 f"For {op_name}, sizes should be a non-empty tuple of ints, "
103 f"but got {sizes}."
104 )
105 for i, sz in enumerate(sizes):
106 if not isinstance(sz, int):
107 raise ValueError(
108 f"For {op_name}, elements in sizes should be int, "
109 f"but got {type(sz)} at position {i}."
110 )
112 in_alias_map = input_layout.alias_tensor_map
113 input_ndim = len(in_alias_map)
114 output_ndim = len(sizes)
115 num_new_dims = output_ndim - input_ndim
117 if len(input_shape) != input_ndim:
118 raise ValueError(
119 f"For {op_name}, input_shape length ({len(input_shape)}) "
120 f"must match input_ndim ({input_ndim})."
121 )
123 if num_new_dims < 0:
124 raise ValueError(
125 f"For {op_name}, cannot reduce dimensions with expand, "
126 f"input has {input_ndim} dims but requested {output_ndim} dims."
127 )
129 # New dimensions cannot use -1
130 for i in range(num_new_dims):
131 if sizes[i] == -1:
132 raise ValueError(
133 f"For {op_name}, cannot use -1 for new dimension at position {i}, "
134 f"sizes should be positive integers."
135 )
137 return input_layout, input_shape, sizes, num_new_dims
139 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
140 """
141 Infer output layout for torch.Tensor.expand.
143 Rules:
144 1. Input must not have Partial status.
145 2. input_shape and sizes must be tuples of positive ints.
146 3. Cannot reduce dimensions (output_ndim < input_ndim).
147 4. -1 cannot be used for new (prepended) dimensions.
148 5. Same-size dimension: preserve original sharding.
149 6. True broadcast (1 → N): the input dimension must be replicated.
150 7. New dimensions are replicated.
152 Args:
153 cache_values (list): [input_layout, input_shape, sizes]
155 Returns:
156 tuple: ((output_layout,), None)
158 Raises:
159 ValueError: If any rule above is violated.
160 """
161 if not self._allow_partial_inputs and cache_values:
162 self._check_partial_inputs([cache_values[0]])
164 input_layout, input_shape, sizes, num_new_dims = self._validate_input_layouts(
165 cache_values, self.op_name
166 )
168 in_alias_map = input_layout.alias_tensor_map
169 input_ndim = len(in_alias_map)
171 output_map = []
173 # New dimensions: always replicated
174 for _ in range(num_new_dims):
175 output_map.append("None")
177 # Process existing dimensions
178 for i in range(input_ndim):
179 output_dim_idx = num_new_dims + i
180 requested_size = sizes[output_dim_idx]
181 input_size = input_shape[i]
183 if requested_size in (-1, input_size):
184 # Dimension unchanged — preserve original sharding
185 output_map.append(in_alias_map[i])
186 elif input_size == 1 and requested_size > 1:
187 # True broadcast: 1 → N — must be replicated
188 if in_alias_map[i] != "None":
189 raise ValueError(
190 f"For {self.op_name}, cannot expand dimension {i} "
191 f"which is sharded "
192 f"(size {input_size} → {requested_size}), "
193 f"got mapping {in_alias_map[i]}."
194 )
195 output_map.append("None")
196 else:
197 raise ValueError(
198 f"For {self.op_name}, cannot expand dimension {i} "
199 f"from size {input_size} to {requested_size}."
200 )
202 output_layout = Layout(
203 mesh_shape=input_layout.mesh_shape,
204 alias_name=input_layout.alias_name,
205 rank_list=input_layout.rank_list
206 )
207 output_layout = output_layout(*output_map)
208 return ((output_layout,), None)
211class ExpandAsDistributedOp(DistributedOp):
212 """Distributed implementation for torch.Tensor.expand_as."""
214 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
215 """
216 Preprocess arguments for ExpandAs operator.
218 Args:
219 args (tuple): Input arguments (input_tensor, target_tensor).
220 kwargs (dict): Keyword arguments (none for expand_as).
222 Returns:
223 tuple: (local_args, local_kwargs, cache_values)
224 """
225 args, kwargs = _normalize_expand_as_args(*args, **kwargs)
226 input_tensor = args[0]
227 target_tensor = args[1]
228 local_args = (input_tensor.to_local(), target_tensor.to_local())
229 cache_values = [input_tensor.layout, input_tensor.shape, target_tensor.shape]
230 return local_args, {}, cache_values
232 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
233 """
234 Infer output layout for expand_as.
236 Rules:
237 1. Input must not have Partial status.
238 2. target_shape must have at least as many dims as input_global_shape.
239 3. Matching-size dimensions preserve input sharding.
240 4. Singleton (size 1) dimensions being expanded to >1 must be unsharded in input.
241 5. Non-singleton dimensions must match exactly.
243 Args:
244 cache_values (list): [input_layout, input_global_shape, target_shape]
246 Returns:
247 tuple: ((output_layout,), None)
249 Raises:
250 ValueError: If any rule above is violated.
251 """
252 if not cache_values:
253 raise ValueError(
254 f"For {self.op_name}, cache_values should contain input layout, "
255 f"but got empty cache_values."
256 )
257 input_layout = cache_values[0]
258 if not self._allow_partial_inputs:
259 self._check_partial_inputs([input_layout])
261 in_alias_map = input_layout.alias_tensor_map
262 input_ndim = len(in_alias_map)
264 input_global_shape = cache_values[1]
265 target_shape = cache_values[2]
267 if not isinstance(target_shape, tuple):
268 raise ValueError(
269 f"For {self.op_name}, target_shape should be tuple, "
270 f"but got {type(target_shape)}."
271 )
272 if not isinstance(input_global_shape, tuple):
273 raise ValueError(
274 f"For {self.op_name}, input_global_shape should be tuple, "
275 f"but got {type(input_global_shape)}."
276 )
278 target_ndim = len(target_shape)
280 if target_ndim < input_ndim:
281 raise ValueError(
282 f"For {self.op_name}, target shape {target_shape} (ndim={target_ndim}) "
283 f"cannot be smaller than input shape {input_global_shape} (ndim={input_ndim})."
284 )
286 # Align dimensions: right-align input to target shape
287 num_leading_implicit = target_ndim - input_ndim
288 aligned_input_shape = (1,) * num_leading_implicit + input_global_shape
289 aligned_tensor_map = ("None",) * num_leading_implicit + in_alias_map
291 # Validate expansion rules and build output tensor_map
292 output_tensor_map = []
293 for i, (in_size, tgt_size, shard_spec) in enumerate(
294 zip(aligned_input_shape, target_shape, aligned_tensor_map)
295 ):
296 if in_size == tgt_size:
297 # Dimension unchanged - preserve sharding pattern
298 output_tensor_map.append(shard_spec)
299 elif in_size == 1 and tgt_size > 1:
300 # Dimension is expanded (broadcast) - must be unsharded
301 if shard_spec != "None":
302 raise ValueError(
303 f"For {self.op_name}, cannot expand sharded dimension {i} "
304 f"which is going to broadcast (global size 1 -> {tgt_size}), "
305 f"got mapping {shard_spec}."
306 )
307 output_tensor_map.append("None")
308 else:
309 raise ValueError(
310 f"For {self.op_name}, cannot expand dimension {i} from size "
311 f"{in_size} to {tgt_size}."
312 )
314 output_layout = Layout(
315 mesh_shape=input_layout.mesh_shape,
316 alias_name=input_layout.alias_name,
317 rank_list=input_layout.rank_list
318 )
319 output_layout = output_layout(*output_tensor_map)
320 return ((output_layout,), None)