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

1# Copyright 2025-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 ArgMaxWithValue operator. 

17""" 

18 

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_argmax_with_value_args(x, axis, keep_dims=False): 

25 return (x, axis, keep_dims), {} 

26 

27 

28class ArgMaxWithValueDistributedOp(DistributedOp): 

29 """Distributed implementation for ArgMaxWithValue operator.""" 

30 

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

32 """ 

33 Preprocess arguments for ArgMaxWithValue operator. 

34 

35 Args: 

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

37 kwargs (dict): Keyword arguments (axis, keep_dims). 

38 

39 Returns: 

40 tuple: (local_args, local_kwargs, cache_values) 

41 """ 

42 args, kwargs = _normalize_argmax_with_value_args(*args, **kwargs) 

43 input_tensor = args[0] 

44 axis = args[1] 

45 keep_dims = args[2] 

46 

47 local_args = (input_tensor.to_local(), axis, keep_dims) 

48 local_kwargs = {} 

49 

50 cache_values = [input_tensor.layout, axis, keep_dims] 

51 return local_args, local_kwargs, cache_values 

52 

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

54 """ 

55 Infer output layouts for ArgMaxWithValue operator. 

56 

57 Rules: 

58 1. Input must not have Partial status. 

59 2. axis must be an integer within the valid range [-ndim, ndim-1]. 

60 3. The axis dimension must not be sharded (including StridedShard multi-axis mappings). 

61 4. If keep_dims is False, the reduced dimension is removed from the output layout. 

62 5. Both output layouts (values, indices) are identical. 

63 

64 Args: 

65 cache_values (list): [input_layout, axis, keep_dims] 

66 

67 Returns: 

68 tuple: ((values_layout, indices_layout), None) 

69 

70 Raises: 

71 ValueError: If any rule above is violated. 

72 """ 

73 input_layout = cache_values[0] 

74 axis = cache_values[1] 

75 keep_dims = cache_values[2] 

76 

77 # Check partial inputs 

78 if not self._allow_partial_inputs: 

79 self._check_partial_inputs([input_layout]) 

80 

81 if not isinstance(axis, int): 

82 raise ValueError( 

83 f"For {self.op_name}, axis should be int, but got {type(axis)}" 

84 ) 

85 

86 rank = len(input_layout.tensor_map) 

87 

88 if axis < 0: 

89 axis += rank 

90 if axis < 0 or axis >= rank: 

91 raise ValueError( 

92 f"For {self.op_name}, axis out of range " 

93 f"(expected to be in range of [{-rank}, {rank - 1}], but got {cache_values[1]})" 

94 ) 

95 

96 # Check if the axis dimension is sharded. 

97 # Use alias_tensor_map to support StridedShard multi-axis mappings. 

98 alias_map = input_layout.alias_tensor_map 

99 mapping = alias_map[axis] 

100 if isinstance(mapping, tuple): 

101 is_sharded = any(m != "None" for m in mapping) 

102 else: 

103 is_sharded = mapping != "None" 

104 

105 if is_sharded: 

106 raise ValueError( 

107 f"For {self.op_name}, cannot perform sharding on axis dim " 

108 f"(dim {axis} mapped to {mapping}). " 

109 f"Please redistribute the tensor to Replicate on this dimension." 

110 ) 

111 

112 # Build output tensor map 

113 if not keep_dims: 

114 tensor_map = alias_map[:axis] + alias_map[axis + 1:] 

115 else: 

116 tensor_map = alias_map[:axis] + ("None",) + alias_map[axis + 1:] 

117 

118 output_layout = input_layout(*tensor_map) 

119 return ((output_layout, output_layout), None) 

120