You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2017/12/04 17:47:04 UTC

[GitHub] eric-haibin-lin closed pull request #8225: eye for dense and sparse

eric-haibin-lin closed pull request #8225: eye for dense and sparse
URL: https://github.com/apache/incubator-mxnet/pull/8225
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/python/mxnet/ndarray/__init__.py b/python/mxnet/ndarray/__init__.py
index 586dc9eda6..64b486a103 100644
--- a/python/mxnet/ndarray/__init__.py
+++ b/python/mxnet/ndarray/__init__.py
@@ -27,7 +27,7 @@
 from .op import *
 from .ndarray import *
 # pylint: enable=wildcard-import
-from .utils import load, save, zeros, empty, array
+from .utils import load, save, zeros, empty, array, eye
 from .sparse import _ndarray_cls
 from .ndarray import _GRAD_REQ_MAP
 
diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py
index a45a6a8247..53bdb45899 100644
--- a/python/mxnet/ndarray/ndarray.py
+++ b/python/mxnet/ndarray/ndarray.py
@@ -3351,6 +3351,51 @@ def imdecode(str_img, clip_rect=(0, 0, 0, 0), out=None, index=0, channels=3, mea
                                    out=out)
 
 
+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 zeros(shape, ctx=None, dtype=None, **kwargs):
     """Returns a new array filled with all zeros, with the given shape and type.
 
diff --git a/python/mxnet/ndarray/sparse.py b/python/mxnet/ndarray/sparse.py
index 229044e289..5b95a33b35 100644
--- a/python/mxnet/ndarray/sparse.py
+++ b/python/mxnet/ndarray/sparse.py
@@ -52,6 +52,7 @@
 from .ndarray import _STORAGE_TYPE_STR_TO_ID, _STORAGE_TYPE_ROW_SPARSE, _STORAGE_TYPE_CSR
 from .ndarray import _STORAGE_TYPE_UNDEFINED, _STORAGE_TYPE_DEFAULT
 from .ndarray import zeros as _zeros_ndarray
+from .ndarray import eye as _eye_ndarray
 from .ndarray import array as _array
 
 
@@ -1150,6 +1151,62 @@ def _ndarray_cls(handle, writable=True, stype=_STORAGE_TYPE_UNDEFINED):
 _set_ndarray_class(_ndarray_cls)
 
 
+def _init_op(ndarray_op, internal_op, stype, shape,
+             ctx, dtype, *args, **kwargs):
+    """Common interface to initialize the NDArray."""
+    if stype == 'default':
+        return ndarray_op(*args, ctx=ctx, dtype=dtype, **kwargs)
+    if ctx is None:
+        ctx = Context.default_ctx
+    dtype = mx_real_t if dtype is None else dtype
+    if stype == 'row_sparse' or stype == 'csr':
+        aux_types = _STORAGE_AUX_TYPES[stype]
+    else:
+        raise ValueError("unknown storage type" + stype)
+    out = _ndarray_cls(_new_alloc_handle(stype, shape, ctx, True, dtype, aux_types))
+    return internal_op(*args, ctx=ctx, dtype=dtype, out=out, **kwargs)
+
+
+def eye(stype, N, M=0, k=0, ctx=None, dtype=None, **kwargs):
+    """Return a 2-D array with ones on the diagonal and zeros elsewhere
+       of given shape and type.
+
+    Parameters
+    ----------
+    stype: string
+        The storage type of the return array, such as 'row_sparse', 'csr', etc
+    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 or CSRNDArray
+        A created array
+    Examples
+    --------
+    >>> mx.nd.sparse.eye('csr', 2).asnumpy()
+    array([[ 1.,  0.],
+           [ 0.,  1.]], dtype=float32)
+    >>> mx.nd.sparse.eye('csr', 2, 3, 1).asnumpy()
+    array([[ 0.,  1.,  0.],
+           [ 0.,  0.,  1.]], dtype=float32)
+    """
+    ncols = M if M > 0 else N
+    shape = (N, ncols)
+    return _init_op(_eye_ndarray, _internal._eye, stype, shape,
+                    ctx, dtype, N, M, k, **kwargs)
+
+
 def zeros(stype, shape, ctx=None, dtype=None, **kwargs):
     """Return a new array of given shape and type, filled with zeros.
 
