Class FunctionLibraryDefinition

tensorflow/core/framework/function.h:314–505

Helper to maintain a map between function names in a given FunctionDefLibrary and function definitions. This class is thread-safe.

//
// This class is thread-safe.
class FunctionLibraryDefinition : public OpRegistryInterface {
 public:
  // Ops created for function arguments bear the name given by `kArgOp`; those
  // created for return values bear the name given by `kRetOp`.
  static constexpr const char* const kArgOp = "_Arg";
  static constexpr const char* const kDeviceArgOp = "_DeviceArg";
  static constexpr const char* const kRetOp = "_Retval";
  static constexpr const char* const kDeviceRetOp = "_DeviceRetval";
  static constexpr const char* const kIntsOnDeviceAttr =
      "experimental_ints_on_device";

  static constexpr const char* const kGradientOp = "SymbolicGradient";
  static constexpr const char* const kFuncAttr = "f";

  // Note: This constructor grabs `lib_def`'s lock in shared mode.
  FunctionLibraryDefinition(const FunctionLibraryDefinition& lib_def);
  FunctionLibraryDefinition(const OpRegistryInterface* default_registry,
                            const FunctionDefLibrary& lib_def);
  ~FunctionLibraryDefinition() override;

  FunctionLibraryDefinition& operator=(const FunctionLibraryDefinition&) =
      delete;

  // Returns True if the library contains `func`, False otherwise.
  bool Contains(const string& func) const;

  // Returns nullptr if "func" is not defined in "lib_def". Otherwise,
  // returns its definition proto.
  //
  // NB: This function returns a borrowed pointer, which can be invalidated by a
  // subsequent call to `ReplaceFunction()` with the given name.
  const FunctionDef* Find(const string& func) const LOCKS_EXCLUDED(mu_);

  // Adds function definition 'fdef' to this function library.
  // Returns status 'ok' on success, or error otherwise. This is a no-op if
  // 'fdef' already exists in this function library.
  // If 'fdef' is successfully added to the library, it will be accessible
  // from 'LookUp' and included in the proto returned by 'ToProto'.
  // This operation is atomic.
  Status AddFunctionDef(const FunctionDef& fdef) LOCKS_EXCLUDED(mu_);

  // Adds gradient definition 'grad' to this function library.
  // This is a no-op if 'grad' already exists in this function library.
  // If 'grad' is successfully added, it will be accessible via 'FindGradient'
  // and included in the proto returned by 'ToProto'.
  // This operation is atomic.
  Status AddGradientDef(const GradientDef& grad) LOCKS_EXCLUDED(mu_);

  // Replaces the function corresponding to `func` with `fdef`. Returns
  // a non-OK status if "func" was not found in the library, OK otherwise.
  // Please be careful when replacing function: make sure all previous pointers
  // returned by `Find()` are no longer in use.
  Status ReplaceFunction(const string& func, const FunctionDef& fdef)
      LOCKS_EXCLUDED(mu_);

  // Replaces the gradient corresponding to `grad.function_name()`. Returns
  // a non-OK status if "grad.function_name()" was not found in the library, OK
  // otherwise.