Checks and sets default value for initial gradients based on out_grads If output is the tensor whose size is greater than 1, out_grad's shape must be same as output's. If output is a scalar tensor, out_grad will also be a scaler or empty(will be initted to `oneflow.ones([1])`).
| 48 | // If output is a scalar tensor, out_grad will also be a scaler or empty(will be initted to |
| 49 | // `oneflow.ones([1])`). |
| 50 | Maybe<one::TensorTuple> CheckAndInitOutGrads(const one::TensorTuple& outputs, |
| 51 | const one::TensorTuple& out_grads, |
| 52 | bool is_grads_batched) { |
| 53 | size_t grad_size = out_grads.empty() ? outputs.size() : out_grads.size(); |
| 54 | auto gradients = std::make_shared<one::TensorTuple>(grad_size); |
| 55 | CHECK_EQ_OR_RETURN(outputs.size(), gradients->size()) |
| 56 | << "RuntimeError: got " << outputs.size() << " tensors and " << gradients->size() |
| 57 | << " gradients"; |
| 58 | for (int i = 0; i < outputs.size(); ++i) { |
| 59 | CHECK_OR_RETURN(outputs.at(i)->requires_grad()) |
| 60 | << "\nRuntimeError: element " << i |
| 61 | << " of tensors does not require grad and does not have a grad_fn"; |
| 62 | if (!outputs.at(i)->grad_fn_node()) { |
| 63 | CHECK_OR_RETURN(outputs.at(i)->is_leaf()) |
| 64 | << "output[" << i << "] doesn't have grad_fn and it is not leaf tensor!\n" |
| 65 | << "It is a bug with oneflow, please submit an issue on GitHub: " |
| 66 | "https://github.com/Oneflow-Inc/oneflow/issues"; |
| 67 | JUST(one::AddAccumulateFunctionNode(outputs.at(i))); |
| 68 | } |
| 69 | if (out_grads.empty() || !out_grads.at(i)) { |
| 70 | CHECK_OR_RETURN(IsScalarTensor(*outputs.at(i))) |
| 71 | << "Grad can be implicitly created only for scalar outputs"; |
| 72 | gradients->at(i) = JUST(one::functional::OnesLike(outputs.at(i))); |
| 73 | } else { |
| 74 | if (is_grads_batched) { |
| 75 | if (*(outputs.at(i)->shape()) != *JUST(out_grads.at(i)->shape()->Slice(1))) { |
| 76 | THROW(RuntimeError) << "If `is_grads_batched=True`, we interpret the first " |
| 77 | << "dimension of each grad_output as the batch dimension. " |
| 78 | << "The sizes of the remaining dimensions are expected to match " |
| 79 | << "the shape of corresponding output, but a mismatch " |
| 80 | << "was detected: grad_output[" << i << "] has a shape of " |
| 81 | << out_grads.at(i)->shape()->ToString() << " and output[" << i |
| 82 | << "] has a shape of " << outputs.at(i)->shape()->ToString() << "."; |
| 83 | } |
| 84 | |
| 85 | } else { |
| 86 | CHECK_EQ_OR_RETURN(*(outputs.at(i)->shape()), *(out_grads.at(i)->shape())) |
| 87 | << "out_grad's shape must be same as output's (" << outputs.at(i)->shape()->ToString() |
| 88 | << " vs " << out_grads.at(i)->shape()->ToString() << ")"; |
| 89 | } |
| 90 | if (JUST(oneflow::VectorAt(outputs, i))->dtype() |
| 91 | != JUST(oneflow::VectorAt(out_grads, i))->dtype()) { |
| 92 | JUST(oneflow::VectorAt(*gradients, i)) = |
| 93 | JUST(one::functional::Cast(out_grads[i], outputs[i]->dtype(), /*pin_memory=*/false)); |
| 94 | } else { |
| 95 | JUST(oneflow::VectorAt(*gradients, i)) = out_grads[i]; |
| 96 | } |
| 97 | } |
| 98 | } |
| 99 | if (LazyMode::is_enabled()) { JUST(MarkOutputGradients(outputs, *gradients)); } |
| 100 | return gradients; |
| 101 | } |
| 102 | |
| 103 | } // namespace |
| 104 |
no test coverage detected