| 954 | |
| 955 | template <typename Gradient, typename BackwardFunction, typename TapeTensor> |
| 956 | Status ForwardAccumulator<Gradient, BackwardFunction, TapeTensor>::Accumulate( |
| 957 | const string& op_type, const std::vector<TapeTensor>& input_tensors, |
| 958 | const std::vector<TapeTensor>& output_tensors, |
| 959 | gtl::ArraySlice<int64> input_tensor_id, |
| 960 | gtl::ArraySlice<tensorflow::DataType> input_dtypes, |
| 961 | const ForwardFunction<Gradient>* forward_function, |
| 962 | const std::function<BackwardFunction*()>& backward_function_getter, |
| 963 | const std::function<void(BackwardFunction*)>& backward_function_deleter) { |
| 964 | if (backward_tape_ != nullptr) { |
| 965 | // If backward_tape_ is not null, then this call to Accumulate is the result |
| 966 | // of a still-active call to Accumulate which is running operations. We |
| 967 | // forward these operations to backward_tape_ so the outer Accumulate call |
| 968 | // can do its work. |
| 969 | // |
| 970 | // Rather than re-entering and delegating Accumulate like this, we could |
| 971 | // instead allow ForwardAccumulator some control over the current tape set |
| 972 | // (so it can deactivate itself and activate its GradientTape). Currently |
| 973 | // that is managed by the language binding and would require relatively |
| 974 | // messy callbacks. |
| 975 | backward_tape_->RecordOperation(op_type, output_tensors, input_tensor_id, |
| 976 | input_dtypes, backward_function_getter, |
| 977 | backward_function_deleter); |
| 978 | return Status::OK(); |
| 979 | } |
| 980 | if (!ShouldRecord(input_tensor_id, input_dtypes)) { |
| 981 | return Status::OK(); |
| 982 | } |
| 983 | |
| 984 | // We may need to allocate zero inputs for trainable dtypes we don't have JVPs |
| 985 | // for. Make sure they get cleaned up. |
| 986 | std::vector<Gradient*> new_zeros; |
| 987 | auto delete_new_zeros = gtl::MakeCleanup([&new_zeros, this] { |
| 988 | for (Gradient* tensor : new_zeros) { |
| 989 | this->vspace_.DeleteGradient(tensor); |
| 990 | } |
| 991 | }); |
| 992 | std::vector<Gradient*> in_grads; |
| 993 | in_grads.reserve(input_tensors.size()); |
| 994 | for (size_t target_index = 0; target_index < input_tensors.size(); |
| 995 | ++target_index) { |
| 996 | const auto current_grad = |
| 997 | accumulated_gradients_.find(input_tensors[target_index].GetID()); |
| 998 | if (current_grad == accumulated_gradients_.end()) { |
| 999 | if (IsDtypeTrainable(input_tensors[target_index].GetDType())) { |
| 1000 | // ForwardAccumulator defaults to zeros for unwatched Tensors, unlike |
| 1001 | // GradientTape which uses ones. |
| 1002 | Gradient* zero = vspace_.Zeros(input_tensors[target_index]); |
| 1003 | new_zeros.push_back(zero); |
| 1004 | in_grads.push_back(zero); |
| 1005 | } else { |
| 1006 | in_grads.push_back(nullptr); |
| 1007 | } |
| 1008 | } else { |
| 1009 | in_grads.push_back(current_grad->second); |
| 1010 | } |
| 1011 | } |
| 1012 | |
| 1013 | // Avoid infinite recursion. Whichever forward function we run, it'll end up |
no test coverage detected