A class to record the schema of an op. More...
#include <operator_schema.h>
Public Types | |
typedef std::function< vector< TensorShape >const OperatorDef &, const vector< TensorShape > &)> | TensorInferenceFunctionType |
Public Member Functions | |
OpSchema (const string &file, const int line) | |
const string & | file () const |
Returns the file that the op schema is registered from. | |
int | line () const |
Returns the line in file that the op schema is registered from. | |
const char * | doc () const |
Returns the docstring of the op schema. | |
bool | Verify (const OperatorDef &def) const |
Verifies if an operator definition protobuf matches the pattern specified in the schema. | |
OpSchema & | NumInputs (int n) |
A single input. | |
OpSchema & | NumInputs (int min, int max) |
Input could be in range [min, max], inclusive. | |
OpSchema & | NumInputs (set< int > allowed_input_nums) |
Input could be one of the values specified in allowed_input_nums. | |
OpSchema & | NumInputs (std::function< bool(int)> func) |
Input is checked with a specified function. | |
OpSchema & | NumOutputs (int n) |
A single output. | |
OpSchema & | NumOutputs (int min, int max) |
Output could be in range [min, max], inclusive. | |
OpSchema & | NumOutputs (set< int > allowed_output_nums) |
Output could be one of the values specified in allowed_output_nums. | |
OpSchema & | NumOutputs (std::function< bool(int)> func) |
Output is checked with a specified function. | |
OpSchema & | NumInputsOutputs (std::function< bool(int, int)> func) |
Relationship between inputs and outputs is checked with a specified function. | |
OpSchema & | OutputCalculator (std::function< int(int)> calc) |
Set the output calculator to a user-defined function. | |
OpSchema & | SameNumberOfOutput () |
Set the number of outputs to be the same as the number of inputs. | |
OpSchema & | AllowInplace (std::function< bool(int, int)> inplace) |
OpSchema & | AllowInplace (set< std::pair< int, int >> inplace) |
OpSchema & | AllowOneToOneInplace () |
OpSchema & | EnforceInplace (std::function< bool(int, int)> inplace) |
OpSchema & | EnforceInplace (set< std::pair< int, int >> inplace) |
OpSchema & | EnforceOneToOneInplace () |
OpSchema & | TensorInferenceFunction (TensorInferenceFunctionType function) |
Sets the tensor inference function, which is a std::function object defined in operator_schema.h. | |
OpSchema & | IdenticalTypeAndShape () |
Seets the tensor inference function to produce the same output as the input. | |
OpSchema & | IdenticalTypeAndShapeOfInput (int idx) |
OpSchema & | IdenticalTypeAndShapeOfInputDim (int idx, int dim) |
OpSchema & | ScalarType (::caffe2::TensorProto_DataType dt) |
vector< TensorShape > | InferTensor (const OperatorDef &def, const vector< TensorShape > input_type_shape) const |
A function to allow one to infer the type and shape from the op schema. | |
OpSchema & | SetDoc (const string &doc) |
OpSchema & | Arg (const char *name, const char *description) |
OpSchema & | Input (const int n, const char *name, const char *description) |
OpSchema & | Output (const int n, const char *name, const char *description) |
OpSchema & | FillUsing (std::function< void(OpSchema &)> populator) |
int | CalculateOutput (int num_input) const |
A function to allow one to get the number of outputs based on the number of inputs, if this schema supports it. | |
const std::vector< std::pair< const char *, const char * > > & | arg_desc () |
const std::vector< std::pair< const char *, const char * > > & | input_desc () |
const std::vector< std::pair< const char *, const char * > > & | output_desc () |
Friends | |
std::ostream & | operator<< (std::ostream &out, const OpSchema &schema) |
A class to record the schema of an op.
OpSchema records the common interface of an op specified by its name. This is optional for each operator implemented in Caffe2 but is strongly recommended.
To register an OpSchema, one can use the macro OPERATOR_SCHEMA(name) and then append the various functions in the class. For example, for an op that itakes in two inputs, one output, and the first input and output could be in-place, can be written as
OPERATOR_SCHEMA(name) .NumInputs(2).NumOutputs(1).AllowInplace({{0, 0}});
Definition at line 36 of file operator_schema.h.