(
reducer, fidentity, inp, axes=None, keepdims=False, oshape=None, odtype=None
)
| 72 | |
| 73 | |
| 74 | def _reduce( |
| 75 | reducer, fidentity, inp, axes=None, keepdims=False, oshape=None, odtype=None |
| 76 | ): |
| 77 | def _reduce_nokeepdim(reducer, fidentity, inp, axes=None, oshape=None, odtype=None): |
| 78 | reduced_shape = _infer_reduce_shape(inp.shape, axes) |
| 79 | |
| 80 | _check_shape(reduced_shape, oshape) |
| 81 | _check_dtype(inp.dtype, odtype) |
| 82 | |
| 83 | reduce_out = ir_utils.make_ir_type_according_meta(reduced_shape, inp.dtype) |
| 84 | init_val = ir_utils.ir_constant_tuple(fidentity(inp.dtype)) |
| 85 | reduce_op = hlo.ReduceOp( |
| 86 | [reduce_out], [inp.tensor], init_val, ir_utils.dense_int_elements(axes) |
| 87 | ) |
| 88 | scalar_type = ir_utils.make_ir_type_according_meta(tuple(), inp.dtype) |
| 89 | reducer_region = reduce_op.regions[0].blocks.append(scalar_type, scalar_type) |
| 90 | with ir.InsertionPoint(reducer_region): |
| 91 | reducer_ret = reducer(*reducer_region.arguments) |
| 92 | hlo.ReturnOp(reducer_ret.results) |
| 93 | |
| 94 | return HLOTensor(reduce_op.result) |
| 95 | |
| 96 | axes = _normalize_reduce_axes(inp.shape, axes) |
| 97 | maykeepdim_shape = _infer_reduce_shape(inp.shape, axes, keepdims) |
| 98 | _check_shape(maykeepdim_shape, oshape) |
| 99 | |
| 100 | oup = _reduce_nokeepdim(reducer, fidentity, inp, axes, oshape, odtype) |
| 101 | if _shape_equal(oup.shape, maykeepdim_shape): |
| 102 | return oup |
| 103 | else: |
| 104 | return reshape(oup, maykeepdim_shape) |
| 105 | |
| 106 | |
| 107 | sum = partial(_reduce, hlo.AddOp, _get_sum_identity) |
no test coverage detected