Home | History | Annotate | Download | only in framework 
      1 /* Copyright 2015 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 // Class and associated machinery for specifying an Op's OpDef and shape
     17 // inference function for Op registration.
     18 
     19 #ifndef TENSORFLOW_FRAMEWORK_OP_DEF_BUILDER_H_
     20 #define TENSORFLOW_FRAMEWORK_OP_DEF_BUILDER_H_
     21 
     22 #include <string>
     23 #include <vector>
     24 #include "tensorflow/core/framework/op_def.pb.h"
     25 #include "tensorflow/core/lib/core/status.h"
     26 #include "tensorflow/core/lib/core/stringpiece.h"
     27 #include "tensorflow/core/platform/macros.h"
     28 
     29 namespace tensorflow {
     30 
     31 namespace shape_inference {
     32 class InferenceContext;
     33 }
     34 typedef std::function<Status(shape_inference::InferenceContext* c)>
     35     OpShapeInferenceFn;
     36 
     37 struct OpRegistrationData {
     38  public:
     39   OpRegistrationData() {}
     40   OpRegistrationData(const OpDef& def) : op_def(def) {}
     41   OpRegistrationData(const OpDef& def, const OpShapeInferenceFn& fn,
     42                      bool is_function = false)
     43       : op_def(def), shape_inference_fn(fn), is_function_op(is_function) {}
     44 
     45   OpDef op_def;
     46   OpShapeInferenceFn shape_inference_fn;
     47   bool is_function_op = false;
     48 };
     49 
     50 // Builder class passed to the REGISTER_OP() macro.
     51 class OpDefBuilder {
     52  public:
     53   // Constructs an OpDef with just the name field set.
     54   explicit OpDefBuilder(StringPiece op_name);
     55 
     56   // Adds an attr to this OpDefBuilder (and returns *this). The spec has
     57   // format "<name>:<type>" or "<name>:<type>=<default>"
     58   // where <name> matches regexp [a-zA-Z][a-zA-Z0-9_]*
     59   // (by convention only using capital letters for attrs that can be inferred)
     60   // <type> can be:
     61   //   "string", "int", "float", "bool", "type", "shape", or "tensor"
     62   //   "numbertype", "realnumbertype", "quantizedtype"
     63   //       (meaning "type" with a restriction on valid values)
     64   //   "{int32,int64}" or {realnumbertype,quantizedtype,string}"
     65   //       (meaning "type" with a restriction containing unions of value types)
     66   //   "{\"foo\", \"bar\n baz\"}", or "{'foo', 'bar\n baz'}"
     67   //       (meaning "string" with a restriction on valid values)
     68   //   "list(string)", ..., "list(tensor)", "list(numbertype)", ...
     69   //       (meaning lists of the above types)
     70   //   "int >= 2" (meaning "int" with a restriction on valid values)
     71   //   "list(string) >= 2", "list(int) >= 2"
     72   //       (meaning "list(string)" / "list(int)" with length at least 2)
     73   // <default>, if included, should use the Proto text format
     74   // of <type>.  For lists use [a, b, c] format.
     75   //
     76   // Note that any attr specifying the length of an input or output will
     77   // get a default minimum of 1 unless the >= # syntax is used.
     78   //
     79   // TODO(josh11b): Perhaps support restrictions and defaults as optional
     80   // extra arguments to Attr() instead of encoding them in the spec string.
     81   // TODO(josh11b): Would like to have better dtype handling for tensor attrs:
     82   // * Ability to say the type of an input/output matches the type of
     83   //   the tensor.
     84   // * Ability to restrict the type of the tensor like the existing
     85   //   restrictions for type attrs.
     86   // Perhaps by linking the type of the tensor to a type attr?
     87   OpDefBuilder& Attr(StringPiece spec);
     88 
     89   // Adds an input or output to this OpDefBuilder (and returns *this).
     90   // The spec has form "<name>:<type-expr>" or "<name>:Ref(<type-expr>)"
     91   // where <name> matches regexp [a-z][a-z0-9_]* and <type-expr> can be:
     92   // * For a single tensor: <type>
     93   // * For a sequence of tensors with the same type: <number>*<type>
     94   // * For a sequence of tensors with different types: <type-list>
     95   // Where:
     96   //   <type> is either one of "float", "int32", "string", ...
     97   //                 or the name of an attr (see above) with type "type".
     98   //   <number> is the name of an attr with type "int".
     99   //   <type-list> is the name of an attr with type "list(type)".
    100   // TODO(josh11b): Indicate Ref() via an optional argument instead of
    101   // in the spec?
    102   // TODO(josh11b): SparseInput() and SparseOutput() matching the Python
    103   // handling?
    104   OpDefBuilder& Input(StringPiece spec);
    105   OpDefBuilder& Output(StringPiece spec);
    106 
    107   // Turns on the indicated boolean flag in this OpDefBuilder (and
    108   // returns *this).
    109   OpDefBuilder& SetIsCommutative();
    110   OpDefBuilder& SetIsAggregate();
    111   OpDefBuilder& SetIsStateful();
    112   OpDefBuilder& SetAllowsUninitializedInput();
    113 
    114   // Deprecate the op at a certain GraphDef version.
    115   OpDefBuilder& Deprecated(int version, StringPiece explanation);
    116 
    117   // Adds docs to this OpDefBuilder (and returns *this).
    118   // Docs have the format:
    119   //   <1-line summary>
    120   //   <rest of the description>
    121   //   <name>: <description of name>
    122   //   <name>: <description of name>
    123   //     <if long, indent the description on subsequent lines>
    124   // Where <name> is the name of an attr, input, or output.  Please
    125   // wrap docs at 72 columns so that it may be indented in the
    126   // generated output.  For tensor inputs or outputs (not attrs), you
    127   // may start the description with an "=" (like name:= <description>)
    128   // to suppress the automatically-generated type documentation in
    129   // generated output.
    130 #ifndef TF_LEAN_BINARY
    131   OpDefBuilder& Doc(StringPiece text);
    132 #else
    133   OpDefBuilder& Doc(StringPiece text) { return *this; }
    134 #endif
    135 
    136   // Sets the shape function to be used for shape inference.
    137   //
    138   // Note that currently (October 2016), python code still requires a
    139   // RegisterShape call to invoke this; see call_cpp_shape_fn in
    140   // python/framework/common_shapes.py
    141   OpDefBuilder& SetShapeFn(Status (*fn)(shape_inference::InferenceContext*));
    142 
    143   // Sets op_reg_data->op_def to the requested OpDef and
    144   // op_reg_data->shape_inference_fn to the requested shape inference function,
    145   // or returns an error.
    146   // Must be called after all of the above methods.
    147   //
    148   // Note that OpDefBuilder only reports parsing errors.  You should also
    149   // call ValidateOpDef() to detect other problems.
    150   Status Finalize(OpRegistrationData* op_reg_data) const;
    151 
    152  private:
    153   OpDef* op_def() { return &op_reg_data_.op_def; }
    154 
    155   OpRegistrationData op_reg_data_;
    156   std::vector<string> attrs_;
    157   std::vector<string> inputs_;
    158   std::vector<string> outputs_;
    159   string doc_;
    160   std::vector<string> errors_;
    161 };
    162 
    163 }  // namespace tensorflow
    164 
    165 #endif  // TENSORFLOW_FRAMEWORK_OP_DEF_BUILDER_H_
    166