| 62 | } |
| 63 | |
| 64 | ValueRefList remove_axis_rule( |
| 65 | const RemoveAxis& remove_axis, Span<ValueRef> inputs, Span<bool> inputs_mask, |
| 66 | const Type<ScalarValue>& scalar_type) { |
| 67 | mgb_assert(!inputs_mask.item()); |
| 68 | bool is_scalar = inputs.item().shape()->ndim == remove_axis.axis.size(); |
| 69 | if (is_scalar && remove_axis.axis.size() == 1) { |
| 70 | return {scalar_type.make(inputs.item())}; |
| 71 | } |
| 72 | auto outputs = imperative::apply(remove_axis, inputs); |
| 73 | if (is_scalar) { |
| 74 | outputs[0] = scalar_type.make(outputs[0]); |
| 75 | } |
| 76 | return outputs; |
| 77 | } |
| 78 | |
| 79 | ValueRefList reduce_rule( |
| 80 | const Reduce& reduce, Span<ValueRef> inputs, Span<bool> inputs_mask, |