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

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

17""" 

18 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.layout import Layout 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_transpose_args(*args): 

26 """Normalize transpose arguments to consistent positional args. 

27 

28 All transpose / permute / TransposeView / TransposeExtView interfaces pass parameters 

29 positionally, so args are returned as-is with empty kwargs. 

30 

31 Args: 

32 *args: Positional arguments from the op call. 

33 

34 Returns: 

35 tuple: (args, {}) — all parameters as positional args, kwargs empty. 

36 """ 

37 return args, {} 

38 

39 

40class TransposeDistributedOp(DistributedOp): 

41 """Distributed implementation for Transpose operator.""" 

42 

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

44 """ 

45 Preprocess arguments for Transpose operator. 

46 

47 Args: 

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

49 kwargs (dict): Keyword arguments (always empty for transpose ops). 

50 

51 Returns: 

52 tuple: (local_args, local_kwargs, cache_values) 

53 """ 

54 args, _ = _normalize_transpose_args(*args) 

55 input_tensor = args[0] 

56 

57 if self.op_name in ("Transpose", "permute", "TransposeView"): 

58 axis = args[1] 

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

60 local_kwargs = {} 

61 cache_values = [input_tensor.layout, axis] 

62 elif self.op_name in ("transpose", "TransposeExtView"): 

63 dim0, dim1 = args[1], args[2] 

64 local_args = (input_tensor.to_local(), dim0, dim1) 

65 local_kwargs = {} 

66 cache_values = [input_tensor.layout, dim0, dim1] 

67 else: 

68 raise ValueError( 

69 f"For TransposeDistributedOp, unsupported op_name: {self.op_name}. " 

70 f"Expected 'Transpose', 'transpose', 'permute', " 

71 f"'TransposeView', or 'TransposeExtView'." 

72 ) 

73 

74 return local_args, local_kwargs, cache_values 

75 

76 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

77 """ 

78 Infer output layout for Transpose operator. 

79 

80 Based on the op_name initialized in the base class, this method switches behavior: 

81 1. op_name == 'Transpose', 'permute' or 'TransposeView': axis-based permutation. 

82 - cache_values: [input_layout, axis] where axis is a tuple of indices. 

83 - Rules: Output tensor_map is permuted according to axis. 

84 2. op_name == 'transpose' or 'TransposeExtView': dim-based swap. 

85 - cache_values: [input_layout, dim0, dim1] where dim0 and dim1 are integers. 

86 - Rules: Output tensor_map has the two dimensions swapped. 

87 

88 Rules: 

89 1. Input must not have Partial status. 

90 2. For axis-based: axis must be a valid permutation of [0, ndim-1]. 

91 3. For dim-based: dim0 and dim1 must be integers within [-ndim, ndim-1]. 

92 4. Output layout inherits mesh info from input, with tensor_map permuted accordingly. 

93 

94 Args: 

95 cache_values (list): [input_layout, ...] where the remaining elements depend on op_name. 

96 

97 Returns: 

98 tuple: ((output_layout,), None) 

99 

100 Raises: 

101 ValueError: If any rule above is violated. 

102 """ 

103 layout = cache_values[0] 

104 if not self._allow_partial_inputs: 

105 self._check_partial_inputs([layout]) 

106 

107 in_tensor_map = layout.alias_tensor_map 

108 ndim = len(in_tensor_map) 

109 

110 if self.op_name in ("Transpose", "permute", "TransposeView"): 

111 axis = cache_values[1] 

112 

113 if not isinstance(axis, (list, tuple)): 

114 raise ValueError( 

115 f"For {self.op_name}, axis should be a list or tuple, " 

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

117 ) 

118 

119 if len(in_tensor_map) != len(axis): 

120 raise ValueError( 

121 f"For {self.op_name}, input tensor shape and permutation " 

122 f"must have the same size. " 

123 f"Got {len(in_tensor_map)} and {len(axis)}." 

124 ) 

125 

126 # check if axis is a permutation 

127 seen = set() 

128 for v in axis: 

129 if not isinstance(v, int): 

130 raise ValueError( 

131 f"For {self.op_name}, axis elements must be integers, " 

132 f"but got {type(v)}." 

133 ) 

134 if v < 0 or v >= ndim or v in seen: 

135 raise ValueError( 

136 f"For {self.op_name}, invalid permutation {axis} for rank {ndim}." 

137 ) 

138 seen.add(v) 

139 

140 out_tensor_map = tuple(in_tensor_map[i] for i in axis) 

141 

142 else: 

143 dim0, dim1 = cache_values[1], cache_values[2] 

144 

145 if not isinstance(dim0, int) or not isinstance(dim1, int): 

146 raise ValueError( 

147 f"For {self.op_name}, dimensions must be integers, " 

148 f"but got {type(dim0)} and {type(dim1)}." 

149 ) 

150 

151 if dim0 < 0: 

152 dim0 += ndim 

153 if dim1 < 0: 

154 dim1 += ndim 

155 

156 if not (0 <= dim0 < ndim and 0 <= dim1 < ndim): 

157 raise ValueError( 

158 f"For {self.op_name}, transpose dimensions out of bounds: " 

159 f"({dim0}, {dim1}) for rank {ndim}." 

160 ) 

161 

162 out_tensor_map_list = list(in_tensor_map) 

163 out_tensor_map_list[dim0], out_tensor_map_list[dim1] = ( 

164 out_tensor_map_list[dim1], out_tensor_map_list[dim0] 

165 ) 

166 out_tensor_map = tuple(out_tensor_map_list) 

167 

168 output_layout = Layout( 

169 mesh_shape=layout.mesh_shape, 

170 alias_name=layout.alias_name, 

171 rank_list=layout.rank_list 

172 ) 

173 

174 return ((output_layout(*out_tensor_map),), None) 

175