| 138 | } |
| 139 | |
| 140 | ValueRefList add_axis_rule( |
| 141 | const AddAxis& add_axis, Span<ValueRef> inputs, Span<bool> inputs_mask, |
| 142 | const Type<ScalarValue>& scalar_type) { |
| 143 | mgb_assert(inputs.size() == 1); |
| 144 | if (inputs_mask.item()) { |
| 145 | mgb_assert(add_axis.axis[0] == 0); |
| 146 | if (add_axis.axis.size() == 1) { |
| 147 | return {inputs[0]}; |
| 148 | } else { |
| 149 | std::vector<int32_t> axis(add_axis.axis.begin() + 1, add_axis.axis.end()); |
| 150 | return imperative::apply(*AddAxis::make(axis, add_axis.scope()), inputs[0]); |
| 151 | } |
| 152 | } else { |
| 153 | return imperative::apply(add_axis, inputs); |
| 154 | } |
| 155 | } |
| 156 | |
| 157 | ValueRefList remote_recv_rule( |
| 158 | const RemoteRecv& remote_recv, Span<ValueRef> inputs, Span<bool> inputs_mask, |