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

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

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_slice_ext_args(x, axis, begin, end, step): 

26 return (x, axis, begin, end, step), {} 

27 

28 

29class SliceExtDistributedOp(DistributedOp): 

30 """Distributed implementation for SliceExt operator.""" 

31 

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

33 """ 

34 Preprocess arguments for SliceExt operator. 

35 

36 Args: 

37 args (tuple): Input arguments (input, axis, begin, end, step). 

38 kwargs (dict): Keyword arguments (empty for this operator). 

39 

40 Returns: 

41 tuple: (local_args, local_kwargs, cache_values) 

42 """ 

43 args, kwargs = _normalize_slice_ext_args(*args, **kwargs) 

44 input_tensor, axis, begin, end, step = args 

45 local_args = (input_tensor.to_local(), axis, begin, end, step) 

46 local_kwargs = {} 

47 cache_values = [input_tensor.layout, axis] 

48 return local_args, local_kwargs, cache_values 

49 

50 # pylint: disable=W0237 

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

52 """ 

53 Infer output layouts for SliceExt operator. 

54 

55 Rules: 

56 1. Input must not have Partial status. 

57 2. The sliced axis must not be sharded. 

58 3. Output layout is identical to the input layout. 

59 

60 Args: 

61 cache_values (list): [input_layout, axis] 

62 

63 Returns: 

64 tuple: ((output_layout,), None) 

65 

66 Raises: 

67 ValueError: If input has Partial status or the sliced axis is sharded. 

68 """ 

69 input_layout = cache_values[0] 

70 axis = cache_values[1] 

71 

72 if not self._allow_partial_inputs: 

73 self._check_partial_inputs([input_layout]) 

74 

75 alias_map = input_layout.alias_tensor_map 

76 ndim = len(alias_map) 

77 

78 if not isinstance(axis, int): 

79 raise ValueError( 

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

81 ) 

82 

83 if axis < -ndim or axis >= ndim: 

84 raise ValueError( 

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

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

87 ) 

88 

89 if axis < 0: 

90 axis += ndim 

91 

92 if alias_map[axis] != "None": 

93 raise ValueError( 

94 f"For {self.op_name}, can not slice tensor at sharded axis[{axis}], " 

95 f"layout: {input_layout}." 

96 ) 

97 

98 return ((copy.deepcopy(input_layout),), None)