Viewing arrays with different shapes

Author: karllessard

ndarray/src/main/java/org/tensorflow/ndarray/BooleanNdArray.java

@@ -68,6 +68,9 @@ public interface BooleanNdArray extends NdArray<Boolean> {
    */
   BooleanNdArray setBoolean(boolean value, long... coordinates);
 
+  @Override
+  BooleanNdArray withShape(Shape shape);
+
   @Override
   BooleanNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/ByteNdArray.java

@@ -68,6 +68,9 @@ public interface ByteNdArray extends NdArray<Byte> {
    */
   ByteNdArray setByte(byte value, long... coordinates);
 
+  @Override
+  ByteNdArray withShape(Shape shape);
+
   @Override
   ByteNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/DoubleNdArray.java

@@ -83,6 +83,9 @@ default DoubleStream streamOfDoubles() {
     return StreamSupport.stream(scalars().spliterator(), false).mapToDouble(DoubleNdArray::getDouble);
   }
 
+  @Override
+  DoubleNdArray withShape(Shape shape);
+
   @Override
   DoubleNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/FloatNdArray.java

@@ -68,6 +68,9 @@ public interface FloatNdArray extends NdArray<Float> {
    */
   FloatNdArray setFloat(float value, long... coordinates);
 
+  @Override
+  FloatNdArray withShape(Shape shape);
+
   @Override
   FloatNdArray slice(Index... coordinates);
 

ndarray/src/main/java/org/tensorflow/ndarray/IntNdArray.java

@@ -83,6 +83,9 @@ default IntStream streamOfInts() {
     return StreamSupport.stream(scalars().spliterator(), false).mapToInt(IntNdArray::getInt);
   }
 
+  @Override
+  IntNdArray withShape(Shape shape);
+
   @Override
   IntNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/LongNdArray.java

@@ -83,6 +83,9 @@ default LongStream streamOfLongs() {
     return StreamSupport.stream(scalars().spliterator(), false).mapToLong(LongNdArray::getLong);
   }
 
+  @Override
+  LongNdArray withShape(Shape shape);
+
   @Override
   LongNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/NdArray.java

@@ -16,14 +16,14 @@
  */
 package org.tensorflow.ndarray;
 
+import org.tensorflow.ndarray.buffer.DataBuffer;
+import org.tensorflow.ndarray.index.Index;
+
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
-import org.tensorflow.ndarray.buffer.DataBuffer;
-import org.tensorflow.ndarray.index.Index;
-
 /**
  * A data structure of N-dimensions.
  *
@@ -101,6 +101,32 @@ public interface NdArray<T> extends Shaped {
    */
   NdArraySequence<? extends NdArray<T>> scalars();
 
+  /**
+   * Returns a new N-dimensional view of this array with the given {@code shape}.
+   *
+   * <p>The provided {@code shape} must comply to the following characteristics:
+   * <ul>
+   *     <li>new shape is known (i.e. has no unknown dimension)</li>
+   *     <li>new shape size is equal to the size of the current shape (i.e. same number of elements)</li>
+   * </ul>
+   * For example,
+   * <pre>{@code
+   *    NdArrays.ofInts(Shape.scalar()).withShape(Shape.of(1, 1));  // ok
+   *    NdArrays.ofInts(Shape.of(2, 3).withShape(Shape.of(3, 2));   // ok
+   *    NdArrays.ofInts(Shape.scalar()).withShape(Shape.of(1, 2));  // not ok, sizes are different (1 != 2)
+   *    NdArrays.ofInts(Shape.of(2, 3)).withShape(Shape.unknown()); // not ok, new shape unknown
+   * }</pre>
+   *
+   * <p>Any changes applied to the returned view affect the data of this array as well, as there
+   * is no copy involved.
+   *
+   * @param shape the new shape to apply
+   * @return a new array viewing the data according to the new shape, or this array if shapes are the same
+   * @throws IllegalArgumentException if the provided {@code shape} is not compliant
+   * @throws UnsupportedOperationException if this array does not support this operation
+   */
+  NdArray<T> withShape(Shape shape);
+
   /**
    * Creates a multi-dimensional view (or slice) of this array by mapping one or more dimensions
    * to the given index selectors.

ndarray/src/main/java/org/tensorflow/ndarray/ShortNdArray.java

@@ -68,6 +68,9 @@ public interface ShortNdArray extends NdArray<Short> {
    */
   ShortNdArray setShort(short value, long... coordinates);
 
+  @Override
+  ShortNdArray withShape(Shape shape);
+
   @Override
   ShortNdArray slice(Index... coordinates);
 

ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java

@@ -18,6 +18,7 @@
 
 import org.tensorflow.ndarray.NdArray;
 import org.tensorflow.ndarray.NdArraySequence;
+import org.tensorflow.ndarray.Shape;
 import org.tensorflow.ndarray.impl.AbstractNdArray;
 import org.tensorflow.ndarray.impl.dimension.RelativeDimensionalSpace;
 import org.tensorflow.ndarray.impl.sequence.FastElementSequence;
@@ -43,18 +44,29 @@ public NdArraySequence<U> elements(int dimensionIdx) {
     DimensionalSpace elemDims = dimensions().from(dimensionIdx + 1);
     try {
       DataBufferWindow<? extends DataBuffer<T>> elemWindow = buffer().window(elemDims.physicalSize());
-      U element = instantiate(elemWindow.buffer(), elemDims);
+      U element = instantiateView(elemWindow.buffer(), elemDims);
       return new FastElementSequence(this, dimensionIdx, element, elemWindow);
     } catch (UnsupportedOperationException e) {
       // If buffer windows are not supported, fallback to slicing (and slower) sequence
       return new SlicingElementSequence<>(this, dimensionIdx, elemDims);
     }
   }
 
+  @Override
+  public U withShape(Shape shape) {
+    if (shape == this.shape()) {
+      return (U)this;
+    }
+    if (shape == null || shape.isUnknown() || shape.size() != this.shape().size()) {
+      throw new IllegalArgumentException("Shape " + shape + " cannot be used to reshape ndarray of shape " + this.shape());
+    }
+    return instantiateView(buffer(), DimensionalSpace.create(shape));
+  }
+
   @Override
   public U slice(long position, DimensionalSpace sliceDimensions) {
     DataBuffer<T> sliceBuffer = buffer().slice(position, sliceDimensions.physicalSize());
-    return instantiate(sliceBuffer, sliceDimensions);
+    return instantiateView(sliceBuffer, sliceDimensions);
   }
 
   @Override
@@ -147,7 +159,7 @@ protected AbstractDenseNdArray(DimensionalSpace dimensions) {
 
   abstract protected DataBuffer<T> buffer();
 
-  abstract U instantiate(DataBuffer<T> buffer, DimensionalSpace dimensions);
+  abstract U instantiateView(DataBuffer<T> buffer, DimensionalSpace dimensions);
 
   long positionOf(long[] coords, boolean isValue) {
     if (coords == null || coords.length == 0) {

ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/BooleanDenseNdArray.java

@@ -73,7 +73,7 @@ protected BooleanDenseNdArray(BooleanDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  BooleanDenseNdArray instantiate(DataBuffer<Boolean> buffer, DimensionalSpace dimensions) {
+  BooleanDenseNdArray instantiateView(DataBuffer<Boolean> buffer, DimensionalSpace dimensions) {
     return new BooleanDenseNdArray((BooleanDataBuffer)buffer, dimensions);
   }