1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #include <memory> 16 17 #include "absl/container/flat_hash_set.h" 18 #include "absl/types/span.h" 19 #include "tensorflow/c/eager/abstract_tensor_handle.h" 20 #include "tensorflow/c/eager/c_api.h" 21 #include "tensorflow/c/eager/c_api_experimental.h" 22 #include "tensorflow/c/eager/c_api_unified_experimental.h" 23 #include "tensorflow/c/eager/c_api_unified_experimental_internal.h" 24 #include "tensorflow/c/eager/gradients.h" 25 #include "tensorflow/c/eager/gradients_internal.h" 26 #include "tensorflow/c/experimental/ops/array_ops.h" 27 #include "tensorflow/c/experimental/ops/math_ops.h" 28 #include "tensorflow/c/experimental/ops/nn_ops.h" 29 #include "tensorflow/c/tf_status_helper.h" 30 #include "tensorflow/c/tf_tensor.h" 31 #include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h" 32 #include "tensorflow/core/platform/errors.h" 33 #include "tensorflow/core/platform/types.h" 34 35 namespace tensorflow { 36 namespace gradients { 37 38 // Get a scalar TensorHandle with given value 39 Status ScalarTensorHandle(AbstractContext* ctx, float value, 40 AbstractTensorHandle** tensor); 41 42 // Get a TensorHandle with given float values and dimensions 43 Status TensorHandleWithDimsFloat(AbstractContext* ctx, float data[], 44 int64_t dims[], int num_dims, 45 AbstractTensorHandle** tensor); 46 47 // Get a TensorHandle with given int values and dimensions 48 Status TensorHandleWithDimsInt(AbstractContext* ctx, int data[], int64_t dims[], 49 int num_dims, AbstractTensorHandle** tensor); 50 51 // Places data from `t` into *result_tensor. 52 Status GetValue(AbstractTensorHandle* t, TF_Tensor** result_tensor); 53 54 // Util function that wraps an AbstractTensorHandle* with given data and dims. 55 AbstractTensorHandlePtr GetTensorHandleUtilFloat(AbstractContext* ctx, 56 float vals[], int64_t dims[], 57 int num_dims); 58 59 // Util function that wraps an AbstractTensorHandle* with given data and dims. 60 AbstractTensorHandlePtr GetTensorHandleUtilInt(AbstractContext* ctx, int vals[], 61 int64_t dims[], int num_dims); 62 63 // Util function that wraps an AbstractTensorHandle* with given data. 64 AbstractTensorHandlePtr GetScalarTensorHandleUtil(AbstractContext* ctx, 65 float val); 66 67 // Performs gradient update for each weight using given learning rate. 68 Status UpdateWeights(AbstractContext* ctx, 69 std::vector<AbstractTensorHandle*>& grads, 70 std::vector<AbstractTensorHandle*>& weights, 71 AbstractTensorHandle* learning_rate); 72 73 using Model = std::function<Status( 74 AbstractContext*, absl::Span<AbstractTensorHandle* const>, 75 absl::Span<AbstractTensorHandle*>, const GradientRegistry&)>; 76 77 // Runs given model in either graph or eager mode depending on value of 78 // use_function. 79 Status RunModel(Model model, AbstractContext* ctx, 80 absl::Span<AbstractTensorHandle* const> inputs, 81 absl::Span<AbstractTensorHandle*> outputs, bool use_function, 82 const GradientRegistry& registry); 83 84 // Builds context and returns inside *ctx. 85 Status BuildImmediateExecutionContext(bool use_tfrt, AbstractContext** ctx); 86 87 } // namespace gradients 88 } // namespace tensorflow 89