Method apply_transformation

imperative/src/impl/transformations/grad.cpp:295–556

}

ValueRefList GradTransformation::apply_transformation(
        const Operator& op, Span<ValueRef> inputs) {
    auto fallback = [&] {
        SmallVector<ValueRef> unwrapped_inputs(inputs.size());
        {
            // overhead
            for (size_t i = 0; i < inputs.size(); ++i) {
                if (auto&& grad_value = as_grad_value(inputs[i])) {
                    unwrapped_inputs[i] = grad_value->m_value;
                } else {
                    unwrapped_inputs[i] = inputs[i];
                }
            }
        }
        return imperative::apply(op, unwrapped_inputs);
    };
    if (op.is<GetAttr>()) {
        // overhead
        if (auto&& grad_value = as_grad_value(inputs.item())) {
            return imperative::apply(op, grad_value->m_value);
        } else {
            return imperative::apply(op, inputs);
        }
    }
    if (m_suppressed) {
        return fallback();
    }
    if (auto* op_val = op.as<ApplyOp>()) {
        size_t nr_require_grad = 0;
        SmallVector<bool> require_grads(inputs.size());
        for (size_t i = 0; i < inputs.size(); ++i) {
            if (is_grad_value(inputs[i])) {
                nr_require_grad++;
                require_grads[i] = true;
            } else {
                require_grads[i] = false;
            }
        }
        if (nr_require_grad == 0) {
            return imperative::apply(op, inputs);
        }
        SmallVector<ValueRef> captured_inputs(inputs.size());
        SmallVector<bool> inputs_require_grad(inputs.size());
        // capture value so that trace could assume input as same
        auto capture_value = [](const ValueRef& value) {
            // TODO: fastpath copy shouldn't be an OpDef
            static auto fastpath_copy = FastpathCopy::make();
            return imperative::apply(ApplyOp(*fastpath_copy), value)[0];
        };
        for (size_t i = 0; i < inputs.size(); ++i) {
            auto& input = inputs[i];
            if (auto&& grad_value = as_grad_value(input)) {
                captured_inputs[i] = capture_value(grad_value->m_value);
                inputs_require_grad[i] = true;
            } else {
                captured_inputs[i] = capture_value(input);
                inputs_require_grad[i] = false;
            }