@@ -1175,17 +1232,8 @@ def zeros(stype, shape, ctx=None, dtype=None, **kwargs):
     >>> mx.nd.sparse.zeros('row_sparse', (1,2), ctx=mx.cpu(), dtype='float16').asnumpy()
     array([[ 0.,  0.]], dtype=float16)
     """
-    if stype == 'default':
-        return _zeros_ndarray(shape, ctx=ctx, dtype=dtype, **kwargs)
-    if ctx is None:
-        ctx = Context.default_ctx
-    dtype = mx_real_t if dtype is None else dtype
-    if stype == 'row_sparse' or stype == 'csr':
-        aux_types = _STORAGE_AUX_TYPES[stype]
-    else:
-        raise ValueError("unknown storage type" + stype)
-    out = _ndarray_cls(_new_alloc_handle(stype, shape, ctx, True, dtype, aux_types))
-    return _internal._zeros(shape=shape, ctx=ctx, dtype=dtype, out=out, **kwargs)
+    return _init_op(_zeros_ndarray, _internal._zeros, stype, shape,
+                    ctx, dtype, shape, **kwargs)
 
 
 def empty(stype, shape, ctx=None, dtype=None):
diff --git a/python/mxnet/ndarray/utils.py b/python/mxnet/ndarray/utils.py
index 4f597c749f..97f8006570 100644
--- a/python/mxnet/ndarray/utils.py
+++ b/python/mxnet/ndarray/utils.py
@@ -26,6 +26,8 @@
 from .ndarray import empty as _empty_ndarray
 from .ndarray import zeros as _zeros_ndarray
 from .sparse import zeros as _zeros_sparse_ndarray
+from .ndarray import eye as _eye_ndarray
+from .sparse import eye as _eye_sparse_ndarray
 from .sparse import empty as _empty_sparse_ndarray
 from .sparse import array as _sparse_array
 from .sparse import _ndarray_cls
@@ -34,7 +36,55 @@
 except ImportError:
     spsp = None
 
-__all__ = ['zeros', 'empty', 'array', 'load', 'save']
+__all__ = ['zeros', 'empty', 'array', 'load', 'save', 'eye']
+
+
+def eye(N, M=0, k=0, ctx=None, dtype=None, stype=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 or CSRNDArray
+        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)>
+    >>> mx.nd.eye(2, stype='csr').asnumpy()
+    array([[ 1.,  0.],
+           [ 0.,  1.]], dtype=float32)
+    >>> mx.nd.eye(2, 3, 1, stype='csr').asnumpy()
+    array([[ 0.,  1.,  0.],
+           [ 0.,  0.,  1.]], dtype=float32)
+    """
+    if stype is None or stype == 'default':
+        return _eye_ndarray(N, M, k, ctx, dtype, **kwargs)
+    else:
+        return _eye_sparse_ndarray(stype, N, M, k, ctx, dtype, **kwargs)
 
 
 def zeros(shape, ctx=None, dtype=None, stype=None, **kwargs):
diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
index ce7776d948..8b3d6a5093 100644
--- a/python/mxnet/symbol/symbol.py
+++ b/python/mxnet/symbol/symbol.py
@@ -47,7 +47,7 @@
 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):
@@ -2700,6 +2700,33 @@ def hypot(left, right):
         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 5be4a2f421..23b8ba49f9 100644
--- a/src/operator/tensor/init_op.cc
+++ b/src/operator/tensor/init_op.cc
@@ -31,8 +31,21 @@ namespace op {
 DMLC_REGISTER_PARAMETER(InitOpParam);
 DMLC_REGISTER_PARAMETER(InitOpWithScalarParam);
 DMLC_REGISTER_PARAMETER(RangeParam);
+DMLC_REGISTER_PARAMETER(EyeParam);
 
 
+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<FInferStorageType>("FInferStorageType", InitStorageType<InitOpParam, false, true>)
+.set_attr<FCompute>("FCompute<cpu>", EyeFill<cpu>)
+.set_attr<FComputeEx>("FComputeEx<cpu>", EyeFillEx<cpu>)
+.add_arguments(EyeParam::__FIELDS__());
+
 NNVM_REGISTER_OP(_zeros)
 .describe("fill target with zeros")
 .set_num_inputs(0)
diff --git a/src/operator/tensor/init_op.cu b/src/operator/tensor/init_op.cu
index aeea2895b0..209ec89f76 100644
--- a/src/operator/tensor/init_op.cu
+++ b/src/operator/tensor/init_op.cu
@@ -43,6 +43,10 @@ void FillZerosCsrImpl(mshadow::Stream<mshadow::gpu> *s, const NDArray& dst) {
   });
 }
 
