mxnet
|
Class to hold custom operator registration. More...
#include <lib_api.h>
Public Member Functions | |
CustomOp (const char *op_name) | |
CustomOp & | setForward (fcomp_t fcomp, const char *ctx) |
CustomOp & | setBackward (fcomp_t fgrad, const char *ctx) |
CustomOp & | setParseAttrs (parseAttrs_t func) |
CustomOp & | setInferType (inferType_t func) |
CustomOp & | setInferSType (inferSType_t func) |
CustomOp & | setInferShape (inferShape_t func) |
CustomOp & | setMutateInputs (mutateInputs_t func) |
CustomOp & | setCreateOpState (createOpState_t func, const char *ctx) |
CustomOp & | setIsSubgraphOp () |
void | mapToVector () |
Public Attributes | |
const char * | name |
operator name More... | |
parseAttrs_t | parse_attrs |
operator functions More... | |
inferType_t | infer_type |
inferSType_t | infer_storage_type |
inferShape_t | infer_shape |
mutateInputs_t | mutate_inputs |
bool | isSGop |
std::vector< const char * > | forward_ctx_cstr |
vector repr of ctx map to be easily loaded from c_api More... | |
std::vector< const char * > | backward_ctx_cstr |
std::vector< const char * > | create_op_ctx_cstr |
std::vector< fcomp_t > | forward_fp |
std::vector< fcomp_t > | backward_fp |
std::vector< createOpState_t > | create_op_fp |
Class to hold custom operator registration.
|
explicit |
void mxnet::ext::CustomOp::mapToVector | ( | ) |
CustomOp& mxnet::ext::CustomOp::setCreateOpState | ( | createOpState_t | func, |
const char * | ctx | ||
) |
CustomOp& mxnet::ext::CustomOp::setInferShape | ( | inferShape_t | func | ) |
CustomOp& mxnet::ext::CustomOp::setInferSType | ( | inferSType_t | func | ) |
CustomOp& mxnet::ext::CustomOp::setInferType | ( | inferType_t | func | ) |
CustomOp& mxnet::ext::CustomOp::setIsSubgraphOp | ( | ) |
CustomOp& mxnet::ext::CustomOp::setMutateInputs | ( | mutateInputs_t | func | ) |
CustomOp& mxnet::ext::CustomOp::setParseAttrs | ( | parseAttrs_t | func | ) |
std::vector<const char*> mxnet::ext::CustomOp::backward_ctx_cstr |
std::vector<fcomp_t> mxnet::ext::CustomOp::backward_fp |
std::vector<const char*> mxnet::ext::CustomOp::create_op_ctx_cstr |
std::vector<createOpState_t> mxnet::ext::CustomOp::create_op_fp |
std::vector<const char*> mxnet::ext::CustomOp::forward_ctx_cstr |
vector repr of ctx map to be easily loaded from c_api
std::vector<fcomp_t> mxnet::ext::CustomOp::forward_fp |
inferShape_t mxnet::ext::CustomOp::infer_shape |
inferSType_t mxnet::ext::CustomOp::infer_storage_type |
inferType_t mxnet::ext::CustomOp::infer_type |
bool mxnet::ext::CustomOp::isSGop |
mutateInputs_t mxnet::ext::CustomOp::mutate_inputs |
const char* mxnet::ext::CustomOp::name |
operator name
parseAttrs_t mxnet::ext::CustomOp::parse_attrs |
operator functions