Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_unbind.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 Unbind 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_unbind_args(input_tensor, dim=0): 

26 return (input_tensor, dim), {} 

27 

28 

29class UnbindDistributedOp(DistributedOp): 

30 """Distributed implementation for Unbind operator.""" 

31 

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

33 """ 

34 Preprocess arguments for Unbind operator. 

35 

36 Args: 

37 args (tuple): Input arguments (input_tensor, dim). 

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_unbind_args(*args, **kwargs) 

44 input_tensor, dim = args 

45 

46 local_args = (input_tensor.to_local(), dim) 

47 local_kwargs = {} 

48 cache_values = [input_tensor.layout, tuple(input_tensor.shape), dim] 

49 return local_args, local_kwargs, cache_values 

50 

51 # pylint: disable=W0237 

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

53 """ 

54 Infer output layouts for Unbind operator. 

55 

56 Rules: 

57 1. Input must not have Partial status. 

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

59 3. The dimension to unbind must not be sharded. 

60 4. Output layout removes the mapping for the unbound dimension; 

61 all output tensors share the same layout. 

62 

63 Args: 

64 cache_values (list): [input_layout, input_shape, dim] 

65 

66 Returns: 

67 tuple: ((output_layouts_tuple,), None) 

68 

69 Raises: 

70 ValueError: If any rule above is violated. 

71 """ 

72 layout = cache_values[0] 

73 shape = cache_values[1] 

74 dim = cache_values[2] 

75 

76 if not self._allow_partial_inputs: 

77 self._check_partial_inputs([layout]) 

78 

79 alias_tensor_map = layout.alias_tensor_map 

80 ndim = len(shape) 

81 

82 if not isinstance(dim, int): 

83 raise ValueError( 

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

85 ) 

86 

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

88 raise ValueError( 

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

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

91 ) 

92 

93 if dim < 0: 

94 dim += ndim 

95 

96 # Check if the dimension to unbind is sharded. 

97 # alias_tensor_map returns "None" for replicated dimensions. 

98 if alias_tensor_map[dim] != "None": 

99 raise ValueError( 

100 f"For {self.op_name}, the dimension {dim} is sharded " 

101 f"(mapped to {alias_tensor_map[dim]}). " 

102 f"Unbinding a sharded dimension is not supported. " 

103 f"Please redistribute the tensor to replicate this dimension first." 

104 ) 

105 

106 # Construct output layout: remove the mapping for the unbound dimension 

107 out_alias_map = alias_tensor_map[:dim] + alias_tensor_map[dim + 1:] 

108 

109 base_layout = Layout( 

110 mesh_shape=layout.mesh_shape, 

111 alias_name=layout.alias_name, 

112 rank_list=layout.rank_list 

113 ) 

114 out_layout = base_layout(*out_alias_map) 

115 

116 num_outputs = shape[dim] 

117 return ((out_layout,) * num_outputs, None)