This is an automated email from the ASF dual-hosted git repository.
wkcn 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 35098b8 support leading dimension of -1 in ravel/unravel (#14356)
35098b8 is described below
commit 35098b8ad27e5151fc266fbf1a766d937bd9ec8e
Author: moin <[email protected]>
AuthorDate: Mon Mar 11 04:53:30 2019 +0100
support leading dimension of -1 in ravel/unravel (#14356)
---
src/operator/tensor/ravel.cc | 6 ++++--
src/operator/tensor/ravel.h | 3 ++-
tests/python/unittest/test_operator.py | 7 +++++++
3 files changed, 13 insertions(+), 3 deletions(-)
diff --git a/src/operator/tensor/ravel.cc b/src/operator/tensor/ravel.cc
index 0a66ea8..94d79c7 100644
--- a/src/operator/tensor/ravel.cc
+++ b/src/operator/tensor/ravel.cc
@@ -31,12 +31,13 @@ DMLC_REGISTER_PARAMETER(RavelParam);
NNVM_REGISTER_OP(_ravel_multi_index)
.add_alias("ravel_multi_index")
-.describe(R"code(Converts a batch of index arrays into an array of flat
indices. The operator follows numpy conventions so a single multi index is
given by a column of the input matrix.
+.describe(R"code(Converts a batch of index arrays into an array of flat
indices. The operator follows numpy conventions so a single multi index is
given by a column of the input matrix. The leading dimension may be left
unspecified by using -1 as placeholder.
Examples::
A = [[3,6,6],[4,5,1]]
ravel(A, shape=(7,6)) = [22,41,37]
+ ravel(A, shape=(-1,6)) = [22,41,37]
)code" ADD_FILELINE)
.set_num_inputs(1)
@@ -55,12 +56,13 @@ Examples::
NNVM_REGISTER_OP(_unravel_index)
.add_alias("unravel_index")
-.describe(R"code(Converts an array of flat indices into a batch of index
arrays. The operator follows numpy conventions so a single multi index is given
by a column of the output matrix.
+.describe(R"code(Converts an array of flat indices into a batch of index
arrays. The operator follows numpy conventions so a single multi index is given
by a column of the output matrix. The leading dimension may be left unspecified
by using -1 as placeholder.
Examples::
A = [22,41,37]
unravel(A, shape=(7,6)) = [[3,6,6],[4,5,1]]
+ unravel(A, shape=(-1,6)) = [[3,6,6],[4,5,1]]
)code" ADD_FILELINE)
.set_num_inputs(1)
diff --git a/src/operator/tensor/ravel.h b/src/operator/tensor/ravel.h
index 6d337dc..256fe33 100644
--- a/src/operator/tensor/ravel.h
+++ b/src/operator/tensor/ravel.h
@@ -110,11 +110,12 @@ struct unravel_index {
DType *unravelled, DType *ravelled) {
index_t idx(ravelled[i]);
#pragma unroll
- for (int j = ndim; j--; ) {
+ for (int j = ndim-1; j > 0; --j) {
index_t tmp = idx / shape[j];
unravelled[i+j*N] = idx - tmp*shape[j];
idx = tmp;
}
+ unravelled[i] = idx;
}
};
diff --git a/tests/python/unittest/test_operator.py
b/tests/python/unittest/test_operator.py
index 6bb8150..7169395 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -7106,6 +7106,13 @@ def test_ravel():
check_symbolic_forward(b, location={'a': data}, expected=[ravel_npy])
c = mx.sym.unravel_index(a, shape=shape)
check_symbolic_forward(c, location={'a': ravel_npy}, expected=[data])
+ # Test with leading dimension set to -1.
+ shape2 = shape
+ shape2 = (-1,)+shape[1:]
+ b = mx.sym.ravel_multi_index(a, shape=shape2)
+ check_symbolic_forward(b, location={'a': data}, expected=[ravel_npy])
+ c = mx.sym.unravel_index(a, shape=shape2)
+ check_symbolic_forward(c, location={'a': ravel_npy}, expected=[data])
def test_context_num_gpus():
try:
