Function where_grad_rule

imperative/python/src/grad_override.cpp:711–750

}

std::optional<ValueRefList> where_grad_rule(
        const OpDef& op, Span<ValueRef> inputs, Span<bool> inputs_require_grad,
        CustomBackward& backward) {
    auto&& where = op.cast_final_safe<Where>();
    auto&& param = where.param();
    mgb_assert(inputs.size() == 3);
    SmallVector<ValueRef> inps;
    if (inputs_require_grad[1] || inputs_require_grad[2]) {
        inps.push_back(inputs[0]);
    }
    bool data1_requires_grad = inputs_require_grad[1],
         data2_requires_grad = inputs_require_grad[2];
    auto maker = CustomGradMaker(backward, inputs.size());
    maker.output_size(1).output_captured(0, false);
    maker.backward([inputs = std::move(inps), data1_requires_grad, data2_requires_grad,
                    param](Span<ValueRef> grads) {
        mgb_assert(grads.size() == 1);
        ValueRef grad = grads[0];
        SmallVector<ValueRef> ret(3);
        if (!grad) {
            return ret;
        }
        if (data1_requires_grad == false && data2_requires_grad == false) {
            return ret;
        }

        auto&& grad_op = WhereBackward::make();
        ValueRefList args_(2);
        args_[0] = grads[0];
        args_[1] = inputs[0];
        auto back_grad = imperative::apply(*grad_op, args_);
        if (data1_requires_grad)
            ret[1] = back_grad[0];
        if (data2_requires_grad)
            ret[2] = back_grad[1];
        return ret;
    });
    maker.finalize();
    return imperative::apply(op, inputs);
}

struct Init {
    Init() {