From 288db02d562f523f8f3c9a0c34e2ac30df9a253c Mon Sep 17 00:00:00 2001
From: Tao Lv <tao.a.lv@intel.com>
Date: Sat, 17 Aug 2019 08:17:42 +0800
Subject: [PATCH 1/2] fix dtype in arange_like operator

---
 src/operator/tensor/init_op.cc | 2 +-
 src/operator/tensor/init_op.h  | 4 ----
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/operator/tensor/init_op.cc b/src/operator/tensor/init_op.cc
index 239cff798e59..110cfecd83ce 100644
--- a/src/operator/tensor/init_op.cc
+++ b/src/operator/tensor/init_op.cc
@@ -129,7 +129,7 @@ Examples::
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<RangeLikeParam>)
 .set_attr<mxnet::FInferShape>("FInferShape", RangeLikeShape)
-.set_attr<nnvm::FInferType>("FInferType", InitType<RangeLikeParam, 1>)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
 .set_attr<nnvm::FIgnoreInputs>("FIgnoreInputs",
     [](const NodeAttrs& attrs) { return std::vector<uint32_t>(1, 0); })
 .set_attr<FCompute>("FCompute<cpu>", RangeCompute<cpu, RangeLikeParam>)
diff --git a/src/operator/tensor/init_op.h b/src/operator/tensor/init_op.h
index 51c84363489a..f3c405d7103c 100644
--- a/src/operator/tensor/init_op.h
+++ b/src/operator/tensor/init_op.h
@@ -179,7 +179,6 @@ struct RangeLikeParam : public dmlc::Parameter<RangeLikeParam> {
   double step;
   int repeat;
   std::string ctx;
-  int dtype;
   dmlc::optional<int> axis;
 
   DMLC_DECLARE_PARAMETER(RangeLikeParam) {
@@ -197,9 +196,6 @@ struct RangeLikeParam : public dmlc::Parameter<RangeLikeParam> {
     .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)
-    MXNET_ADD_ALL_TYPES
-    .describe("Target data type.");
     DMLC_DECLARE_FIELD(axis)
     .set_default(dmlc::optional<int>())
     .describe("Arange elements according to the size of a certain axis of input array."

From c3a92f999a7b79787987c0fc4d493ed377dc9a93 Mon Sep 17 00:00:00 2001
From: Tao Lv <tao.a.lv@intel.com>
Date: Wed, 21 Aug 2019 21:23:25 +0800
Subject: [PATCH 2/2] add unit test

---
 tests/python/gpu/test_operator_gpu.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
index f8d8b4496afc..db550e4254b7 100644
--- a/tests/python/gpu/test_operator_gpu.py
+++ b/tests/python/gpu/test_operator_gpu.py
@@ -2331,6 +2331,22 @@ def test_math():
             for op in ops:
                 run_math(op, shape, dtype, check_value=check_value)
 
+@with_seed()
+def test_arange_like_dtype():
+    dtypes = [np.float16, np.float32, np.float64]
+
+    for t in dtypes:
+        x = mx.sym.Variable('x', dtype=t)
+        y = mx.sym.reshape(x, shape=(0, 0, -1))
+        z = mx.sym.contrib.arange_like(y, axis=-1)
+    
+        mod = z.simple_bind(ctx=mx.gpu(0), x=(3, 4, 5, 6), grad_req='null')
+        mod.arg_arrays[0][:] = np.random.normal(size=mod.arg_arrays[0].shape).astype(t)
+        out = mod.forward(is_train=False)
+        for v in out:
+            assert v.dtype == t
+    
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()