1 /* Copyright 2022 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_CORE_GRAPPLER_OPTIMIZERS_DATA_SPLIT_UTILS_H_ 17 #define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_SPLIT_UTILS_H_ 18 19 #include <vector> 20 21 #include "absl/container/flat_hash_set.h" 22 #include "tensorflow/core/framework/function.h" 23 #include "tensorflow/core/platform/statusor.h" 24 25 namespace tensorflow { 26 namespace grappler { 27 namespace split_utils { 28 29 // Return value of `SplitFunction`, which is described below. 30 struct SplitResults { 31 FunctionDef first_function; 32 FunctionDef second_function; 33 std::vector<DataType> first_function_output_types; 34 }; 35 36 // Splits a FunctionDef into two FunctionDefs, called `first` and `second`, such 37 // that calling `function(*args)` is equivalent to calling 38 // `second(first(*args))`. The set `nodes_in_first_function` specifies nodes 39 // that are copied to `first`, and the other nodes are copied to `second`. Any 40 // edges from `first` to `second` will be represented by an output of `first` 41 // and a corresponding input of `second`. The caller must pass 42 // `nodes_in_first_function` such that there will not be any edges from `second` 43 // to `first`. 44 // 45 // For example, if you have the following function (using Python syntax): 46 // 47 // def f(x): 48 // y = tf.math.add(x, 1., name='add') 49 // return tf.multiply(y, 2, name='mul') 50 // 51 // Calling SplitFunction(f, {'add'}) results in: 52 // 53 // def first_function(x): 54 // return tf.math.add(x, 1., name='add') 55 // def second_function(y): 56 // return tf.multiply(y, 2, name='mul') 57 // 58 // The `num_captured_inputs` argument controls which arguments of `function` 59 // will be arguments of `second`. If it is zero, the only arguments of `second` 60 // are the outputs of `first`. If it is above zero, the last 61 // `num_caputured_inputs` arguments of `function` will also be arguments of 62 // `second`. 63 // 64 // Splitting functions in certain cases is unimplemented, in which case an 65 // Unimplemented status will be returned. Grappler passes must gracefully handle 66 // Unimplemented statuses without returning the error to its caller. 67 StatusOr<SplitResults> SplitFunction( 68 const FunctionDef& function, 69 const absl::flat_hash_set<absl::string_view>& nodes_in_first_function, 70 int64_t num_captured_inputs, const FunctionLibraryDefinition& library); 71 72 } // namespace split_utils 73 } // namespace grappler 74 } // namespace tensorflow 75 76 #endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_SPLIT_UTILS_H_ 77