Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_argsort.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 Argsort operator. 

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_argsort_args(x, dim=-1, descending=False, stable=False): 

26 """Normalize torch.argsort arguments to (args, kwargs). 

27 

28 torch.argsort(input, dim=-1, descending=False, *, stable=False) 

29 Only `stable` is keyword-only; all other params are positional. 

30 """ 

31 return (x, dim, descending), {'stable': stable} 

32 

33 

34class ArgsortDistributedOp(DistributedOp): 

35 """Distributed implementation for torch.argsort.""" 

36 

37 _MS_PRIMITIVE_OP_NAMES = frozenset({'ArgSort'}) 

38 

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

40 """ 

41 Preprocess arguments for Argsort operator. 

42 

43 Args: 

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

45 kwargs (dict): Keyword arguments (dim, descending, stable). 

46 

47 Returns: 

48 tuple: (local_args, local_kwargs, cache_values) 

49 """ 

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

51 input_tensor = args[0] 

52 dim = args[1] 

53 descending = args[2] 

54 stable = kwargs['stable'] 

55 

56 if self.op_name in self._MS_PRIMITIVE_OP_NAMES: 

57 local_args = (input_tensor.to_local(), dim, descending, stable) 

58 local_kwargs = {} 

59 else: 

60 local_args = (input_tensor.to_local(),) 

61 local_kwargs = {'dim': dim, 'descending': descending, 'stable': stable} 

62 

63 cache_values = [input_tensor.layout, dim] 

64 return local_args, local_kwargs, cache_values 

65 

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

67 """ 

68 Infer output layout for Argsort operator. 

69 

70 Rules: 

71 1. Input must not have Partial status. 

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

73 3. The sort dimension must not be sharded. 

74 4. Output layout is identical to the input layout. 

75 

76 Args: 

77 cache_values (list): [input_layout, dim] where dim is the sort dimension. 

78 

79 Returns: 

80 tuple: ((output_layout,), None) 

81 

82 Raises: 

83 ValueError: If input has Partial status, dim is out of range, or the sort 

84 dimension is sharded. 

85 """ 

86 layout = cache_values[0] 

87 dim = cache_values[1] 

88 

89 if not self._allow_partial_inputs: 

90 self._check_partial_inputs([layout]) 

91 

92 if not isinstance(dim, int): 

93 raise ValueError( 

94 f"For {self.op_name}, dimension should be int, but got {type(dim)}" 

95 ) 

96 

97 alias_map = layout.alias_tensor_map 

98 ndim = len(alias_map) 

99 

100 if dim < -ndim or dim >= ndim: 

101 raise ValueError( 

102 f"For {self.op_name}, dimension out of range " 

103 f"(expected to be in range of [{-ndim}, {ndim - 1}], but got {dim})" 

104 ) 

105 

106 if dim < 0: 

107 dim += ndim 

108 

109 if alias_map[dim] != "None": 

110 raise ValueError( 

111 f"For {self.op_name}, sorting along a sharded dimension " 

112 f"(dim {dim} mapped to {alias_map[dim]}) is not supported. " 

113 f"Please redistribute the tensor to Replicate on this dimension before sorting." 

114 ) 

115 

116 return ((copy.deepcopy(layout),), None)