fix reshaping for slicing arrays

Author: devcrocod

gradle.properties

@@ -1,5 +1,5 @@
 kotlin.code.style=official
-multik_version=0.2.0-dev-3
+multik_version=0.2.0-dev-4
 
 # Kotlin
 systemProp.kotlin_version=1.7.10

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

@@ -107,10 +107,14 @@ public class NDArray<T, D : Dimension> constructor(
         requirePositiveShape(dim1)
         require(dim1 == size) { "Cannot reshape array of size $size into a new shape ($dim1)" }
 
+        // TODO(get rid of copying)
+        val newData = if (consistent) this.data else this.deepCopy().data
+        val newBase = if (consistent) base ?: this else null
+
         return if (this.dim.d == 1 && this.shape.first() == dim1) {
             this as D1Array<T>
         } else {
-            D1Array(this.data, this.offset, intArrayOf(dim1), dim = D1, base = base ?: this)
+            D1Array(newData, this.offset, intArrayOf(dim1), dim = D1, base = newBase)
         }
     }
 
@@ -119,10 +123,14 @@ public class NDArray<T, D : Dimension> constructor(
         newShape.forEach { requirePositiveShape(it) }
         require(dim1 * dim2 == size) { "Cannot reshape array of size $size into a new shape ($dim1, $dim2)" }
 
+        // TODO(get rid of copying)
+        val newData = if (consistent) this.data else this.deepCopy().data
+        val newBase = if (consistent) base ?: this else null
+
         return if (this.shape.contentEquals(newShape)) {
             this as D2Array<T>
         } else {
-            D2Array(this.data, this.offset, newShape, dim = D2, base = base ?: this)
+            D2Array(newData, this.offset, newShape, dim = D2, base = newBase)
         }
     }
 
@@ -131,10 +139,14 @@ public class NDArray<T, D : Dimension> constructor(
         newShape.forEach { requirePositiveShape(it) }
         require(dim1 * dim2 * dim3 == size) { "Cannot reshape array of size $size into a new shape ($dim1, $dim2, $dim3)" }
 
+        // TODO(get rid of copying)
+        val newData = if (consistent) this.data else this.deepCopy().data
+        val newBase = if (consistent) base ?: this else null
+
         return if (this.shape.contentEquals(newShape)) {
             this as D3Array<T>
         } else {
-            D3Array(this.data, this.offset, newShape, dim = D3, base = base ?: this)
+            D3Array(newData, this.offset, newShape, dim = D3, base = newBase)
         }
     }
 
@@ -143,10 +155,14 @@ public class NDArray<T, D : Dimension> constructor(
         newShape.forEach { requirePositiveShape(it) }
         require(dim1 * dim2 * dim3 * dim4 == size) { "Cannot reshape array of size $size into a new shape ($dim1, $dim2, $dim3, $dim4)" }
 
+        // TODO(get rid of copying)
+        val newData = if (consistent) this.data else this.deepCopy().data
+        val newBase = if (consistent) base ?: this else null
+
         return if (this.shape.contentEquals(newShape)) {
             this as D4Array<T>
         } else {
-            D4Array(this.data, this.offset, newShape, dim = D4, base = base ?: this)
+            D4Array(newData, this.offset, newShape, dim = D4, base = newBase)
         }
     }
 
@@ -157,10 +173,14 @@ public class NDArray<T, D : Dimension> constructor(
             "Cannot reshape array of size $size into a new shape ${newShape.joinToString(prefix = "(", postfix = ")")}"
         }
 
+        // TODO(get rid of copying)
+        val newData = if (consistent) this.data else this.deepCopy().data
+        val newBase = if (consistent) base ?: this else null
+
         return if (this.shape.contentEquals(newShape)) {
             this as NDArray<T, DN>
         } else {
-            NDArray(this.data, this.offset, newShape, dim = DN(newShape.size), base = base ?: this)
+            NDArray(newData, this.offset, newShape, dim = DN(newShape.size), base = newBase)
         }
     }
 
@@ -201,13 +221,11 @@ public class NDArray<T, D : Dimension> constructor(
         for (axis in axes.sorted()) {
             newShape.add(axis, 1)
         }
-        return NDArray(
-            this.data,
-            this.offset,
-            newShape.toIntArray(),
-            dim = DN(newShape.size),
-            base = base ?: this
-        )
+        // TODO(get rid of copying)
+        val newData = if (consistent) this.data else this.deepCopy().data
+        val newBase = if (consistent) base ?: this else null
+
+        return NDArray(newData, this.offset, 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

@@ -306,39 +306,48 @@ public fun <T, D : Dimension> MultiArray<T, D>.clip(min: T, max: T): NDArray<T,
  * Returns a ndarray with an expanded shape.
  */
 @JvmName("expandDimsD1")
-public fun <T, D : Dimension> MultiArray<T, D1>.expandDims(axis: Int): MultiArray<T, D2> {
+public fun <T> MultiArray<T, D1>.expandDims(axis: Int): MultiArray<T, D2> {
     val newShape = shape.toMutableList().apply { add(axis, 1) }.toIntArray()
-    return D2Array(this.data, this.offset, newShape, dim = D2, base = base ?: this)
+    // 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)
 }
 
 /**
  * Returns a ndarray with an expanded shape.
  */
 @JvmName("expandDimsD2")
-public fun <T, D : Dimension> MultiArray<T, D2>.expandDims(axis: Int): MultiArray<T, D3> {
+public fun <T> MultiArray<T, D2>.expandDims(axis: Int): MultiArray<T, D3> {
     val newShape = shape.toMutableList().apply { add(axis, 1) }.toIntArray()
-    return D3Array(this.data, this.offset, newShape, dim = D3, base = base ?: this)
+    // 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)
 }
 
 /**
  * Returns a ndarray with an expanded shape.
  */
 @JvmName("expandDimsD3")
-public fun <T, D : Dimension> MultiArray<T, D3>.expandDims(axis: Int): MultiArray<T, D4> {
+public fun <T> MultiArray<T, D3>.expandDims(axis: Int): MultiArray<T, D4> {
     val newShape = shape.toMutableList().apply { add(axis, 1) }.toIntArray()
-    return D4Array(this.data, this.offset, newShape, dim = D4, base = base ?: this)
+    // 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)
 }
 
 /**
  * Returns a ndarray with an expanded shape.
  */
 @JvmName("expandDimsD4")
-public fun <T, D : Dimension> MultiArray<T, D4>.expandDims(axis: Int): MultiArray<T, DN> = this.unsqueeze()
+public fun <T> MultiArray<T, D4>.expandDims(axis: Int): MultiArray<T, DN> = this.unsqueeze()
 
 /**
  * Returns a ndarray with an expanded shape.
  *
  * @see MultiArray.unsqueeze
  */
 @JvmName("expandDimsDN")
-public fun <T, D : Dimension> MultiArray<T, D>.expandDims(vararg axes: Int): MultiArray<T, DN> = this.unsqueeze()
+public fun <T, D : Dimension> MultiArray<T, D>.expandNDims(vararg axes: Int): MultiArray<T, DN> = this.unsqueeze()