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

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

17""" 

18from typing import Tuple 

19 

20from hyper_parallel.core.shard.ops.parallel_reshape import ReshapeDistributedOp 

21 

22 

23def _normalize_flatten_args(x, start_dim=0, end_dim=-1): 

24 return (x, start_dim, end_dim), {} 

25 

26 

27class FlattenDistributedOp(ReshapeDistributedOp): 

28 """Distributed implementation for torch.flatten.""" 

29 

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

31 """ 

32 Preprocess arguments for Flatten operator. 

33 

34 Args: 

35 args (tuple): Input tensor followed by optional start_dim and end_dim. 

36 kwargs (dict): Optional keyword arguments. 

37 

38 Returns: 

39 tuple: (local_args, local_kwargs, cache_values) 

40 """ 

41 args, _ = _normalize_flatten_args(*args, **kwargs) 

42 input_tensor, start_dim, end_dim = args[0], args[1], args[2] 

43 local_args = (input_tensor.to_local(), start_dim, end_dim) 

44 local_kwargs = {} 

45 cache_values = [input_tensor.layout, start_dim, end_dim, tuple(input_tensor.shape)] 

46 return local_args, local_kwargs, cache_values 

47 

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

49 """ 

50 Infer output layout for Flatten operator. 

51 

52 Rules: 

53 1. Partial input is allowed and preserved by reshape layout inference. 

54 2. input_shape must be provided and match the input rank. 

55 3. start_dim and end_dim must be integers within the valid input rank. 

56 4. If start_dim >= end_dim after normalization, the output layout is the 

57 input layout. 

58 5. Otherwise, dimensions from start_dim through end_dim are merged using 

59 reshape-compatible sharding rules. 

60 

61 Args: 

62 cache_values (list): [input_layout, start_dim, end_dim, input_shape]. 

63 

64 Returns: 

65 tuple: ((output_layout,), None) 

66 

67 Raises: 

68 ValueError: If cache_values are invalid, dimensions are out of range, 

69 or the flatten would change sharded slices incompatibly. 

70 """ 

71 input_layout, start_dim, end_dim, input_shape = ( 

72 cache_values[0], cache_values[1], cache_values[2], cache_values[3] 

73 ) 

74 if input_layout is None: 

75 raise ValueError( 

76 f"For {self.op_name}, flatten requires a valid input tensor layout." 

77 ) 

78 if not isinstance(input_shape, (list, tuple)): 

79 raise ValueError( 

80 f"For {self.op_name}, input_shape should be list or tuple, " 

81 f"but got {type(input_shape)}." 

82 ) 

83 if len(input_shape) != len(input_layout.tensor_map): 

84 raise ValueError( 

85 f"For {self.op_name}, input shape rank should match layout rank, " 

86 f"but got {len(input_shape)} and {len(input_layout.tensor_map)}." 

87 ) 

88 if not isinstance(start_dim, int) or not isinstance(end_dim, int): 

89 raise ValueError( 

90 f"For {self.op_name}, start_dim and end_dim should be int, " 

91 f"but got {type(start_dim)} and {type(end_dim)}." 

92 ) 

93 

94 ndim = len(input_shape) 

95 

96 if ndim == 0: 

97 out_layout = input_layout.__class__.from_device_mesh(input_layout.mesh) 

98 out_layout.set_placements(input_layout.placements) 

99 out_layout.placement_to_tensor_map(1) 

100 return ((out_layout,), None) 

101 

102 if start_dim < 0: 

103 start_dim += ndim 

104 if end_dim < 0: 

105 end_dim += ndim 

106 

107 start_dim_invalid = start_dim < 0 or start_dim >= ndim 

108 end_dim_invalid = end_dim < 0 or end_dim >= ndim 

109 if start_dim_invalid or end_dim_invalid: 

110 raise ValueError( 

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

112 f"(start_dim={start_dim}, end_dim={end_dim}, ndim={ndim})." 

113 ) 

114 

115 if start_dim >= end_dim: 

116 return ((input_layout,), None) 

117 

118 flattened_size = 1 

119 for i in range(start_dim, end_dim + 1): 

120 flattened_size *= input_shape[i] 

121 dst_shape = list(input_shape[:start_dim]) + [flattened_size] + list(input_shape[end_dim + 1:]) 

122 

123 out_layout, _ = self._infer_reshape_layout(input_layout, dst_shape, input_shape) 

124 return ((out_layout,), None)