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

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

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_cumsum_args(x, dim, dtype=None): 

26 return (x, dim), {'dtype': dtype} 

27 

28 

29class CumsumDistributedOp(DistributedOp): 

30 """Distributed implementation for torch.cumsum.""" 

31 _MS_PRIMITIVE_OP_NAMES = frozenset({'CumsumExt'}) 

32 

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

34 """ 

35 Preprocess arguments for Cumsum operator. 

36 

37 Args: 

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

39 kwargs (dict): Keyword arguments (dim, dtype). 

40 

41 Returns: 

42 tuple: (local_args, local_kwargs, cache_values) 

43 """ 

44 args, kwargs = _normalize_cumsum_args(*args, **kwargs) 

45 input_tensor = args[0] 

46 dim = args[1] 

47 dtype = kwargs['dtype'] 

48 

49 local_input = input_tensor.to_local() 

50 if self.op_name in self._MS_PRIMITIVE_OP_NAMES: 

51 local_args = (local_input, dim, dtype) 

52 local_kwargs = {} 

53 else: 

54 local_args = (local_input, dim) 

55 local_kwargs = {} 

56 if dtype is not None: 

57 local_kwargs['dtype'] = dtype 

58 

59 cache_values = [input_tensor.layout, dim] 

60 return local_args, local_kwargs, cache_values 

61 

62 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221 

63 """ 

64 Infer output layout for Cumsum operator. 

65 

66 Rules: 

67 1. Input must not have Partial status. 

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

69 3. The cumsum dimension must not be sharded, including StridedShard mappings. 

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

71 

72 Args: 

73 cache_values (list): [input_layout, dim] 

74 

75 Returns: 

76 tuple: ((output_layout,), None) 

77 

78 Raises: 

79 ValueError: If input has Partial status, dim is not an int, dim is out of range, 

80 or the cumsum dimension is sharded. 

81 """ 

82 layout = cache_values[0] 

83 dim = cache_values[1] 

84 

85 if not self._allow_partial_inputs: 

86 self._check_partial_inputs([layout]) 

87 

88 if not isinstance(dim, int): 

89 raise ValueError( 

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

91 ) 

92 

93 alias_map = layout.alias_tensor_map 

94 ndim = len(alias_map) 

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

96 raise ValueError( 

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

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

99 ) 

100 if dim < 0: 

101 dim += ndim 

102 

103 mapping = alias_map[dim] 

104 if mapping != "None": 

105 raise ValueError( 

106 f"For {self.op_name}, cumsum along a sharded dimension " 

107 f"(dim {dim} mapped to {mapping}) is not supported." 

108 ) 

109 

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