You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by jx...@apache.org on 2018/02/26 18:36:59 UTC
[incubator-mxnet] branch master updated: eye operator,
for default storage type (#9770)
This is an automated email from the ASF dual-hosted git repository.
jxie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new db24ac1 eye operator, for default storage type (#9770)
db24ac1 is described below
commit db24ac102cb857d768d6cfd32858504c6327e3b8
Author: Ziyue Huang <zy...@gmail.com>
AuthorDate: Tue Feb 27 02:36:51 2018 +0800
eye operator, for default storage type (#9770)
* eye
* more test
* change to two kernels
* address comments
* Update init_op.h
---
python/mxnet/ndarray/ndarray.py | 41 +++++++++++++++-
python/mxnet/symbol/symbol.py | 25 +++++++++-
src/operator/tensor/init_op.cc | 11 +++++
src/operator/tensor/init_op.cu | 3 ++
src/operator/tensor/init_op.h | 89 +++++++++++++++++++++++++++++++++++
tests/python/unittest/test_ndarray.py | 9 ++++
6 files changed, 176 insertions(+), 2 deletions(-)
diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py
index 4c5273f..59389dd 100644
--- a/python/mxnet/ndarray/ndarray.py
+++ b/python/mxnet/ndarray/ndarray.py
@@ -43,7 +43,7 @@ from . import op
from ._internal import NDArrayBase
__all__ = ["NDArray", "concatenate", "_DTYPE_NP_TO_MX", "_DTYPE_MX_TO_NP", "_GRAD_REQ_MAP",
- "ones", "add", "arange", "divide", "equal", "full", "greater", "greater_equal",
+ "ones", "add", "arange", "eye", "divide", "equal", "full", "greater", "greater_equal",
"imdecode", "lesser", "lesser_equal", "maximum", "minimum", "moveaxis", "modulo",
"multiply", "not_equal", "onehot_encode", "power", "subtract", "true_divide",
"waitall", "_new_empty_handle"]
@@ -3411,6 +3411,45 @@ def zeros(shape, ctx=None, dtype=None, **kwargs):
return _internal._zeros(shape=shape, ctx=ctx, dtype=dtype, **kwargs)
# pylint: enable= no-member, protected-access
+def eye(N, M=0, k=0, ctx=None, dtype=None, **kwargs):
+ """Return a 2-D array with ones on the diagonal and zeros elsewhere.
+ Parameters
+ ----------
+ N: int
+ Number of rows in the output.
+ M: int, optional
+ Number of columns in the output. If 0, defaults to N.
+ k: int, optional
+ Index of the diagonal: 0 (the default) refers to the main diagonal,
+ a positive value refers to an upper diagonal,
+ and a negative value to a lower diagonal.
+ ctx: Context, optional
+ An optional device context (default is the current default context)
+ dtype: str or numpy.dtype, optional
+ An optional value type (default is `float32`)
+ Returns
+ -------
+ NDArray
+ A created array
+ Examples
+ --------
+ >>> mx.nd.eye(2)
+ [[ 1. 0.]
+ [ 0. 1.]]
+ <NDArray 2x2 @cpu(0)>
+ >>> mx.nd.eye(2, 3, 1)
+ [[ 0. 1. 0.]
+ [ 0. 0. 1.]]
+ <NDArray 2x3 @cpu(0)>
+ """
+ # pylint: disable= unused-argument
+ if ctx is None:
+ ctx = Context.default_ctx
+ dtype = mx_real_t if dtype is None else dtype
+ # pylint: disable= no-member, protected-access
+ return _internal._eye(N=N, M=M, k=k, ctx=ctx, dtype=dtype, **kwargs)
+ # pylint: enable= no-member, protected-access
+
def empty(shape, ctx=None, dtype=None):
"""Returns a new array of given shape and type, without initializing entries.
diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
index 220e1bd..59402e6 100644
--- a/python/mxnet/symbol/symbol.py
+++ b/python/mxnet/symbol/symbol.py
@@ -47,7 +47,7 @@ from . import op
from ._internal import SymbolBase, _set_symbol_class
__all__ = ["Symbol", "var", "Variable", "Group", "load", "load_json",
- "pow", "maximum", "minimum", "hypot", "zeros", "ones", "full", "arange"]
+ "pow", "maximum", "minimum", "hypot", "eye", "zeros", "ones", "full", "arange"]
class Symbol(SymbolBase):
@@ -2731,6 +2731,29 @@ def hypot(left, right):
else:
raise TypeError('types (%s, %s) not supported' % (str(type(left)), str(type(right))))
+def eye(N, M=0, k=0, dtype=None, **kwargs):
+ """Returns a new symbol of 2-D shpae, filled with ones on the diagonal
+ and zeros elsewhere.
+ Parameters
+ ----------
+ N: int
+ Number of rows in the output.
+ M: int, optional
+ Number of columns in the output. If 0, defaults to N.
+ k: int, optional
+ Index of the diagonal: 0 (the default) refers to the main diagonal,
+ a positive value refers to an upper diagonal,
+ and a negative value to a lower diagonal.
+ dtype : str or numpy.dtype, optional
+ The value type of the inner value, default to ``np.float32``.
+ Returns
+ -------
+ out : Symbol
+ The created Symbol.
+ """
+ if dtype is None:
+ dtype = _numpy.float32
+ return _internal._eye(N, M, k, dtype=dtype, **kwargs)
def zeros(shape, dtype=None, **kwargs):
"""Returns a new symbol of given shape and type, filled with zeros.
diff --git a/src/operator/tensor/init_op.cc b/src/operator/tensor/init_op.cc
index 5be4a2f..d5a328e 100644
--- a/src/operator/tensor/init_op.cc
+++ b/src/operator/tensor/init_op.cc
@@ -31,6 +31,7 @@ namespace op {
DMLC_REGISTER_PARAMETER(InitOpParam);
DMLC_REGISTER_PARAMETER(InitOpWithScalarParam);
DMLC_REGISTER_PARAMETER(RangeParam);
+DMLC_REGISTER_PARAMETER(EyeParam);
NNVM_REGISTER_OP(_zeros)
@@ -45,6 +46,16 @@ NNVM_REGISTER_OP(_zeros)
.set_attr<FComputeEx>("FComputeEx<cpu>", FillComputeZerosEx<cpu>)
.add_arguments(InitOpParam::__FIELDS__());
+NNVM_REGISTER_OP(_eye)
+.describe("Return a 2-D array with ones on the diagonal and zeros elsewhere.")
+.set_num_inputs(0)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<EyeParam>)
+.set_attr<nnvm::FInferShape>("FInferShape", InitEyeShape<EyeParam>)
+.set_attr<nnvm::FInferType>("FInferType", InitType<EyeParam>)
+.set_attr<FCompute>("FCompute<cpu>", EyeFill<cpu>)
+.add_arguments(EyeParam::__FIELDS__());
+
NNVM_REGISTER_OP(_ones)
.describe("fill target with ones")
.set_num_inputs(0)
diff --git a/src/operator/tensor/init_op.cu b/src/operator/tensor/init_op.cu
index 37660e1..5841408fa 100644
--- a/src/operator/tensor/init_op.cu
+++ b/src/operator/tensor/init_op.cu
@@ -48,6 +48,9 @@ NNVM_REGISTER_OP(_zeros)
.set_attr<FCompute>("FCompute<gpu>", FillCompute<gpu, 0>)
.set_attr<FComputeEx>("FComputeEx<gpu>", FillComputeZerosEx<gpu>);
+NNVM_REGISTER_OP(_eye)
+.set_attr<FCompute>("FCompute<gpu>", EyeFill<gpu>);
+
NNVM_REGISTER_OP(_ones)
.set_attr<FCompute>("FCompute<gpu>", FillCompute<gpu, 1>);
diff --git a/src/operator/tensor/init_op.h b/src/operator/tensor/init_op.h
index 0b9e6a2..475409e 100644
--- a/src/operator/tensor/init_op.h
+++ b/src/operator/tensor/init_op.h
@@ -32,6 +32,7 @@
#include <dmlc/optional.h>
#include <vector>
#include <string>
+#include <algorithm>
#include <limits>
#include "../mshadow_op.h"
#include "../elemwise_op_common.h"
@@ -60,6 +61,62 @@ struct InitOpParam : public dmlc::Parameter<InitOpParam> {
}
};
+struct EyeParam : public dmlc::Parameter<EyeParam> {
+ nnvm::dim_t N;
+ nnvm::dim_t M;
+ nnvm::dim_t k;
+ std::string ctx;
+ int dtype;
+
+ DMLC_DECLARE_PARAMETER(EyeParam) {
+ DMLC_DECLARE_FIELD(N)
+ .describe("Number of rows in the output.");
+ DMLC_DECLARE_FIELD(M)
+ .set_default(0)
+ .describe("Number of columns in the output. If 0, defaults to N");
+ DMLC_DECLARE_FIELD(k)
+ .set_default(0)
+ .describe("Index of the diagonal. 0 (the default) refers to the main diagonal."
+ "A positive value refers to an upper diagonal."
+ "A negative value to a lower diagonal.");
+ DMLC_DECLARE_FIELD(ctx)
+ .set_default("")
+ .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
+ "Only used for imperative calls.");
+ DMLC_DECLARE_FIELD(dtype).set_default(mshadow::kFloat32)
+ .add_enum("float32", mshadow::kFloat32)
+ .add_enum("float64", mshadow::kFloat64)
+ .add_enum("float16", mshadow::kFloat16)
+ .add_enum("uint8", mshadow::kUint8)
+ .add_enum("int32", mshadow::kInt32)
+ .add_enum("int64", mshadow::kInt64)
+ .describe("Target data type.");
+ }
+};
+
+template<typename ParamType>
+inline bool InitEyeShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape> *in_attrs,
+ std::vector<TShape> *out_attrs) {
+ const ParamType& param = nnvm::get<ParamType>(attrs.parsed);
+ CHECK_EQ(in_attrs->size(), 0U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::Shape2(param.N, param.M > 0 ? param.M : param.N));
+ return true;
+}
+
+template<int req>
+struct eye_dns_fill {
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data,
+ const nnvm::dim_t init_col,
+ const nnvm::dim_t k,
+ const nnvm::dim_t num_cols) {
+ KERNEL_ASSIGN(out_data[(i+init_col-k)*num_cols+i+init_col], req, static_cast<DType>(1));
+ }
+};
+
+
struct RangeParam : public dmlc::Parameter<RangeParam> {
double start;
dmlc::optional<double> stop;
@@ -336,6 +393,38 @@ void FillComputeZerosEx(const nnvm::NodeAttrs& attrs,
}
}
+
+template<typename xpu>
+void EyeFill(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ CHECK_EQ(inputs.size(), 0U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ const EyeParam& param = nnvm::get<EyeParam>(attrs.parsed);
+ const TBlob& out_data = outputs[0];
+ const nnvm::dim_t num_cols = param.M > 0 ? param.M : param.N;
+
+ const nnvm::dim_t cnnz = std::max(num_cols - std::abs(param.k), (nnvm::dim_t)0);
+ const nnvm::dim_t rnnz = std::max(param.N - std::abs(param.k), (nnvm::dim_t)0);
+ const nnvm::dim_t nnz = param.k > 0 ? std::min(cnnz, param.N) :
+ std::min(rnnz, num_cols);
+ using namespace mxnet_op;
+ MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
+ Fill(s, out_data, req[0], static_cast<DType>(0));
+ if (nnz > 0) {
+ Kernel<eye_dns_fill<req_type>, xpu>::Launch(s, nnz, out_data.dptr<DType>(),
+ std::max(static_cast<nnvm::dim_t>(0), param.k), param.k, num_cols);
+ }
+ });
+ });
+}
+
+
struct range_fwd {
template<typename DType>
MSHADOW_XINLINE static void Map(int i, int repeat, DType start, DType step,
diff --git a/tests/python/unittest/test_ndarray.py b/tests/python/unittest/test_ndarray.py
index 6c10487..52697c2 100644
--- a/tests/python/unittest/test_ndarray.py
+++ b/tests/python/unittest/test_ndarray.py
@@ -836,6 +836,15 @@ def test_output():
assert_almost_equal(out.asnumpy(), ones.asnumpy() * 2)
arange_out = mx.nd.arange(0, 20, dtype='int64')
assert_almost_equal(arange_out.asnumpy(), np.arange(0, 20))
+ N_array = np.random.randint(1, high=8, size=10)
+ M_array = np.random.randint(1, high=8, size=10)
+ k_array = np.random.randint(-10, high=10, size=10)
+ for i in range(10):
+ N = N_array[i]
+ M = M_array[i]
+ k = k_array[i]
+ assert_almost_equal(np.eye(N, M, k), mx.nd.eye(N, M, k).asnumpy())
+ assert_almost_equal(np.eye(N, k=k), mx.nd.eye(N, k=k).asnumpy())
@with_seed()
--
To stop receiving notification emails like this one, please contact
jxie@apache.org.