Function subtensor_grad_rule

imperative/python/src/grad_override.cpp:316–378

}

std::optional<ValueRefList> subtensor_grad_rule(
        const OpDef& op, Span<ValueRef> inputs, Span<bool> inputs_require_grad,
        CustomBackward& backward) {
    auto&& subtensor = op.cast_final_safe<Subtensor>();
    auto&& grad_op = SetSubtensor::make(subtensor.items);
    SmallVector<ValueRef> inputs2;
    if (inputs_require_grad[0]) {
        inputs2.push_back(get_shape(inputs[0]));
        for (size_t i = 1; i < inputs.size(); ++i) {
            inputs2.push_back(inputs[i]);
        }
    }
    CompNodeValue::ref_t device = inputs[0].device();
    auto get_subtensor_index = [&](int idx) {
        HostTensorStorage storage(*device);
        storage.ensure_size(dtype::Int32().size());
        auto* ptr = reinterpret_cast<dt_int32*>(storage.ptr());
        ptr[0] = idx;
        return imperative::apply(
                CreateTensor(
                        CreateTensor::Unique, *device, dtype::Int32(), ValueShape({1})),
                HostStorage::make(storage))[0];
    };
    auto slice_items = subtensor.slice_items;
    auto items = subtensor.items;
    for (int i = 0; i < slice_items.size(); i++) {
        auto&& [axis, b_flag, e_flag, s_flag, idx_flag] = items[i];
        auto&& [b_val, e_val, s_val, ax_val] = slice_items[i];
        if (b_flag) {
            inputs2.push_back(get_subtensor_index(b_val));
        };
        if (e_flag) {
            inputs2.push_back(get_subtensor_index(e_val));
        };
        if (s_flag) {
            inputs2.push_back(get_subtensor_index(s_val));
        };
        if (idx_flag) {
            inputs2.push_back(get_subtensor_index(ax_val));
        };
    };
    auto maker = CustomGradMaker(backward, inputs.size());
    maker.output_size(1).output_captured(0, false);
    maker.backward([inputs = std::move(inputs2),
                    grad_op_ = std::move(grad_op)](Span<ValueRef> grads) {
        mgb_assert(grads.size() == 1);
        ValueRef grad = grads[0];
        SmallVector<ValueRef> ret(1);
        if (grad && inputs[0]) {
            ValueRefList args_(inputs.size() + 1);
            auto&& zeros = make_empty_tensor(grad.device(), inputs[0], grad.dtype());
            args_[0] = zeros;
            args_[1] = grad;
            for (size_t i = 1; i < inputs.size(); ++i) {
                args_[i + 1] = inputs[i];
            }
            ret[0] = imperative::apply(ApplyOp(*grad_op_), args_)[0];
        }