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

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 Outer 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_outer_args(vec1, vec2): 

26 return (vec1, vec2), {} 

27 

28 

29def _get_alias_shard_set(dim_alias): 

30 if isinstance(dim_alias, str): 

31 return {dim_alias} if dim_alias != "None" else set() 

32 return set(dim_alias) 

33 

34 

35class OuterDistributedOp(DistributedOp): 

36 """Distributed implementation for torch.outer.""" 

37 

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

39 """ 

40 Preprocess arguments for Outer operator. 

41 

42 Args: 

43 args (tuple): Input arguments (input, vec2). 

44 kwargs (dict): Keyword arguments (unused for outer). 

45 

46 Returns: 

47 tuple: (local_args, local_kwargs, cache_values) where local_args contains 

48 local tensors for input and vec2, and cache_values contains their layouts. 

49 """ 

50 args, kwargs = _normalize_outer_args(*args, **kwargs) 

51 input_tensor, vec2_tensor = args[0], args[1] 

52 local_args = (input_tensor.to_local(), vec2_tensor.to_local()) 

53 local_kwargs = {} 

54 cache_values = [input_tensor.layout, vec2_tensor.layout] 

55 return local_args, local_kwargs, cache_values 

56 

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

58 """ 

59 Infer output layout for Outer operator. 

60 

61 PyTorch semantics: 

62 - Computes the outer product of two 1-D tensors. 

63 - If input is of size N and vec2 is of size M, the output is of size (N, M). 

64 - Input tensors must be 1-D. 

65 

66 Rules: 

67 1. Inputs must not have Partial status. 

68 2. Exactly two input layouts are required, both must be non-None. 

69 3. Both inputs must be exactly 1-D. 

70 4. The two inputs cannot be sharded along the same device mesh dimension. 

71 5. Output dim 0 inherits the sharding of input; output dim 1 inherits the sharding of vec2. 

72 

73 Args: 

74 cache_values (list): [input_layout, vec2_layout] 

75 

76 Returns: 

77 tuple: ((output_layout,), None) 

78 

79 Raises: 

80 ValueError: If any rule above is violated. 

81 """ 

82 layout1, layout2 = cache_values[0], cache_values[1] 

83 

84 if layout1 is None or layout2 is None: 

85 raise ValueError( 

86 f"For {self.op_name}, both inputs should be DTensors with valid layouts, " 

87 f"but got layout1={layout1}, layout2={layout2}." 

88 ) 

89 

90 if not self._allow_partial_inputs: 

91 self._check_partial_inputs([layout1, layout2]) 

92 

93 alias_map1 = layout1.alias_tensor_map 

94 alias_map2 = layout2.alias_tensor_map 

95 

96 if len(alias_map1) != 1 or len(alias_map2) != 1: 

97 raise ValueError( 

98 f"For {self.op_name}, both inputs should be exactly 1-D tensors, " 

99 f"but got {len(alias_map1)}-D and {len(alias_map2)}-D." 

100 ) 

101 

102 dim0_alias = alias_map1[0] 

103 dim1_alias = alias_map2[0] 

104 

105 set1 = _get_alias_shard_set(dim0_alias) 

106 set2 = _get_alias_shard_set(dim1_alias) 

107 

108 if set1.intersection(set2): 

109 raise ValueError( 

110 f"For {self.op_name}, the two inputs should not be sharded on the " 

111 f"same device mesh dimension, " 

112 f"but got conflict on mesh dimension(s): {set1.intersection(set2)}." 

113 ) 

114 

115 output_alias_map = (dim0_alias, dim1_alias) 

116 

117 output_layout = Layout( 

118 mesh_shape=layout1.mesh_shape, 

119 alias_name=layout1.alias_name, 

120 rank_list=layout1.rank_list, 

121 ) 

122 output_layout = output_layout(*output_alias_map) 

123 

124 return ((output_layout,), None) 

125