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

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 torch.vstack operator. 

17""" 

18 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.dtensor import DTensor 

22from hyper_parallel.core.dtensor.layout import Layout 

23from .parallel_concat import ConcatDistributedOp 

24 

25 

26def _normalize_vstack_args(tensors, *, out=None): 

27 """Normalize torch.vstack arguments. 

28 

29 vstack takes a single positional argument (sequence of tensors) 

30 and an optional keyword-only ``out``. 

31 """ 

32 return (tensors,), {"out": out} 

33 

34 

35def _promote_tensor_map(alias_tensor_map): 

36 """Apply atleast_2d semantic promotion to a tensor_map tuple. 

37 

38 Returns the promoted alias_tensor_map. 

39 """ 

40 ndim = len(alias_tensor_map) 

41 if ndim == 0: 

42 return ("None", "None") 

43 if ndim == 1: 

44 return ("None",) + alias_tensor_map 

45 return alias_tensor_map 

46 

47 

48class VstackDistributedOp(ConcatDistributedOp): 

49 """Distributed implementation for torch.vstack(). 

50 

51 vstack = cat(atleast_2d(*tensors), dim=0). 

52 

53 Inherits from ConcatDistributedOp to reuse cat's layout validation 

54 (Partial check, dim-0 Replicate constraint, same-layout requirement). 

55 Only the atleast_2d promotion and DTensor enforcement are added. 

56 """ 

57 

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

59 """Preprocess arguments for vstack operator. 

60 

61 Responsibilities: 

62 - Normalize args/kwargs 

63 - Enforce all-DTensor and out=None 

64 - Extract local tensors 

65 - Cache original layouts (no promotion — that's infer_layout's job) 

66 

67 Args: 

68 args: Raw positional args from call site. 

69 kwargs: Raw keyword args from call site. 

70 

71 Returns: 

72 tuple: (local_args, local_kwargs, cache_values) 

73 cache_values = [original_layout_0, ..., original_layout_n-1] 

74 """ 

75 args, kwargs = _normalize_vstack_args(*args, **kwargs) 

76 tensors = args[0] 

77 out = kwargs["out"] 

78 

79 if out is not None: 

80 raise ValueError( 

81 f"For {self.op_name}, out keyword is not supported. " 

82 f"vstack currently only supports out=None." 

83 ) 

84 

85 # Enforce all-DTensor policy 

86 for i, t in enumerate(tensors): 

87 if not isinstance(t, DTensor): 

88 raise ValueError( 

89 f"For {self.op_name}, all inputs must be DTensor, " 

90 f"but input {i} is {type(t).__name__}." 

91 ) 

92 

93 local_tensors = tuple(t.to_local() for t in tensors) 

94 

95 local_args = (local_tensors,) 

96 local_kwargs = {} 

97 cache_values = [t.layout for t in tensors] 

98 return local_args, local_kwargs, cache_values 

99 

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

101 """Infer output layout for vstack operator. 

102 

103 1. Check Partial on original layouts (before promotion) 

104 2. Promote each layout via atleast_2d 

105 3. Delegate to parent ConcatDistributedOp with promoted layouts + [0] 

106 

107 Args: 

108 cache_values: [original_layout_0, ..., original_layout_n-1] 

109 

110 Returns: 

111 tuple: ((output_layout,), None) 

112 

113 Raises: 

114 ValueError: If layout constraints are violated. 

115 """ 

116 original_layouts = cache_values 

117 

118 # Check Partial on original layouts before promotion 

119 if not self._allow_partial_inputs: 

120 self._check_partial_inputs(original_layouts) 

121 

122 # Apply atleast_2d promotion to each layout 

123 promoted_layouts = [] 

124 for layout in original_layouts: 

125 promoted_map = _promote_tensor_map(layout.alias_tensor_map) 

126 if promoted_map == layout.alias_tensor_map: 

127 # ndim >= 2, unchanged → reuse original 

128 promoted_layouts.append(layout) 

129 else: 

130 promoted = Layout( 

131 mesh_shape=layout.mesh_shape, 

132 alias_name=layout.alias_name, 

133 rank_list=layout.rank_list, 

134 ) 

135 promoted = promoted(*promoted_map) 

136 promoted_layouts.append(promoted) 

137 

138 # Delegate to parent: partial already checked, promoted layouts + dim=0 

139 return super().infer_layout(promoted_layouts + [0]) 

140 

141 # get_expand_impl not overridden — returns None from parent.