| 498 | } |
| 499 | |
| 500 | std::optional<ValueRefList> pixelShuffle_grad_rule( |
| 501 | const OpDef& op, Span<ValueRef> inputs, Span<bool> inputs_require_grad, |
| 502 | CustomBackward& backward) { |
| 503 | auto&& pixelShuffle = op.cast_final_safe<PixelShuffle>(); |
| 504 | mgb_assert(inputs.size() == 1); |
| 505 | bool flag = inputs_require_grad[0]; |
| 506 | auto&& grad_op = PixelShuffleBackward::make(pixelShuffle.factor); |
| 507 | auto maker = CustomGradMaker(backward, inputs.size()); |
| 508 | maker.output_size(1).output_captured(0, false); |
| 509 | maker.backward([grad_op_ = std::move(grad_op), flag_ = flag](Span<ValueRef> grads) { |
| 510 | mgb_assert(grads.size() == 1); |
| 511 | ValueRef grad = grads[0]; |
| 512 | SmallVector<ValueRef> ret(1); |
| 513 | if (grad && flag_) { |
| 514 | ret[0] = imperative::apply(*grad_op_, grad)[0]; |
| 515 | } |
| 516 | return ret; |
| 517 | }); |
| 518 | maker.finalize(); |
| 519 | return imperative::apply(op, inputs); |
| 520 | } |
| 521 | |
| 522 | std::optional<ValueRefList> indexing_grad_rule( |
| 523 | const OpDef& op, Span<ValueRef> inputs, Span<bool> inputs_require_grad, |
nothing calls this directly
no test coverage detected