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

1# Copyright 2025 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 ExpandDims operator. 

17""" 

18from typing import Tuple 

19 

20from hyper_parallel.core.dtensor.layout import Layout 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_expand_dims_args(x, axis=None, dim=None): 

25 if axis is None: 

26 axis = dim 

27 return (x, axis), {} 

28 

29 

30class ExpandDimsDistributedOp(DistributedOp): 

31 """Distributed implementation for ExpandDims operator.""" 

32 

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

34 """ 

35 Preprocess arguments for ExpandDims operator. 

36 

37 Args: 

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

39 kwargs (dict): Keyword arguments (axis or dim). 

40 

41 Returns: 

42 tuple: (local_args, local_kwargs, cache_values) 

43 """ 

44 args, _ = _normalize_expand_dims_args(*args, **kwargs) 

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

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

47 cache_values = [input_tensor.layout, axis] 

48 return local_args, {}, cache_values 

49 

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

51 """ 

52 Infer output layout for ExpandDims operator. 

53 

54 Rules: 

55 1. Input must not have Partial status. 

56 2. axis must be an integer within the valid range [-(rank + 1), rank]. 

57 3. The inserted dimension is replicated. 

58 4. Existing input dimension mappings are shifted and otherwise preserved. 

59 

60 Args: 

61 cache_values (list): [input_layout, axis] 

62 

63 Returns: 

64 tuple: ((output_layout,), None) 

65 

66 Raises: 

67 ValueError: If input has Partial status, input layout is missing, 

68 axis is missing or invalid, or axis is out of range. 

69 """ 

70 if not cache_values: 

71 raise ValueError( 

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

73 f"but got empty cache_values." 

74 ) 

75 

76 x_layout = cache_values[0] 

77 if not self._allow_partial_inputs: 

78 self._check_partial_inputs([x_layout]) 

79 

80 if x_layout.mesh_shape is None: 

81 raise ValueError( 

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

83 f"but got None." 

84 ) 

85 

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

87 

88 if axis is None: 

89 raise ValueError(f"For {self.op_name}, axis parameter is required.") 

90 if not isinstance(axis, int): 

91 raise ValueError( 

92 f"For {self.op_name}, axis should be int, but got {type(axis)}." 

93 ) 

94 

95 in_rank = len(x_layout.alias_tensor_map) 

96 original_axis = axis 

97 if axis < 0: 

98 axis = axis + in_rank + 1 

99 

100 if axis < 0 or axis > in_rank: 

101 raise ValueError( 

102 f"For {self.op_name}, axis {original_axis} out of range for input rank {in_rank}. " 

103 f"Valid range is [{-in_rank - 1}, {in_rank}]." 

104 ) 

105 

106 x_map = list(x_layout.alias_tensor_map) 

107 x_map.insert(axis, "None") 

108 

109 output_layout = Layout( 

110 mesh_shape=x_layout.mesh_shape, 

111 alias_name=x_layout.alias_name, 

112 rank_list=x_layout.rank_list 

113 ) 

114 output_layout = output_layout(*x_map) 

115 

116 if self._allow_partial_inputs: 

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

118 if partial_op is not None: 

119 dev_axis_name = x_layout.alias_name[i] 

120 output_layout.set_partial_by_dev_axis(dev_axis_name, partial_op) 

121 

122 return ((output_layout,), None)