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

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

17""" 

18 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.layout import Layout 

22from hyper_parallel.platform import get_platform 

23from .parallel_ops import DistributedOp 

24 

25platform = get_platform() 

26 

27 

28def _normalize_histc_args(x, bins=100, min_val=0, max_val=0): 

29 return (x, bins, min_val, max_val), {} 

30 

31 

32class HistcExtDistributedOp(DistributedOp): 

33 """ 

34 Distributed implementation for HistcExt operator. 

35 

36 HistcExt computes the histogram of a tensor. In distributed setting: 

37 - Each device computes a local histogram 

38 - Local histograms are aggregated using AllReduce(SUM) 

39 - Output is always replicated (1D tensor with shape (bins,)) 

40 """ 

41 

42 def __init__(self, op_name: str = "HistcExt") -> None: 

43 """Initialize HistcExtDistributedOp.""" 

44 super().__init__(op_name) 

45 

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

47 """ 

48 Preprocess arguments for HistcExt operator. 

49 

50 Args: 

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

52 kwargs (dict): Keyword arguments (bins, min, max). 

53 

54 Returns: 

55 tuple: (local_args, local_kwargs, cache_values) 

56 """ 

57 # Map external API parameter names (min, max) to internal names to avoid 

58 # shadowing Python builtins. 

59 if "min" in kwargs: 

60 kwargs["min_val"] = kwargs.pop("min") 

61 if "max" in kwargs: 

62 kwargs["max_val"] = kwargs.pop("max") 

63 args, kwargs = _normalize_histc_args(*args, **kwargs) 

64 x, bins, min_val, max_val = args 

65 local_args = (x.to_local(), bins, min_val, max_val) 

66 local_kwargs = {} 

67 cache_values = [x.layout, bins, min_val, max_val] 

68 return local_args, local_kwargs, cache_values 

69 

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

71 """ 

72 Infer output layout for HistcExt operator. 

73 

74 Rules: 

75 1. Input layout must not be None. 

76 2. bins must be a positive integer. 

77 3. min and max must be numbers with min <= max. 

78 4. Output is always a 1D replicated tensor of shape (bins,). 

79 5. When input is sharded, output carries Partial(sum) on sharded device axes. 

80 

81 Args: 

82 cache_values (list): [input_layout, bins, min, max] 

83 

84 Returns: 

85 tuple: ((output_layout,), None) 

86 

87 Raises: 

88 ValueError: If any rule above is violated. 

89 """ 

90 x_layout = cache_values[0] 

91 bins = cache_values[1] 

92 min_val = cache_values[2] 

93 max_val = cache_values[3] 

94 

95 if not self._allow_partial_inputs: 

96 self._check_partial_inputs([x_layout]) 

97 

98 if x_layout is None or x_layout.mesh_shape is None: 

99 raise ValueError( 

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

101 f"but got {x_layout}" 

102 ) 

103 

104 if not isinstance(bins, int): 

105 raise ValueError( 

106 f"For {self.op_name}, bins should be an integer, " 

107 f"but got {type(bins).__name__}" 

108 ) 

109 if bins <= 0: 

110 raise ValueError( 

111 f"For {self.op_name}, bins should be a positive integer, " 

112 f"but got {bins}" 

113 ) 

114 if not isinstance(min_val, (int, float)): 

115 raise ValueError( 

116 f"For {self.op_name}, min should be a number, " 

117 f"but got {type(min_val).__name__}" 

118 ) 

119 if not isinstance(max_val, (int, float)): 

120 raise ValueError( 

121 f"For {self.op_name}, max should be a number, " 

122 f"but got {type(max_val).__name__}" 

123 ) 

124 if min_val > max_val: 

125 raise ValueError( 

126 f"For {self.op_name}, min should be less than or equal to max, " 

127 f"but got min={min_val}, max={max_val}" 

128 ) 

129 

130 output_layout = Layout( 

131 mesh_shape=x_layout.mesh_shape, 

132 alias_name=x_layout.alias_name, 

133 rank_list=x_layout.rank_list, 

134 ) 

135 out_layout = output_layout("None",) 

136 

137 has_sharding = any( 

138 alias is not None and alias != "None" 

139 for alias in x_layout.alias_tensor_map 

140 ) 

141 

142 if has_sharding: 

143 for alias, tensor_map_val in zip(x_layout.alias_name, x_layout.alias_tensor_map): 

144 if tensor_map_val is not None and tensor_map_val != "None": 

145 out_layout.set_partial_by_dev_axis(alias, "sum") 

146 

147 return ((out_layout,), None)