Permalink
e59cb57 Dec 29, 2016
@tqchen @piiswrong
67 lines (58 sloc) 1.97 KB
/*!
* Copyright (c) 2016 by Contributors
* \file op_attr_types.h
* \brief Additional operator attributes
* beside the ones provided by NNVM
*/
#ifndef MXNET_OP_ATTR_TYPES_H_
#define MXNET_OP_ATTR_TYPES_H_
#include <mshadow/tensor.h>
#include <nnvm/op_attr_types.h>
#include <vector>
#include <functional>
#include "./base.h"
#include "./operator.h"
#include "./ndarray.h"
namespace mxnet {
using nnvm::NodeAttrs;
/*!
* \brief Create a Layer style, forward/backward operator.
* This is easy to write code that contains state.
*
* This is not the only way to register an op execution function.
* More simpler or specialized operator form can be registered
*
* \note Register under "FCreateLayerOp"
*/
using FCreateLayerOp = std::function<
Operator* (const NodeAttrs& n,
Context ctx,
const std::vector<TShape>& in_shape,
const std::vector<int>& in_type)>;
/*!
* \brief The resource request from the operator
*
* \note Register under "FResourceRequest"
*/
using FResourceRequest = std::function<
std::vector<ResourceRequest> (const NodeAttrs& n)>;
/*!
* \brief Register an operator called as a NDArray function
*
* \note Register under "FNDArrayFunction"
*/
using FNDArrayFunction = std::function<void (const nnvm::NodeAttrs& attrs,
const std::vector<NDArray>& inputs,
std::vector<NDArray>* outputs)>;
/*!
* \brief Resiger a compute function for simple stateless forward only operator
*
* \note Register under "FCompute<cpu>" and "FCompute<gpu>"
*/
using FCompute = std::function<void (const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs)>;
} // namespace mxnet
#endif // MXNET_OP_ATTR_TYPES_H_