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

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

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25# pylint: disable=unused-argument 

26def _normalize_concat_args(tensors, dim=0, **kwargs): 

27 """ 

28 Normalize arguments for Concat operator. 

29 """ 

30 return (tensors, dim), {} 

31 

32 

33class ConcatDistributedOp(DistributedOp): 

34 """Distributed implementation for Concat.""" 

35 

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

37 """ 

38 Preprocess arguments for Concat operator. 

39 

40 Args: 

41 args (tuple): Input arguments, first element is the input tensor sequence. 

42 kwargs (dict): Keyword arguments, may contain dim. 

43 

44 Returns: 

45 tuple: (local_args, local_kwargs, cache_values) 

46 """ 

47 args, _ = _normalize_concat_args(*args, **kwargs) 

48 tensors = args[0] 

49 dim = args[1] 

50 

51 local_tensors = tuple(t.to_local() if hasattr(t, "to_local") else t for t in tensors) 

52 layouts = [getattr(t, "layout", None) for t in tensors] 

53 

54 local_args = (local_tensors, dim) 

55 local_kwargs = {} 

56 cache_values = layouts + [dim] 

57 return local_args, local_kwargs, cache_values 

58 

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

60 """ 

61 Infer output layouts for Concat operator. 

62 

63 Rules: 

64 1. Inputs must not have Partial status. 

65 2. At least one input must be a DTensor. 

66 3. All input DTensors must have the same layout. 

67 4. dim must be an integer within the valid range [-ndim, ndim-1]. 

68 5. The concatenation dimension must not be sharded. 

69 6. Output layout is identical to the input layout. 

70 

71 Args: 

72 cache_values (list): [input_layout, ..., dim] where non-DTensor inputs 

73 use None as their layout sentinel. 

74 

75 Returns: 

76 tuple: ((output_layout,), None) 

77 

78 Raises: 

79 ValueError: If inputs are invalid, layouts mismatch, dim is out of range, 

80 or the concatenation dimension is sharded. 

81 """ 

82 layouts = cache_values[:-1] 

83 dim = cache_values[-1] 

84 valid_layouts = [layout for layout in layouts if layout is not None] 

85 

86 if not valid_layouts: 

87 raise ValueError(f"For {self.op_name}, cat requires at least one input DTensor.") 

88 

89 self._check_partial_inputs(valid_layouts) 

90 

91 base_layout = valid_layouts[0] 

92 

93 for layout in valid_layouts: 

94 if layout != base_layout: 

95 raise ValueError( 

96 f"For {self.op_name}, All input tensors must have the same layout. " 

97 f"Expected layout: {base_layout}, Mismatched layout: {layout}" 

98 ) 

99 

100 if not isinstance(dim, int): 

101 raise ValueError( 

102 f"For {self.op_name}, dimension should be int, but got {type(dim)}" 

103 ) 

104 

105 ndim = len(base_layout.alias_tensor_map) 

106 if dim < -ndim or dim >= ndim: 

107 raise ValueError( 

108 f"For {self.op_name}, dimension out of range " 

109 f"(expected to be in range of [{-ndim}, {ndim - 1}], but got {dim})" 

110 ) 

111 

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

113 

114 mapping = base_layout.alias_tensor_map[actual_dim] 

115 if mapping != "None": 

116 raise ValueError( 

117 f"For {self.op_name}, Concatenation along a sharded dimension " 

118 f"(dim={dim}, normalized_dim={actual_dim}) is not supported." 

119 ) 

120 

121 return ((copy.deepcopy(base_layout),), None)