Fix bug when slicing on a segmented dimension (#2)

Author: karllessard

ndarray/src/main/java/org/tensorflow/ndarray/impl/dimension/DimensionalSpace.java

@@ -144,7 +144,7 @@ public DimensionalSpace from(int dimensionStart) {
       throw new IndexOutOfBoundsException();
     }
     Dimension[] newDimensions = Arrays.copyOfRange(dimensions, dimensionStart, dimensions.length);
-    if (segmentationIdx > dimensionStart) {
+    if (segmentationIdx >= dimensionStart) {
       return new DimensionalSpace(newDimensions, segmentationIdx - dimensionStart);
     }
     return new DimensionalSpace(newDimensions);

ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java

@@ -34,6 +34,7 @@
 import java.nio.BufferUnderflowException;
 import org.junit.jupiter.api.Test;
 import org.tensorflow.ndarray.buffer.DataBuffer;
+import org.tensorflow.ndarray.index.Indices;
 
 public abstract class NdArrayTestBase<T> {
 
@@ -335,4 +336,26 @@ public void equalsAndHashCode() {
     assertNotEquals(array1, array4);
     assertNotEquals(array1.hashCode(), array4.hashCode());
   }
+
+  @Test
+  public void iterateScalarsOnSegmentedElements() {
+    NdArray<T> originalTensor = allocate(Shape.of(2, 3));
+
+    originalTensor
+        .setObject(valueOf(0L), 0, 0)
+        .setObject(valueOf(1L), 0, 1)
+        .setObject(valueOf(2L), 0, 2)
+        .setObject(valueOf(3L), 1, 0)
+        .setObject(valueOf(4L), 1, 1)
+        .setObject(valueOf(5L), 1, 2);
+
+    NdArray<T> slice = originalTensor.slice(Indices.all(), Indices.sliceFrom(1));
+    assertEquals(Shape.of(2, 2), slice.shape());
+
+    slice.elements(0).forEachIndexed((eCoord, e) -> {
+      e.scalars().forEachIndexed((sCoord, s) -> {
+        assertEquals(valueOf((eCoord[0] * originalTensor.shape().get(1)) + sCoord[0] + 1), s.getObject());
+      });
+    });
+  }
 }