+NNVM_REGISTER_OP(_eye)
+.set_attr<FCompute>("FCompute<gpu>", EyeFill<gpu>)
+.set_attr<FComputeEx>("FComputeEx<gpu>", EyeFillEx<gpu>);
+
 NNVM_REGISTER_OP(_zeros)
 .set_attr<FCompute>("FCompute<gpu>", FillCompute<gpu, 0>)
 .set_attr<FComputeEx>("FComputeEx<gpu>", FillComputeZerosEx<gpu>);
diff --git a/src/operator/tensor/init_op.h b/src/operator/tensor/init_op.h
index 4d899704a1..9004aa8154 100644
--- a/src/operator/tensor/init_op.h
+++ b/src/operator/tensor/init_op.h
@@ -33,6 +33,7 @@
 #include <vector>
 #include <string>
 #include <limits>
+#include <algorithm>
 #include "../mshadow_op.h"
 #include "../elemwise_op_common.h"
 #include "../mxnet_op.h"
@@ -41,6 +42,155 @@
 namespace mxnet {
 namespace op {
 
+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 num_cols,
+                                  const nnvm::dim_t k) {
+    if ((i % num_cols) == ((i / num_cols) + k)) {
+      KERNEL_ASSIGN(out_data[i], req, static_cast<DType>(1));
+    } else {
+      KERNEL_ASSIGN(out_data[i], req, static_cast<DType>(0));
+    }
+  }
+};
+
+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;
+  using namespace mxnet_op;
+  MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {
+    MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
+      Kernel<eye_dns_fill<req_type>, xpu>::Launch(
+          s, out_data.Size(), out_data.dptr<DType>(),
+          num_cols, param.k);
+    });
+  });
+}
+
+// fill indptr array of eye in CSR format
+struct eye_csr_indptr_fill {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i, DType* out_data, const nnvm::dim_t nnz,
+                                  const nnvm::dim_t init_row) {
+    const nnvm::dim_t data = (i - init_row) > 0 ? (i - init_row) : 0;
+    out_data[i] = static_cast<DType>(data <= nnz ? data : nnz);
+  }
+};
+
+// fill indices array of eye in CSR format
+struct eye_csr_idx_fill {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i, DType* out_data, const nnvm::dim_t init_col) {
+    out_data[i] = static_cast<DType>(i + init_col);
+  }
+};
+
+
+template<typename xpu>
+void EyeFillCsrImpl(mshadow::Stream<xpu> *s, const NDArray& out,
+                    const nnvm::dim_t N, const nnvm::dim_t M, const nnvm::dim_t k) {
+  const nnvm::dim_t num_cols = M > 0 ? M : N;
+  const nnvm::dim_t cnnz = std::max(num_cols - std::abs(k), (nnvm::dim_t)0);
+  const nnvm::dim_t rnnz = std::max(N - std::abs(k), (nnvm::dim_t)0);
+  const nnvm::dim_t nnz = k > 0 ? std::min(cnnz, N) :
+                                  std::min(rnnz, num_cols);
+  using namespace mxnet_op;
+  out.CheckAndAllocAuxData(csr::kIndPtr, mshadow::Shape1(N + 1));
+  out.CheckAndAllocAuxData(csr::kIdx, mshadow::Shape1(nnz));
+  out.CheckAndAllocData(mshadow::Shape1(nnz));
+
+  MSHADOW_IDX_TYPE_SWITCH(out.aux_type(csr::kIndPtr), RType, {
+    MSHADOW_IDX_TYPE_SWITCH(out.aux_type(csr::kIdx), IType, {
+      MSHADOW_TYPE_SWITCH(out.dtype(), DType, {
+        Kernel<eye_csr_indptr_fill, xpu>::Launch(
+          s, N + 1, out.aux_data(csr::kIndPtr).dptr<RType>(),
+          nnz, std::abs(std::min((nnvm::dim_t)0, k)));
+        Kernel<eye_csr_idx_fill, xpu>::Launch(
+          s, nnz, out.aux_data(csr::kIdx).dptr<IType>(),
+          std::max(std::min(k, num_cols-1), (nnvm::dim_t)0));
+        Kernel<set_to_int<1>, xpu>::Launch(
+          s, nnz, out.data().dptr<DType>());
+      });
+    });
+  });
+}
+
+template<typename xpu>
+void EyeFillEx(const nnvm::NodeAttrs& attrs,
+               const OpContext& ctx,
+               const std::vector<NDArray>& inputs,
+               const std::vector<OpReqType>& req,
+               const std::vector<NDArray>& outputs) {
+  using namespace mshadow;
+  using namespace mshadow::expr;
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+  CHECK_EQ(outputs.size(), 1);
+  auto stype = outputs[0].storage_type();
+  const EyeParam& param = nnvm::get<EyeParam>(attrs.parsed);
+  if (req[0] == kNullOp) return;
+  CHECK_EQ(req[0], kWriteTo) << "kWriteTo is expected for EyeFillEx";
+  if (stype == kCSRStorage) {
+    EyeFillCsrImpl<xpu>(s, outputs[0], param.N, param.M, param.k);
+  } else {
+    LOG(FATAL) << "Not implemented: " << operator_string(attrs, ctx, inputs, req, outputs);
+  }
+}
+
 struct InitOpParam : public dmlc::Parameter<InitOpParam> {
   TShape shape;
   std::string ctx;
diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
index cecda21a07..b33a0e7933 100644
--- a/tests/python/gpu/test_operator_gpu.py
+++ b/tests/python/gpu/test_operator_gpu.py
@@ -35,7 +35,7 @@
 #from test_rnn import *
 from test_gluon_rnn import *
 from test_sparse_ndarray import test_create_csr, test_create_row_sparse, test_sparse_nd_slice
-from test_sparse_ndarray import test_create_sparse_nd_empty, test_create_sparse_nd_from_sparse
+from test_sparse_ndarray import test_create_sparse_nd_empty, test_create_sparse_nd_from_sparse, test_sparse_eye
 from test_sparse_ndarray import test_create_sparse_nd_from_dense, test_create_sparse_nd_infer_shape
 from test_sparse_ndarray import test_sparse_nd_check_format, test_sparse_nd_copy
 from test_sparse_operator import *
diff --git a/tests/python/unittest/test_ndarray.py b/tests/python/unittest/test_ndarray.py
index 5512b07c77..600566c947 100644
--- a/tests/python/unittest/test_ndarray.py
+++ b/tests/python/unittest/test_ndarray.py
@@ -733,6 +733,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=3, size=3)
+    M_array = np.random.randint(1, high=3, size=3)
+    k_array = np.random.randint(-5, high=5, size=3)
+    for i in range(3):
+        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())
+
 
 def test_ndarray_fluent():
     has_grad = set(['flatten', 'expand_dims', 'flip', 'tile', 'transpose', 'sum', 'nansum', 'prod',
diff --git a/tests/python/unittest/test_sparse_ndarray.py b/tests/python/unittest/test_sparse_ndarray.py
index e59e476601..29b0a72d61 100644
--- a/tests/python/unittest/test_sparse_ndarray.py
+++ b/tests/python/unittest/test_sparse_ndarray.py
@@ -705,6 +705,20 @@ def check_rsp_empty(shape, dtype, ctx):
     check_rsp_empty(shape, dtype, ctx)
     check_rsp_empty(shape_3d, dtype, ctx)
 
+    
+def test_sparse_eye():
+    stypes = ['csr', 'row_sparse', 'default']
+    for stype in stypes:
+        N_array = np.random.randint(1, high=3, size=3)
+        M_array = np.random.randint(1, high=3, size=3)
+        k_array = np.random.randint(-5, high=5, size=3)
+        for i in range(3):
+            N = N_array[i]
+            M = M_array[i]
+            k = k_array[i]
+            same(mx.nd.sparse.eye(stype, N, M, k).asnumpy(), np.eye(N, M, k))
+
+
 def test_synthetic_dataset_generator():
     def test_powerlaw_generator(csr_arr, final_row=1):
         """Test power law distribution


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services