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 

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24# pylint: disable=unused-argument 

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

26 """ 

27 Normalize arguments for Concat operator. 

28 """ 

29 return (tensors, dim), {} 

30 

31 

32class ConcatDistributedOp(DistributedOp): 

33 """Distributed implementation for Concat.""" 

34 

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

36 """ 

37 Preprocess arguments for Concat operator. 

38 

39 Args: 

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

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

42 

43 Returns: 

44 tuple: (local_args, local_kwargs, cache_values) 

45 """ 

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

47 tensors = args[0] 

48 dim = args[1] 

49 

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

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

52 

53 local_args = (local_tensors, dim) 

54 local_kwargs = {} 

55 cache_values = layouts + [dim] 

56 return local_args, local_kwargs, cache_values 

57 

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

59 """ 

60 Infer output layouts for Concat operator. 

61 

62 Rules: 

63 1. Inputs must not have Partial status. 

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

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

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

67 5. The concatenation dimension must not be sharded. 

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

69 

70 Args: 

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

72 use None as their layout sentinel. 

73 

74 Returns: 

75 tuple: ((output_layout,), None) 

76 

77 Raises: 

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

79 or the concatenation dimension is sharded. 

80 """ 

81 layouts = cache_values[:-1] 

82 dim = cache_values[-1] 

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

84 

85 if not valid_layouts: 

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

87 

88 self._check_partial_inputs(valid_layouts) 

89 

90 base_layout = valid_layouts[0] 

91 

92 for layout in valid_layouts: 

93 if layout != base_layout: 

94 raise ValueError( 

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

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

97 ) 

98 

99 if not isinstance(dim, int): 

100 raise ValueError( 

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

102 ) 

103 

104 ndim = len(base_layout.alias_tensor_map) 

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

106 raise ValueError( 

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

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

109 ) 

110 

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

112 

113 mapping = base_layout.alias_tensor_map[actual_dim] 

114 if mapping != "None": 

115 raise ValueError( 

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

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

118 ) 

119 

120 return ((base_layout,), None)