Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_nonzero.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 Nonzero 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_nonzero_args(x, as_tuple=False): 

26 return (x,), {'as_tuple': as_tuple} 

27 

28 

29class NonzeroDistributedOp(DistributedOp): 

30 """Distributed implementation for torch.nonzero.""" 

31 

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

33 """ 

34 Preprocess arguments for Nonzero operator. 

35 

36 Args: 

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

38 kwargs (dict): Keyword arguments (as_tuple). 

39 

40 Returns: 

41 tuple: (local_args, local_kwargs, cache_values) 

42 """ 

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

44 input_tensor = args[0] 

45 as_tuple = kwargs['as_tuple'] 

46 

47 local_args = (input_tensor.to_local(),) 

48 local_kwargs = {'as_tuple': as_tuple} 

49 

50 cache_values = [input_tensor.layout, as_tuple] 

51 return local_args, local_kwargs, cache_values 

52 

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

54 """ 

55 Infer output layouts for Nonzero operator. 

56 

57 Rules: 

58 1. Input must not have Partial status. 

59 2. Input must be fully replicated (all dimensions mapped to "None"). 

60 nonzero produces data-dependent output shapes, which would differ 

61 across sharded ranks. 

62 3. If as_tuple=True: returns a tuple of 1D replicated layouts, one per 

63 input dimension. 

64 4. If as_tuple=False: returns a single 2D replicated layout. 

65 

66 Args: 

67 cache_values (list): [input_layout, as_tuple] 

68 

69 Returns: 

70 tuple: ((output_layout(s),), None) 

71 

72 Raises: 

73 ValueError: If input has Partial status or is sharded. 

74 """ 

75 input_layout = cache_values[0] 

76 as_tuple = cache_values[1] 

77 

78 if input_layout is None: 

79 raise ValueError( 

80 f"For {self.op_name}, input_layout should be a valid Layout, but got None." 

81 ) 

82 

83 # Rule 1: Input must not have Partial status 

84 if not self._allow_partial_inputs: 

85 self._check_partial_inputs([input_layout]) 

86 

87 if not isinstance(as_tuple, bool): 

88 raise ValueError( 

89 f"For {self.op_name}, as_tuple should be bool, but got {type(as_tuple)}." 

90 ) 

91 

92 alias_map = input_layout.alias_tensor_map 

93 input_ndim = len(alias_map) 

94 

95 # Rule 2: Input must be fully replicated due to data-dependent dynamic shapes 

96 for dim, dim_sharding in enumerate(alias_map): 

97 if dim_sharding != "None": 

98 raise ValueError( 

99 f"For {self.op_name}, input tensor should be fully replicated, " 

100 f"but got dim {dim} mapped to {dim_sharding}. " 

101 f"nonzero produces dynamic shapes that depend on data values, " 

102 f"which causes shape mismatches across ranks if the tensor is sharded." 

103 ) 

104 

105 mesh_shape = input_layout.mesh_shape 

106 alias_name = input_layout.alias_name 

107 rank_list = input_layout.rank_list 

108 

109 def _create_replicated_layout(ndim): 

110 """Helper to create a fully replicated layout for a given dimension.""" 

111 layout = Layout( 

112 mesh_shape=mesh_shape, 

113 alias_name=alias_name, 

114 rank_list=rank_list 

115 ) 

116 alias_map = tuple("None" for _ in range(ndim)) 

117 return layout(*alias_map) 

118 

119 # Rule 3 & 4: Construct the return layout based on as_tuple flag 

120 if as_tuple: 

121 out_layout = _create_replicated_layout(1) 

122 return (tuple(out_layout for _ in range(input_ndim)), None) 

123 

124 return ((_create_replicated_layout(2),), None)