[MXNET-1379] update reshape operator

scala-package/core/src/main/scala/org/apache/mxnet/LibInfo.scala

@@ -86,9 +86,10 @@ private[mxnet] class LibInfo {
  @native def mxNDArrayAt(handle: NDArrayHandle,
  idx: MXUint,
  out: NDArrayHandleRef): Int
- @native def mxNDArrayReshape(handle: NDArrayHandle,
+ @native def mxNDArrayReshape64(handle: NDArrayHandle,
  nDim: Int,
- dims: Array[Int],
+ dims: Array[Long],
+ reverse: Boolean,
  reshapeHandle: NDArrayHandleRef): Int
  @native def mxNDArraySyncCopyFromCPU(handle: NDArrayHandle,
  source: Array[MXFloat],

scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala

@@ -950,8 +950,19 @@ class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle,
  * @return a reshaped NDArray that shares memory with current one.
  */
  def reshape(dims: Array[Int]): NDArray = {
+ reshape(dims.map(_.toLong))
+ }
+
+ /**
+ * Return a reshaped NDArray that shares memory with current one.
+ * @param dims New shape.
+ * @param reverse whether to inplace reshape
+ * @return a reshaped NDArray that shares memory with current one.
+ */
+ def reshape(dims: Array[Long], reverse: Option[Boolean] = None): NDArray = {
  val reshapeHandle = new NDArrayHandleRef
- checkCall(_LIB.mxNDArrayReshape(handle, dims.length, dims, reshapeHandle))
+ checkCall(_LIB.mxNDArrayReshape64(handle,
+ dims.length, dims, reverse.getOrElse(false), reshapeHandle))
  new NDArray(handle = reshapeHandle.value, writable = this.writable)
  }
 

scala-package/core/src/test/scala/org/apache/mxnet/NDArraySuite.scala

@@ -878,14 +878,18 @@ class NDArraySuite extends FunSuite with BeforeAndAfterAll with Matchers {
  }
 
  test("reshape") {
- val arr = NDArray.array(Array(1f, 2f, 3f, 4f, 5f, 6f), shape = Shape(3, 2))
+ var arr = NDArray.array(Array(1f, 2f, 3f, 4f, 5f, 6f), shape = Shape(3, 2))
 
- val arr1 = arr.reshape(Array(2, 3))
+ var arr1 = arr.reshape(Array(2, 3))
  assert(arr1.shape === Shape(2, 3))
  assert(arr1.toArray === Array(1f, 2f, 3f, 4f, 5f, 6f))
 
  arr.set(1f)
  assert(arr1.toArray === Array(1f, 1f, 1f, 1f, 1f, 1f))
+
+ arr = NDArray.ones(1, 384, 1)
+ arr1 = arr.reshape(Array(0, -3))
+ assert(arr1.shape === Shape(1, 384))
  }
 
  test("dispose deps") {

scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc

@@ -404,14 +404,15 @@ JNIEXPORT jint JNICALL Java_org_apache_mxnet_LibInfo_mxNDArrayAt
  return ret;
 }
 
-JNIEXPORT jint JNICALL Java_org_apache_mxnet_LibInfo_mxNDArrayReshape
- (JNIEnv *env, jobject obj, jlong ndArrayPtr, jint ndim, jintArray dims, jobject reshapedHandle) {
+JNIEXPORT jint JNICALL Java_org_apache_mxnet_LibInfo_mxNDArrayReshape64
+ (JNIEnv *env, jobject obj, jlong ndArrayPtr, jint ndim,
+ jlongArray dims, jboolean reverse, jobject reshapedHandle) {
  NDArrayHandle out;
- jint *pdims = env->GetIntArrayElements(dims, NULL);
- int ret = MXNDArrayReshape(reinterpret_cast<NDArrayHandle>(ndArrayPtr), ndim,
- reinterpret_cast<int *>(pdims), &out);
+ jlong *pdims = env->GetLongArrayElements(dims, NULL);
+ int ret = MXNDArrayReshape64(reinterpret_cast<NDArrayHandle>(ndArrayPtr), ndim,
+ reinterpret_cast<dim_t *>(pdims), reverse, &out);
  SetLongField(env, reshapedHandle, reinterpret_cast<jlong>(out));
- env->ReleaseIntArrayElements(dims, pdims, 0);
+ env->ReleaseLongArrayElements(dims, pdims, 0);
  return ret;
 }