Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
arange_op.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_OPERATORS_ARRAY_ARGMAX_OP_H_
14 #define DRAGON_OPERATORS_ARRAY_ARGMAX_OP_H_
15 
16 #include "core/operator.h"
17 
18 namespace dragon {
19 
20 template <class Context>
21 class ArangeOp final : public Operator<Context> {
22  public:
23  ArangeOp(const OperatorDef& def, Workspace* ws)
24  : Operator<Context>(def, ws) {
25  GET_ARG_WITH_DESC(int64_t, start, 0);
26  GET_ARG_WITH_DESC(int64_t, stop, 0);
27  GET_ARG_WITH_DESC(int64_t, step, 1);
28  }
30 
31  void RunOnDevice() override;
32  template <typename T> void RunImpl();
33 
34  protected:
35  int64_t astart_, astop_, astep_, dim_;
36  DECLARE_ARG_WITH_DESC(int64_t, start);
37  DECLARE_ARG_WITH_DESC(int64_t, stop);
38  DECLARE_ARG_WITH_DESC(int64_t, step);
39 };
40 
41 DEFINE_ARG_WITH_DESC(int64_t, ArangeOp, start);
42 DEFINE_ARG_WITH_DESC(int64_t, ArangeOp, stop);
43 DEFINE_ARG_WITH_DESC(int64_t, ArangeOp, step);
44 
45 } // namespace dragon
46 
47 #endif // DRAGON_OPERATORS_ARRAY_ARANGE_OP_H_
Definition: arange_op.h:21
DECLARE_ARG_WITH_DESC(int64_t, start)
Definition: workspace.h:20
USE_OPERATOR_FUNCTIONS
Definition: arange_op.h:29
const OperatorDef & def() const
Return the stored def.
Definition: operator.h:114
Workspace * ws() const
Return the parent workspace.
Definition: operator.h:87
ArangeOp(const OperatorDef &def, Workspace *ws)
Definition: arange_op.h:23
Definition: operator.h:149
int64_t astep_
Definition: arange_op.h:35
void RunImpl()
Definition: arange_op.cc:8
int64_t dim_
Definition: arange_op.h:35
DEFINE_ARG_WITH_DESC(float, DropoutOp, prob)
#define GET_ARG_WITH_DESC(type, arg, default_value)
Definition: operator.h:408
int64_t astart_
Definition: arange_op.h:35
int64_t astop_
Definition: arange_op.h:35
Definition: common.h:41
void RunOnDevice() override
Implement the detailed execution.
Definition: arange_op.cc:14