(hlo_op, inps)
| 165 | |
| 166 | |
| 167 | def _elemwise(hlo_op, inps): |
| 168 | inps = _elemwise_dtype_promote(inps, hlo_op=hlo_op) |
| 169 | hinps = [HLOTensor(inp) if not isinstance(inp, HLOTensor) else inp for inp in inps] |
| 170 | |
| 171 | ishapes = [inp.shape for inp in hinps] |
| 172 | oshape = _infer_elemwise_oshape(ishapes) |
| 173 | |
| 174 | broadcasted_inps = [hinp.broadcast_to(oshape) for hinp in hinps] |
| 175 | results = hlo_op(*[binp.tensor for binp in broadcasted_inps]).results |
| 176 | assert len(results) == 1, f"elemwise op {hlo_op} should have only one output" |
| 177 | |
| 178 | return HLOTensor(results[0]) |
| 179 | |
| 180 | |
| 181 | def _elemwise_unary(hlo_op, a): |
no test coverage detected