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

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 atleast_1d 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_atleast_1d_args(*tensors): 

26 return tensors, {} 

27 

28 

29class Atleast1DDistributedOp(DistributedOp): 

30 """Distributed implementation for torch.atleast_1d.""" 

31 

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

33 """ 

34 Preprocess arguments for atleast_1d operator. 

35 

36 torch.atleast_1d(*tensors) takes a variable number of tensor inputs. 

37 All arguments are positional tensors with no keyword-only parameters. 

38 

39 Args: 

40 args (tuple): Positional arguments (tensors). 

41 kwargs (dict): Keyword arguments (none expected). 

42 

43 Returns: 

44 tuple: (local_args, local_kwargs, cache_values) 

45 """ 

46 args, kwargs = _normalize_atleast_1d_args(*args, **kwargs) 

47 tensors = args 

48 

49 local_args = tuple( 

50 t.to_local() if hasattr(t, 'to_local') else t 

51 for t in tensors 

52 ) 

53 local_kwargs = {} 

54 

55 cache_values = [ 

56 t.layout if hasattr(t, 'layout') else None 

57 for t in tensors 

58 ] 

59 

60 return local_args, local_kwargs, cache_values 

61 

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

63 """ 

64 Infer output layouts for atleast_1d operator. 

65 

66 Rules: 

67 1. Inputs must not have Partial status. 

68 2. For 0D → 1D: the newly created dimension is unsharded (-1). 

69 3. For 1D or higher: the layout is preserved unchanged. 

70 4. If a single tensor is provided, a single Layout is returned. 

71 If multiple tensors are provided, a tuple of Layouts is returned. 

72 

73 Args: 

74 cache_values (list): List of Layout objects (one per input tensor). 

75 None entries represent non-DTensor inputs and produce None outputs. 

76 

77 Returns: 

78 tuple: ((output_layout_or_tuple,), None) 

79 

80 Raises: 

81 ValueError: If no inputs are provided or any input has Partial status. 

82 """ 

83 layouts = cache_values 

84 

85 if not layouts: 

86 raise ValueError( 

87 f"For {self.op_name}, at least one input tensor is required, " 

88 f"but got an empty input list." 

89 ) 

90 

91 # Check partial inputs (atleast_1d does not support partial) 

92 if not self._allow_partial_inputs: 

93 self._check_partial_inputs(layouts) 

94 

95 output_layouts = [] 

96 

97 # Process each layout for the case of multiple input tensors 

98 for input_layout in layouts: 

99 if input_layout is None: 

100 output_layouts.append(None) 

101 continue 

102 

103 in_tensor_map = input_layout.tensor_map 

104 input_ndim = len(in_tensor_map) 

105 

106 # Build output tensor map 

107 if input_ndim == 0: 

108 # 0D -> 1D: the newly created dimension is unsharded 

109 output_map = (-1,) 

110 else: 

111 # 1D or higher: preserve original layout 

112 output_map = in_tensor_map 

113 

114 # Construct output layout using the same mesh properties 

115 mesh_shape = input_layout.mesh_shape 

116 alias_name = input_layout.alias_name 

117 rank_list = input_layout.rank_list 

118 

119 def idx_to_alias(idx, aliases): 

120 if idx == -1: 

121 return "None" 

122 # Map index back to alias name string 

123 return aliases[len(aliases) - idx - 1] 

124 

125 output_alias_map = tuple(idx_to_alias(idx, alias_name) for idx in output_map) 

126 

127 out_layout = Layout( 

128 mesh_shape=mesh_shape, 

129 alias_name=alias_name, 

130 rank_list=rank_list 

131 ) 

132 # Re-apply the alias mapping to generate the full internal layout 

133 out_layout = out_layout(*output_alias_map) 

134 

135 output_layouts.append(out_layout) 

136 

137 # If there's only one input, return a single Layout. 

138 # If there are multiple inputs, return a tuple of Layouts. 

139 if len(output_layouts) == 1: 

140 return ((output_layouts[0],), None) 

141 

142 return (tuple(output_layouts), None)