From dac31390f4865dd15f48b4809585ac3f62827d22 Mon Sep 17 00:00:00 2001
From: Karl Lessard <karl.lessard@gmail.com>
Date: Sat, 17 Feb 2024 15:26:08 -0500
Subject: [PATCH 1/3] Viewing arrays with different shapes

---
 .../tensorflow/ndarray/BooleanNdArray.java    |  3 ++
 .../org/tensorflow/ndarray/ByteNdArray.java   |  3 ++
 .../org/tensorflow/ndarray/DoubleNdArray.java |  3 ++
 .../org/tensorflow/ndarray/FloatNdArray.java  |  3 ++
 .../org/tensorflow/ndarray/IntNdArray.java    |  3 ++
 .../org/tensorflow/ndarray/LongNdArray.java   |  3 ++
 .../java/org/tensorflow/ndarray/NdArray.java  | 32 +++++++++++++++++--
 .../org/tensorflow/ndarray/ShortNdArray.java  |  3 ++
 .../impl/dense/AbstractDenseNdArray.java      | 18 +++++++++--
 .../impl/dense/BooleanDenseNdArray.java       |  2 +-
 .../ndarray/impl/dense/ByteDenseNdArray.java  |  2 +-
 .../ndarray/impl/dense/DenseNdArray.java      |  2 +-
 .../impl/dense/DoubleDenseNdArray.java        |  2 +-
 .../ndarray/impl/dense/FloatDenseNdArray.java |  2 +-
 .../ndarray/impl/dense/IntDenseNdArray.java   |  2 +-
 .../ndarray/impl/dense/LongDenseNdArray.java  |  2 +-
 .../ndarray/impl/dense/ShortDenseNdArray.java |  2 +-
 .../impl/sparse/AbstractSparseNdArray.java    |  5 +++
 .../tensorflow/ndarray/NdArrayTestBase.java   | 27 ++++++++++++++++
 .../tensorflow/ndarray/SparseNdArrayTest.java | 10 ++++--
 20 files changed, 112 insertions(+), 17 deletions(-)

diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/BooleanNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/BooleanNdArray.java
index 5b4bedb..bd16a9a 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/BooleanNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/ByteNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/ByteNdArray.java
index 0e6f118..47e5a0d 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/ByteNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/DoubleNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/DoubleNdArray.java
index b0e6dab..da42bab 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/DoubleNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/FloatNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/FloatNdArray.java
index 8d4fbf5..34e4201 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/FloatNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/IntNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/IntNdArray.java
index e6a5cf0..71f19b9 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/IntNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/LongNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/LongNdArray.java
index e7bd266..a55b2ab 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/LongNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/NdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/NdArray.java
index f1e84d4..a75da48 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/NdArray.java
+++ b/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.
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/ShortNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/ShortNdArray.java
index f9335b4..92b608f 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/ShortNdArray.java
+++ b/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);
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
index 30af952..8b5c5ba 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
+++ b/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,7 +44,7 @@ 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
@@ -51,10 +52,21 @@ public NdArraySequence<U> elements(int dimensionIdx) {
     }
   }
 
