Builder class passed to the REGISTER_OP() macro.
| 51 | |
| 52 | // Builder class passed to the REGISTER_OP() macro. |
| 53 | class OpDefBuilder { |
| 54 | public: |
| 55 | // Constructs an OpDef with just the name field set. |
| 56 | explicit OpDefBuilder(string op_name); |
| 57 | |
| 58 | // Adds an attr to this OpDefBuilder (and returns *this). The spec has |
| 59 | // format "<name>:<type>" or "<name>:<type>=<default>" |
| 60 | // where <name> matches regexp [a-zA-Z][a-zA-Z0-9_]* |
| 61 | // (by convention only using capital letters for attrs that can be inferred) |
| 62 | // <type> can be: |
| 63 | // "string", "int", "float", "bool", "type", "shape", or "tensor" |
| 64 | // "numbertype", "realnumbertype", "quantizedtype" |
| 65 | // (meaning "type" with a restriction on valid values) |
| 66 | // "{int32,int64}" or {realnumbertype,quantizedtype,string}" |
| 67 | // (meaning "type" with a restriction containing unions of value types) |
| 68 | // "{\"foo\", \"bar\n baz\"}", or "{'foo', 'bar\n baz'}" |
| 69 | // (meaning "string" with a restriction on valid values) |
| 70 | // "list(string)", ..., "list(tensor)", "list(numbertype)", ... |
| 71 | // (meaning lists of the above types) |
| 72 | // "int >= 2" (meaning "int" with a restriction on valid values) |
| 73 | // "list(string) >= 2", "list(int) >= 2" |
| 74 | // (meaning "list(string)" / "list(int)" with length at least 2) |
| 75 | // <default>, if included, should use the Proto text format |
| 76 | // of <type>. For lists use [a, b, c] format. |
| 77 | // |
| 78 | // Note that any attr specifying the length of an input or output will |
| 79 | // get a default minimum of 1 unless the >= # syntax is used. |
| 80 | // |
| 81 | // TODO(josh11b): Perhaps support restrictions and defaults as optional |
| 82 | // extra arguments to Attr() instead of encoding them in the spec string. |
| 83 | // TODO(josh11b): Would like to have better dtype handling for tensor attrs: |
| 84 | // * Ability to say the type of an input/output matches the type of |
| 85 | // the tensor. |
| 86 | // * Ability to restrict the type of the tensor like the existing |
| 87 | // restrictions for type attrs. |
| 88 | // Perhaps by linking the type of the tensor to a type attr? |
| 89 | OpDefBuilder& Attr(string spec); |
| 90 | |
| 91 | // Adds an input or output to this OpDefBuilder (and returns *this). |
| 92 | // The spec has form "<name>:<type-expr>" or "<name>:Ref(<type-expr>)" |
| 93 | // where <name> matches regexp [a-z][a-z0-9_]* and <type-expr> can be: |
| 94 | // * For a single tensor: <type> |
| 95 | // * For a sequence of tensors with the same type: <number>*<type> |
| 96 | // * For a sequence of tensors with different types: <type-list> |
| 97 | // Where: |
| 98 | // <type> is either one of "float", "int32", "string", ... |
| 99 | // or the name of an attr (see above) with type "type". |
| 100 | // <number> is the name of an attr with type "int". |
| 101 | // <type-list> is the name of an attr with type "list(type)". |
| 102 | // TODO(josh11b): Indicate Ref() via an optional argument instead of |
| 103 | // in the spec? |
| 104 | // TODO(josh11b): SparseInput() and SparseOutput() matching the Python |
| 105 | // handling? |
| 106 | OpDefBuilder& Input(string spec); |
| 107 | OpDefBuilder& Output(string spec); |
| 108 | |
| 109 | // Turns on the indicated boolean flag in this OpDefBuilder (and |
| 110 | // returns *this). |