| 77 | } |
| 78 | |
| 79 | ValueRefList reduce_rule( |
| 80 | const Reduce& reduce, Span<ValueRef> inputs, Span<bool> inputs_mask, |
| 81 | const Type<ScalarValue>& scalar_type) { |
| 82 | bool keepdim = reduce.keepdim; |
| 83 | auto axis = reduce.axis; |
| 84 | if (inputs.size() == 1) { |
| 85 | if (axis == INT_MAX || (inputs[0].shape()->ndim == 1 && keepdim == false)) { |
| 86 | // CompNode device = *inputs[0].device(); |
| 87 | return {scalar_type.make(imperative::apply(reduce, inputs)[0])}; |
| 88 | } |
| 89 | return imperative::apply(reduce, inputs); |
| 90 | } |
| 91 | mgb_assert(inputs.size() == 2); |
| 92 | bool is_scalar = is_scalar_shape(inputs[1]); |
| 93 | if (is_scalar) { |
| 94 | CompNode device = *inputs[0].device(); |
| 95 | return {scalar_type.make( |
| 96 | imperative::apply(reduce, inputs[0], make_scalar_shape(device))[0])}; |
| 97 | } |
| 98 | return imperative::apply(reduce, inputs); |
| 99 | } |
| 100 | |
| 101 | ValueRefList collective_comm_rule( |
| 102 | const CollectiveComm& collective_comm, Span<ValueRef> inputs, |
nothing calls this directly
no test coverage detected