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

1# Copyright 2025 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""" 

16Element-wise distributed operator implementation. 

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _unwrap_local_value(value): 

26 """Convert DTensor-like values to local tensors while preserving scalar slots.""" 

27 return value.to_local() if hasattr(value, "_layout") else value 

28 

29 

30class TupleElementWiseDistributedOp(DistributedOp): 

31 """ 

32 Distributed implementation for tuple element-wise operators. 

33 

34 Inherits from DistributedOp and provides element-wise specific implementations. 

35 """ 

36 

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

38 """ 

39 Preprocess arguments for tuple element-wise operators. 

40 

41 Args: 

42 args (tuple): Positional arguments passed to the operator. 

43 kwargs (dict): Keyword arguments passed to the operator. 

44 

45 Returns: 

46 tuple: (local_args, local_kwargs, cache_values) 

47 """ 

48 expanded_args = [] 

49 local_args = [] 

50 for arg in args: 

51 if isinstance(arg, (tuple, list)): 

52 expanded_args.extend(arg) 

53 local_args.append(tuple(_unwrap_local_value(item) for item in arg)) 

54 else: 

55 expanded_args.append(arg) 

56 local_args.append(_unwrap_local_value(arg)) 

57 

58 local_kwargs = {key: _unwrap_local_value(value) for key, value in kwargs.items()} 

59 cache_values = [getattr(arg, "layout", None) for arg in expanded_args] 

60 cache_values.extend(getattr(value, "layout", None) for value in kwargs.values()) 

61 

62 return tuple(local_args), local_kwargs, cache_values 

63 

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

65 """ 

66 Infer output layouts for element-wise operations. 

67 

68 Rules: 

69 1. Inputs must not have Partial status. 

70 2. Tuple/list positional arguments are inferred from their expanded elements. 

71 3. Output layouts are identical to the expanded input layouts. 

72 

73 Args: 

74 cache_values (list): Expanded input layouts, using None for non-DTensor slots. 

75 

76 Returns: 

77 tuple: (output_layouts, None) 

78 

79 Raises: 

80 ValueError: If input has Partial status. 

81 """ 

82 if not cache_values: 

83 return None 

84 

85 self._check_partial_inputs(cache_values) 

86 

87 output_layouts = tuple( 

88 copy.deepcopy(layout) if layout is not None else None 

89 for layout in cache_values 

90 ) 

91 return output_layouts, None