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

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 .parallel_ops import DistributedOp 

23 

24 

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

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

27 

28 

29class HistcExtDistributedOp(DistributedOp): 

30 """ 

31 Distributed implementation for HistcExt operator. 

32 

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

34 - Each device computes a local histogram 

35 - Local histograms are aggregated using AllReduce(SUM) 

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

37 """ 

38 

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

40 """Initialize HistcExtDistributedOp.""" 

41 super().__init__(op_name) 

42 

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

44 """ 

45 Preprocess arguments for HistcExt operator. 

46 

47 Args: 

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

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

50 

51 Returns: 

52 tuple: (local_args, local_kwargs, cache_values) 

53 """ 

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

55 # shadowing Python builtins. 

56 if "min" in kwargs: 

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

58 if "max" in kwargs: 

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

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

61 x, bins, min_val, max_val = args 

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

63 local_kwargs = {} 

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

65 return local_args, local_kwargs, cache_values 

66 

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

68 """ 

69 Infer output layout for HistcExt operator. 

70 

71 Rules: 

72 1. Input layout must not be None. 

73 2. bins must be a positive integer. 

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

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

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

77 

78 Args: 

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

80 

81 Returns: 

82 tuple: ((output_layout,), None) 

83 

84 Raises: 

85 ValueError: If any rule above is violated. 

86 """ 

87 x_layout = cache_values[0] 

88 bins = cache_values[1] 

89 min_val = cache_values[2] 

90 max_val = cache_values[3] 

91 

92 if not self._allow_partial_inputs: 

93 self._check_partial_inputs([x_layout]) 

94 

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

96 raise ValueError( 

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

98 f"but got {x_layout}" 

99 ) 

100 

101 if not isinstance(bins, int): 

102 raise ValueError( 

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

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

105 ) 

106 if bins <= 0: 

107 raise ValueError( 

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

109 f"but got {bins}" 

110 ) 

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

112 raise ValueError( 

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

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

115 ) 

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

117 raise ValueError( 

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

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

120 ) 

121 if min_val > max_val: 

122 raise ValueError( 

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

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

125 ) 

126 

127 output_layout = Layout( 

128 mesh_shape=x_layout.mesh_shape, 

129 alias_name=x_layout.alias_name, 

130 rank_list=x_layout.rank_list, 

131 ) 

132 out_layout = output_layout("None",) 

133 

134 has_sharding = any( 

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

136 for alias in x_layout.alias_tensor_map 

137 ) 

138 

139 if has_sharding: 

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

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

142 out_layout.set_partial_by_dev_axis(alias, "sum") 

143 

144 return ((out_layout,), None)