+  @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) {
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/BooleanDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/BooleanDenseNdArray.java
index 0764146..9c134b5 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/BooleanDenseNdArray.java
+++ b/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);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ByteDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ByteDenseNdArray.java
index 172432b..a2525c6 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ByteDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ByteDenseNdArray.java
@@ -73,7 +73,7 @@ protected ByteDenseNdArray(ByteDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  ByteDenseNdArray instantiate(DataBuffer<Byte> buffer, DimensionalSpace dimensions) {
+  ByteDenseNdArray instantiateView(DataBuffer<Byte> buffer, DimensionalSpace dimensions) {
     return new ByteDenseNdArray((ByteDataBuffer)buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DenseNdArray.java
index 819d95d..18b3755 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DenseNdArray.java
@@ -45,7 +45,7 @@ protected DenseNdArray(DataBuffer<T> buffer, Shape shape) {
   }
 
   @Override
-  DenseNdArray<T> instantiate(DataBuffer<T> buffer, DimensionalSpace dimensions) {
+  DenseNdArray<T> instantiateView(DataBuffer<T> buffer, DimensionalSpace dimensions) {
     return new DenseNdArray<>(buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DoubleDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DoubleDenseNdArray.java
index f54b8d0..a967ce1 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DoubleDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/DoubleDenseNdArray.java
@@ -73,7 +73,7 @@ protected DoubleDenseNdArray(DoubleDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  DoubleDenseNdArray instantiate(DataBuffer<Double> buffer, DimensionalSpace dimensions) {
+  DoubleDenseNdArray instantiateView(DataBuffer<Double> buffer, DimensionalSpace dimensions) {
     return new DoubleDenseNdArray((DoubleDataBuffer)buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/FloatDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/FloatDenseNdArray.java
index 196b5ef..a04c192 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/FloatDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/FloatDenseNdArray.java
@@ -73,7 +73,7 @@ protected FloatDenseNdArray(FloatDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  FloatDenseNdArray instantiate(DataBuffer<Float> buffer, DimensionalSpace dimensions) {
+  FloatDenseNdArray instantiateView(DataBuffer<Float> buffer, DimensionalSpace dimensions) {
     return new FloatDenseNdArray((FloatDataBuffer) buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/IntDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/IntDenseNdArray.java
index a7af498..e1a726f 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/IntDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/IntDenseNdArray.java
@@ -73,7 +73,7 @@ protected IntDenseNdArray(IntDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  IntDenseNdArray instantiate(DataBuffer<Integer> buffer, DimensionalSpace dimensions) {
+  IntDenseNdArray instantiateView(DataBuffer<Integer> buffer, DimensionalSpace dimensions) {
     return new IntDenseNdArray((IntDataBuffer)buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/LongDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/LongDenseNdArray.java
index cd56dad..802cbcd 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/LongDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/LongDenseNdArray.java
@@ -73,7 +73,7 @@ protected LongDenseNdArray(LongDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  LongDenseNdArray instantiate(DataBuffer<Long> buffer, DimensionalSpace dimensions) {
+  LongDenseNdArray instantiateView(DataBuffer<Long> buffer, DimensionalSpace dimensions) {
     return new LongDenseNdArray((LongDataBuffer)buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ShortDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ShortDenseNdArray.java
index 291c01a..434b260 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ShortDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/ShortDenseNdArray.java
@@ -73,7 +73,7 @@ protected ShortDenseNdArray(ShortDataBuffer buffer, Shape shape) {
   }
 
   @Override
-  ShortDenseNdArray instantiate(DataBuffer<Short> buffer, DimensionalSpace dimensions) {
+  ShortDenseNdArray instantiateView(DataBuffer<Short> buffer, DimensionalSpace dimensions) {
     return new ShortDenseNdArray((ShortDataBuffer)buffer, dimensions);
   }
 
diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/sparse/AbstractSparseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/sparse/AbstractSparseNdArray.java
index 8e3892d..e4a2fba 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/sparse/AbstractSparseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/sparse/AbstractSparseNdArray.java
@@ -212,6 +212,11 @@ protected long[] getIndicesCoordinates(LongNdArray l) {
    */
   public abstract U toDense();
 
+  @Override
+  public U withShape(Shape shape) {
+    throw new UnsupportedOperationException("Sparse NdArrays cannot be viewed with a different shape");
+  }
+
   /** {@inheritDoc} */
   @Override
   public NdArray<T> slice(Index... indices) {
diff --git a/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java b/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java
index 8a09ec7..35806f7 100644
--- a/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java
+++ b/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java
@@ -384,4 +384,31 @@ public void streamingObjects() {
     values = matrix.streamOfObjects().collect(Collectors.toList());
     assertIterableEquals(List.of(valueOf(1L), valueOf(2L), valueOf(3L), valueOf(4L)), values);
   }
+
+  @Test
+  public void withShape() {
+    Shape originalShape = Shape.scalar();
+    Shape newShape = originalShape.prepend(1).prepend(1); // [1, 1]
+
+    NdArray<T> originalArray = allocate(originalShape);
+    originalArray.setObject(valueOf(10L));
+    assertEquals(valueOf(10L), originalArray.getObject());
+
+    NdArray<T> newArray = originalArray.withShape(newShape);
+    assertNotNull(newArray);
+    assertEquals(newShape, newArray.shape());
+    assertEquals(originalShape, originalArray.shape());
+    assertEquals(valueOf(10L), newArray.getObject(0, 0));
+
+    NdArray<T> sameArray = originalArray.withShape(originalShape);
+    assertSame(originalArray, sameArray);
+
+    assertThrows(IllegalArgumentException.class, () -> originalArray.withShape(Shape.of(2)));
+    assertThrows(IllegalArgumentException.class, () -> originalArray.withShape(Shape.unknown()));
+
+    NdArray<T> originalMatrix = allocate(Shape.of(2, 3));
+    assertThrows(IllegalArgumentException.class, () -> originalMatrix.withShape(Shape.scalar()));
+    NdArray<T> newMatrix = originalMatrix.withShape(Shape.of(3, 2));
+    assertEquals(Shape.of(3, 2), newMatrix.shape());
+  }
 }
diff --git a/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java b/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java
index 43779b3..4c4db5b 100644
--- a/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java
+++ b/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java
@@ -23,9 +23,7 @@
 import org.tensorflow.ndarray.impl.sparse.LongSparseNdArray;
 import org.tensorflow.ndarray.impl.sparse.ShortSparseNdArray;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class SparseNdArrayTest {
   long[][] indicesArray = {{0, 0}, {1, 2}, {2, 3}};
@@ -188,4 +186,10 @@ public void testShort() {
     assertEquals((short) 0, instance.getShort(2, 2));
     assertEquals((short) 0xff00, instance.getShort(2, 3));
   }
+
+  @Test
+  public void withShape() {
+    NdArray<?> sparseArray = NdArrays.sparseOf(indices, NdArrays.vectorOf(1, 2, 3), shape);
+    assertThrows(UnsupportedOperationException.class, () -> sparseArray.withShape(shape.prepend(1)));
+  }
 }

From b8ee47356f95a0c942dac3dedcc6c917dc857448 Mon Sep 17 00:00:00 2001
From: Karl Lessard <karl.lessard@gmail.com>
Date: Sat, 17 Feb 2024 20:51:34 -0500
Subject: [PATCH 2/3] Fix equality check and star import

---
 .../tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java  | 2 +-
 .../test/java/org/tensorflow/ndarray/NdArrayTestBase.java    | 2 +-
 .../test/java/org/tensorflow/ndarray/SparseNdArrayTest.java  | 5 ++++-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
index 8b5c5ba..8e21c11 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
@@ -54,7 +54,7 @@ public NdArraySequence<U> elements(int dimensionIdx) {
 
   @Override
   public U withShape(Shape shape) {
-    if (shape == this.shape()) {
+    if (shape.equals(this.shape())) {
       return (U)this;
     }
     if (shape == null || shape.isUnknown() || shape.size() != this.shape().size()) {
diff --git a/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java b/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java
index 35806f7..577c2b7 100644
--- a/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java
+++ b/ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java
@@ -400,7 +400,7 @@ public void withShape() {
     assertEquals(originalShape, originalArray.shape());
     assertEquals(valueOf(10L), newArray.getObject(0, 0));
 
-    NdArray<T> sameArray = originalArray.withShape(originalShape);
+    NdArray<T> sameArray = originalArray.withShape(Shape.scalar());
     assertSame(originalArray, sameArray);
 
     assertThrows(IllegalArgumentException.class, () -> originalArray.withShape(Shape.of(2)));
diff --git a/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java b/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java
index 4c4db5b..0c5d6b3 100644
--- a/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java
+++ b/ndarray/src/test/java/org/tensorflow/ndarray/SparseNdArrayTest.java
@@ -23,7 +23,10 @@
 import org.tensorflow.ndarray.impl.sparse.LongSparseNdArray;
 import org.tensorflow.ndarray.impl.sparse.ShortSparseNdArray;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class SparseNdArrayTest {
   long[][] indicesArray = {{0, 0}, {1, 2}, {2, 3}};

From b650c3729a35a9b08b078a9fa6967dc2314267bb Mon Sep 17 00:00:00 2001
From: Karl Lessard <karl.lessard@gmail.com>
Date: Sat, 17 Feb 2024 21:36:00 -0500
Subject: [PATCH 3/3] Check null before equality

---
 .../tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
index 8e21c11..baaf23e 100644
--- a/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
+++ b/ndarray/src/main/java/org/tensorflow/ndarray/impl/dense/AbstractDenseNdArray.java
@@ -54,12 +54,12 @@ public NdArraySequence<U> elements(int dimensionIdx) {
 
   @Override
   public U withShape(Shape shape) {
-    if (shape.equals(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());
     }
+    if (shape.equals(this.shape())) {
+      return (U)this;
+    }
     return instantiateView(buffer(), DimensionalSpace.create(shape));
   }