diff options
Diffstat (limited to 'tensorflow/core/framework/op.h')
| -rw-r--r-- | tensorflow/core/framework/op.h | 14 |
1 files changed, 11 insertions, 3 deletions
diff --git a/tensorflow/core/framework/op.h b/tensorflow/core/framework/op.h index 6e20a0fb4a..2a6a34f28e 100644 --- a/tensorflow/core/framework/op.h +++ b/tensorflow/core/framework/op.h @@ -16,7 +16,6 @@ limitations under the License. #ifndef TENSORFLOW_FRAMEWORK_OP_H_ #define TENSORFLOW_FRAMEWORK_OP_H_ -#include <functional> #include <unordered_map> #include "tensorflow/core/framework/op_def.pb.h" @@ -65,7 +64,7 @@ class OpRegistry : public OpRegistryInterface { // we defer calling func() until the first call to LookUp() or // Export() (if one of those has already been called, func() is // called immediately). - void Register(std::function<OpDef(void)> func); + void Register(const OpDef& op_def); const OpDef* LookUp(const string& op_type_name, Status* status) const override; @@ -81,6 +80,9 @@ class OpRegistry : public OpRegistryInterface { // A singleton available at startup. static OpRegistry* Global(); + // Get all registered ops. + void GetRegisteredOps(std::vector<OpDef>* op_defs); + private: // Ensures that all the functions in deferred_ get called, their OpDef's // registered, and returns with deferred_ empty. Returns true the first @@ -94,11 +96,17 @@ class OpRegistry : public OpRegistryInterface { mutable mutex mu_; // Functions in deferred_ may only be called with mu_ held. - mutable std::vector<std::function<OpDef(void)>> deferred_ GUARDED_BY(mu_); + mutable std::vector<OpDef> deferred_ GUARDED_BY(mu_); mutable std::unordered_map<string, OpDef*> registry_ GUARDED_BY(mu_); mutable bool initialized_ GUARDED_BY(mu_); }; +// Treats 'registry_ptr' as a pointer to OpRegistry, and calls +// registry_ptr->Register(op_def) for each op_def that has been registered with +// the current library's global op registry (obtained by calling +// OpRegistry::Global(). +extern "C" void RegisterOps(void* registry_ptr); + // Support for defining the OpDef (specifying the semantics of the Op and how // it should be created) and registering it in the OpRegistry::Global() // registry. Usage: |