mod: can be 'EQ' (equal-to), 'NE' (not equal-to), 'GE' (greater-or-equal-than), 'GT' (greater-than), 'LE' (less-or-equal-than), 'LT' (less-than) comparision_type: can be 'UNSIGNED', 'SIGNED', 'FLOAT'
(lhs, rhs, mode, comparison_type=None)
| 129 | |
| 130 | |
| 131 | def _compare(lhs, rhs, mode, comparison_type=None): |
| 132 | """ |
| 133 | mod: can be |
| 134 | 'EQ' (equal-to), |
| 135 | 'NE' (not equal-to), |
| 136 | 'GE' (greater-or-equal-than), |
| 137 | 'GT' (greater-than), |
| 138 | 'LE' (less-or-equal-than), |
| 139 | 'LT' (less-than) |
| 140 | comparision_type: can be 'UNSIGNED', 'SIGNED', 'FLOAT' |
| 141 | """ |
| 142 | lhs, rhs = _elemwise_dtype_promote(lhs, rhs) |
| 143 | lhs = HLOTensor(lhs) if not isinstance(lhs, HLOTensor) else lhs |
| 144 | rhs = HLOTensor(rhs) if not isinstance(rhs, HLOTensor) else rhs |
| 145 | oshape = _infer_elemwise_oshape([lhs.shape, rhs.shape]) |
| 146 | |
| 147 | lhs = lhs.broadcast_to(oshape) |
| 148 | rhs = rhs.broadcast_to(oshape) |
| 149 | |
| 150 | if comparison_type is None: |
| 151 | err_info = f"invalid dtype for compare {lhs.dtype} .vs {rhs.dtype}" |
| 152 | kind_map = {"i": "SIGNED", "u": "UNSIGNED", "f": "FLOAT"} |
| 153 | assert lhs.dtype.kind == rhs.dtype.kind, err_info |
| 154 | assert lhs.dtype.kind in kind_map, err_info |
| 155 | comparison_type = kind_map[lhs.dtype.kind] |
| 156 | |
| 157 | return HLOTensor( |
| 158 | hlo.CompareOp( |
| 159 | lhs.tensor, |
| 160 | rhs.tensor, |
| 161 | hlo.ComparisonDirectionAttr.get(mode), |
| 162 | compare_type=hlo.ComparisonTypeAttr.get(comparison_type), |
| 163 | ).result |
| 164 | ) |
| 165 | |
| 166 | |
| 167 | def _elemwise(hlo_op, inps): |
nothing calls this directly
no test coverage detected