wehub-resource-sync
137 lines
5.3 KiB
C++
137 lines
5.3 KiB
C++
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#ifndef TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_
|
|
#define TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <unordered_set>
|
|
#include <vector>
|
|
|
|
#include "tensorflow/cc/framework/scope.h"
|
|
|
|
namespace tensorflow {
|
|
|
|
class ShapeRefiner;
|
|
|
|
// NewInternalScope returns a new scope which doesn't take ownership of
|
|
// graph, status, name_map, and refiner.
|
|
// This is intended to enable the C API (which are used by other language
|
|
// bindings) to create a Scope and access C++ functionality (i.e. gradients).
|
|
//
|
|
// Shape inference is disabled if `refiner` is nullptr.
|
|
Scope NewInternalScope(Graph* graph, absl::Status* status,
|
|
ShapeRefiner* refiner);
|
|
|
|
class Scope::Impl {
|
|
public:
|
|
// A NameMap is used to keep track of suffixes for names used in a scope. A
|
|
// name that has not been used so far in a scope will get no suffix. Later
|
|
// uses of the same name will get suffixes _1, _2, _3, etc. Multiple scopes
|
|
// can share the same NameMap. For instance, a new scope created using
|
|
// WithControlDependencies() would share the same NameMap with the parent.
|
|
typedef std::unordered_map<std::string, int> NameMap;
|
|
|
|
Impl(const std::shared_ptr<Graph>& graph,
|
|
const std::shared_ptr<absl::Status>& status,
|
|
const std::shared_ptr<NameMap>& name_map,
|
|
const std::shared_ptr<ShapeRefiner>& refiner);
|
|
|
|
const std::string& name() const { return name_; }
|
|
const std::vector<Operation>& control_deps() const { return control_deps_; }
|
|
|
|
private:
|
|
friend class Scope;
|
|
|
|
// Tag types to choose the constructor to dispatch.
|
|
struct Tags {
|
|
enum class ScopeName;
|
|
enum class OpName;
|
|
enum class ControlDeps;
|
|
enum class Device;
|
|
enum class SingleUseScope;
|
|
enum class ExitOnError;
|
|
enum class KernelLabel;
|
|
enum class Colocate;
|
|
enum class AssignedDevice;
|
|
enum class XlaCluster;
|
|
};
|
|
|
|
Impl(Graph* graph, absl::Status* status, NameMap* name_map,
|
|
ShapeRefiner* refiner, bool disable_shape_inference);
|
|
Impl(const Scope& other, Tags::ScopeName, const std::string& name,
|
|
bool copy_names);
|
|
Impl(const Scope& other, Tags::OpName, const std::string& name,
|
|
const std::string& op_name);
|
|
Impl(const Scope& other, Tags::ControlDeps,
|
|
std::vector<Operation> control_deps, bool clear_control_deps);
|
|
Impl(const Scope& other, Tags::Device, const std::string& device);
|
|
Impl(const Scope& other, Tags::SingleUseScope, const std::string& op_name);
|
|
Impl(const Scope& other, Tags::ExitOnError);
|
|
Impl(const Scope& other, Tags::KernelLabel, const std::string& kernel_label);
|
|
Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op,
|
|
bool clear_colocations);
|
|
Impl(const Scope& other, Tags::AssignedDevice,
|
|
const std::string& assigned_device);
|
|
Impl(const Scope& other, Tags::XlaCluster, const std::string& xla_cluster);
|
|
|
|
std::unordered_set<std::string> GetColocationConstraints(
|
|
const Operation& colocate_with_op) const;
|
|
|
|
// Helper functions to get a unique names.
|
|
std::string GetUniqueName(const std::string& prefix,
|
|
bool check_single_use) const;
|
|
std::string GetNameForOp(const std::string& default_name) const;
|
|
|
|
bool single_use_scope() const { return scope_used_ != nullptr; }
|
|
|
|
// The graph, status, and name maps are shared by all child scopes
|
|
// created from a single 'root' scope. A root scope is created by calling the
|
|
// Scope::NewRootScope function, which creates a new graph, a new status and
|
|
// the name maps.
|
|
std::shared_ptr<Graph> graph_ = nullptr;
|
|
std::shared_ptr<absl::Status> status_ = nullptr;
|
|
std::shared_ptr<NameMap> name_map_ = nullptr;
|
|
std::shared_ptr<ShapeRefiner> refiner_ = nullptr;
|
|
|
|
// If scope_used_ is not nullptr, op_name_ should be empty and
|
|
// GetUniqueNameForOp can only be called once on this scope. More calls to
|
|
// GetUniqueNameForOp will cause an error status to be set on this scope.
|
|
std::shared_ptr<bool> scope_used_ = nullptr;
|
|
|
|
const std::vector<Operation> control_deps_;
|
|
|
|
// The fully-qualified name of this scope (i.e. includes any parent scope
|
|
// names).
|
|
const std::string name_ = "";
|
|
const std::string op_name_ = "";
|
|
const bool exit_on_error_ = false;
|
|
const std::string kernel_label_ = "";
|
|
const std::string device_ = "";
|
|
const std::string assigned_device_ = "";
|
|
const std::string xla_cluster_ = "";
|
|
const std::unordered_set<std::string> colocation_constraints_;
|
|
|
|
// If true, Scope::DoShapeInference() always returns Status:OK().
|
|
// TODO(skyewm): remove this when possible
|
|
const bool disable_shape_inference_;
|
|
};
|
|
|
|
} // namespace tensorflow
|
|
|
|
#endif // TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_
|