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

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_isin_args(elements, test_elements, assume_unique=False, invert=False): 

26 return (elements, test_elements), {'assume_unique': assume_unique, 'invert': invert} 

27 

28 

29class IsinDistributedOp(DistributedOp): 

30 """Distributed implementation for torch.isin.""" 

31 

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

33 """ 

34 Preprocess arguments for Isin operator. 

35 

36 Args: 

37 args (tuple): Input arguments (elements, test_elements). 

38 kwargs (dict): Keyword arguments (assume_unique, invert). 

39 

40 Returns: 

41 tuple: (local_args, local_kwargs, cache_values) 

42 """ 

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

44 elements, test_elements = args[0], args[1] 

45 assume_unique = kwargs['assume_unique'] 

46 invert = kwargs['invert'] 

47 

48 local_args = ( 

49 elements.to_local(), 

50 test_elements.to_local() if hasattr(test_elements, '_layout') else test_elements, 

51 ) 

52 local_kwargs = {'assume_unique': assume_unique, 'invert': invert} 

53 

54 cache_values = [ 

55 elements.layout, 

56 test_elements.layout if hasattr(test_elements, '_layout') else None, 

57 ] 

58 return local_args, local_kwargs, cache_values 

59 

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

61 """ 

62 Infer output layout for torch.isin(elements, test_elements, ...) 

63 

64 PyTorch semantics: 

65 - Returns boolean tensor with SAME SHAPE as `elements` 

66 - Each element is tested against ALL values in `test_elements`, so requires GLOBAL view of `test_elements` 

67 

68 Rules: 

69 1. elements must have a valid DTensor layout. 

70 2. elements must not have Partial status. 

71 3. If test_elements is a DTensor, it must be fully unsharded (replicated across all dimensions) 

72 and must not have Partial status. 

73 4. If test_elements is a plain Tensor or scalar, no sharding validation is needed. 

74 5. Output layout is identical to elements layout. 

75 

76 Args: 

77 cache_values (list): [elements_layout, test_elements_layout_or_None] 

78 

79 Returns: 

80 tuple: ((output_layout,), None) 

81 

82 Raises: 

83 ValueError: If any rule above is violated. 

84 """ 

85 if not cache_values or cache_values[0] is None: 

86 raise ValueError( 

87 f"For {self.op_name}, 'elements' requires a valid tensor layout, " 

88 f"but got {cache_values[0] if cache_values else None}." 

89 ) 

90 

91 elements_layout = cache_values[0] 

92 test_elements_layout = cache_values[1] if len(cache_values) >= 2 else None 

93 

94 if not self._allow_partial_inputs: 

95 check_layouts = [elements_layout] 

96 if test_elements_layout is not None: 

97 check_layouts.append(test_elements_layout) 

98 self._check_partial_inputs(check_layouts) 

99 

100 # test_elements must be unsharded if it is a DTensor 

101 if test_elements_layout is not None: 

102 alias_map = test_elements_layout.alias_tensor_map 

103 if not all(entry == "None" for entry in alias_map): 

104 raise ValueError( 

105 f"For {self.op_name}, 'test_elements' must be unsharded, " 

106 f"but got alias_tensor_map: {alias_map}." 

107 ) 

108 

109 return ((copy.deepcopy(elements_layout),), None)