fix slice test and offset for reshape

Author: devcrocod

multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/ndarray/data/NDArray.kt

@@ -109,12 +109,13 @@ public class NDArray<T, D : Dimension> constructor(
 
         // TODO(get rid of copying)
         val newData = if (consistent) this.data else this.deepCopy().data
-        val newBase = if (consistent) base ?: this else null
+        val newBase = if (consistent) this.base ?: this else null
+        val newOffset = if (consistent) this.offset else 0
 
         return if (this.dim.d == 1 && this.shape.first() == dim1) {
             this as D1Array<T>
         } else {
-            D1Array(newData, this.offset, intArrayOf(dim1), dim = D1, base = newBase)
+            D1Array(newData, newOffset, intArrayOf(dim1), dim = D1, base = newBase)
         }
     }
 
@@ -125,12 +126,13 @@ public class NDArray<T, D : Dimension> constructor(
 
         // TODO(get rid of copying)
         val newData = if (consistent) this.data else this.deepCopy().data
-        val newBase = if (consistent) base ?: this else null
+        val newBase = if (consistent) this.base ?: this else null
+        val newOffset = if (consistent) this.offset else 0
 
         return if (this.shape.contentEquals(newShape)) {
             this as D2Array<T>
         } else {
-            D2Array(newData, this.offset, newShape, dim = D2, base = newBase)
+            D2Array(newData, newOffset, newShape, dim = D2, base = newBase)
         }
     }
 
@@ -142,11 +144,12 @@ public class NDArray<T, D : Dimension> constructor(
         // TODO(get rid of copying)
         val newData = if (consistent) this.data else this.deepCopy().data
         val newBase = if (consistent) base ?: this else null
+        val newOffset = if (consistent) this.offset else 0
 
         return if (this.shape.contentEquals(newShape)) {
             this as D3Array<T>
         } else {
-            D3Array(newData, this.offset, newShape, dim = D3, base = newBase)
+            D3Array(newData, newOffset, newShape, dim = D3, base = newBase)
         }
     }
 
@@ -157,12 +160,13 @@ public class NDArray<T, D : Dimension> constructor(
 
         // TODO(get rid of copying)
         val newData = if (consistent) this.data else this.deepCopy().data
-        val newBase = if (consistent) base ?: this else null
+        val newBase = if (consistent) this.base ?: this else null
+        val newOffset = if (consistent) this.offset else 0
 
         return if (this.shape.contentEquals(newShape)) {
             this as D4Array<T>
         } else {
-            D4Array(newData, this.offset, newShape, dim = D4, base = newBase)
+            D4Array(newData, newOffset, newShape, dim = D4, base = newBase)
         }
     }
 
@@ -175,12 +179,13 @@ public class NDArray<T, D : Dimension> constructor(
 
         // TODO(get rid of copying)
         val newData = if (consistent) this.data else this.deepCopy().data
-        val newBase = if (consistent) base ?: this else null
+        val newBase = if (consistent) this.base ?: this else null
+        val newOffset = if (consistent) this.offset else 0
 
         return if (this.shape.contentEquals(newShape)) {
             this as NDArray<T, DN>
         } else {
-            NDArray(newData, this.offset, newShape, dim = DN(newShape.size), base = newBase)
+            NDArray(newData, newOffset, newShape, dim = DN(newShape.size), base = newBase)
         }
     }
 
@@ -223,9 +228,10 @@ public class NDArray<T, D : Dimension> constructor(
         }
         // TODO(get rid of copying)
         val newData = if (consistent) this.data else this.deepCopy().data
-        val newBase = if (consistent) base ?: this else null
+        val newBase = if (consistent) this.base ?: this else null
+        val newOffset = if (consistent) this.offset else 0
 
-        return NDArray(newData, this.offset, newShape.toIntArray(), dim = DN(newShape.size), base = newBase)
+        return NDArray(newData, newOffset, newShape.toIntArray(), dim = DN(newShape.size), base = newBase)
     }
 
     override infix fun cat(other: MultiArray<T, D>): NDArray<T, D> =

multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/ndarray/operations/Transformation.kt

@@ -310,8 +310,9 @@ public fun <T> MultiArray<T, D1>.expandDims(axis: Int): MultiArray<T, D2> {
     val newShape = shape.toMutableList().apply { add(axis, 1) }.toIntArray()
     // TODO(get rid of copying)
     val newData = if (consistent) this.data else this.deepCopy().data
-    val newBase = if (consistent) base ?: this else null
-    return D2Array(newData, this.offset, newShape, dim = D2, base = newBase)
+    val newBase = if (consistent) this.base ?: this else null
+    val newOffset = if (consistent) this.offset else 0
+    return D2Array(newData, newOffset, newShape, dim = D2, base = newBase)
 }
 
 /**
@@ -322,8 +323,9 @@ public fun <T> MultiArray<T, D2>.expandDims(axis: Int): MultiArray<T, D3> {
     val newShape = shape.toMutableList().apply { add(axis, 1) }.toIntArray()
     // TODO(get rid of copying)
     val newData = if (consistent) this.data else this.deepCopy().data
-    val newBase = if (consistent) base ?: this else null
-    return D3Array(newData, this.offset, newShape, dim = D3, base = newBase)
+    val newBase = if (consistent) this.base ?: this else null
+    val newOffset = if (consistent) this.offset else 0
+    return D3Array(newData, newOffset, newShape, dim = D3, base = newBase)
 }
 
 /**
@@ -334,8 +336,9 @@ public fun <T> MultiArray<T, D3>.expandDims(axis: Int): MultiArray<T, D4> {
     val newShape = shape.toMutableList().apply { add(axis, 1) }.toIntArray()
     // TODO(get rid of copying)
     val newData = if (consistent) this.data else this.deepCopy().data
-    val newBase = if (consistent) base ?: this else null
-    return D4Array(newData, this.offset, newShape, dim = D4, base = newBase)
+    val newBase = if (consistent) this.base ?: this else null
+    val newOffset = if (consistent) this.offset else 0
+    return D4Array(newData, newOffset, newShape, dim = D4, base = newBase)
 }
 
 /**

multik-core/src/commonTest/kotlin/org/jetbrains/kotlinx/multik/ndarray/data/SliceTest.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2020-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 package org.jetbrains.kotlinx.multik.ndarray.data
@@ -71,15 +71,15 @@ class SliceTest {
         val a2 = a.reshape(3, 2)
         val b2 = b.reshape(4, 1)
         assertSame(a, a2.base)
-        assertSame(a, b2.base)
+        assertSame(null, b2.base)
 
         val d1 = b2.squeeze()
         val d2 = d1.unsqueeze()
-        assertSame(a, d1.base)
-        assertSame(a, d2.base)
+        assertSame(b2, d1.base)
+        assertSame(b2, d2.base)
 
         val e = b2.transpose()
-        assertSame(a, e.base)
+        assertSame(b2, e.base)
 
         val f = a2[1]
         assertSame(a, f.base)