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 16 #ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_FLAT_TENSOR_FUNCTION_H_ 17 #define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_FLAT_TENSOR_FUNCTION_H_ 18 19 #include <functional> 20 #include <memory> 21 #include <string> 22 #include <vector> 23 24 #include "tensorflow/c/eager/immediate_execution_context.h" 25 #include "tensorflow/c/eager/immediate_execution_operation.h" 26 #include "tensorflow/c/eager/immediate_execution_tensor_handle.h" 27 #include "tensorflow/core/framework/function.pb.h" 28 #include "tensorflow/core/protobuf/saved_object_graph.pb.h" 29 30 namespace tensorflow { 31 32 // FlatTensorFunction models a TF2 eager runtime view of a callable function, 33 // taking + returning flat lists of tensors, including any captures. 34 // Effectively, it is a thin wrapper around a FunctionDef owned by the 35 // EagerContext, and any TensorHandle captures associated with the function. The 36 // MakeCallOp method handles the logic of marshaling captures after the user 37 // provided inputs automatically. 38 // Note(bmzhao): This class is mainly intended to house low-level reusable 39 // function logic between SignatureDefFunction and ConcreteFunction, which 40 // present higher level interfaces. This type does *not* hold any "function 41 // metadata". 42 class FlatTensorFunction { 43 public: 44 // Factory for creating a FlatTensorFunction. 45 // 46 // Params: 47 // function_def - The function_def associated with the created 48 // FlatTensorFunction. FlatTensorFunction will register this 49 // function_def with `ctx` on creation, and de-register it on 50 // destruction. function_def must be non-null, but 51 // otherwise has no lifetime requirements. 52 // captures - The captured TensorHandles associated with this 53 // FlatTensorFunction. FlatTensorFunction will participate in 54 // ownership of the handles (it explicitly increments the refcount 55 // of each handle, and will decrement them on destruction). 56 // ctx - A handle to the Tensorflow runtime. This MUST be non-null and 57 // outlive TFConcreteFunction. 58 // out - The output FlatTensorFunction. 59 static Status Create(const FunctionDef* function_def, 60 std::vector<ImmediateExecutionTensorHandle*> captures, 61 ImmediateExecutionContext* ctx, 62 std::unique_ptr<FlatTensorFunction>* out); 63 64 // This method creates a "Call" Op used to execute the function. 65 Status MakeCallOp(absl::Span<AbstractTensorHandle* const> inputs, 66 ImmediateOpPtr* out) const; 67 68 ~FlatTensorFunction(); 69 70 private: 71 FlatTensorFunction(const std::string& name, 72 std::vector<ImmediateTensorHandlePtr> captures, 73 ImmediateExecutionContext* ctx); 74 75 FlatTensorFunction(const FlatTensorFunction&) = delete; 76 FlatTensorFunction& operator=(const FlatTensorFunction&) = delete; 77 78 // Name of the FunctionDef corresponding to this TFConcreteFunction 79 std::string name_; 80 std::vector<ImmediateTensorHandlePtr> captures_; 81 ImmediateExecutionContext* ctx_; 82 }; 83 84 } // namespace tensorflow 85 86 #endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_FLAT_TENSOR_FUNCTION_H_ 87