start of adapting ndarray indexing

Signed-off-by: Ryan Nett <[email protected]>

Author: rnett

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

@@ -39,20 +39,38 @@ public RelativeDimensionalSpace mapTo(Index[] indices) {
       throw new ArrayIndexOutOfBoundsException();
     }
     int dimIdx = 0;
+    int indexIdx = 0;
     int newDimIdx = 0;
     int segmentationIdx = -1;
     long initialOffset = 0;
 
-    Dimension[] newDimensions = new Dimension[dimensions.length];
-    while (dimIdx < indices.length) {
+    int newAxes = 0;
+    boolean seenEllipsis = false;
+    for(Index idx : indices){
+      if(idx.isNewAxis()){
+        newAxes += 1;
+      }
+      if(idx.isEllipsis()){
+        if(seenEllipsis){
+          throw new IllegalArgumentException("Only one ellipsis allowed");
+        } else {
+          seenEllipsis = true;
+        }
+      }
+    }
+    int newLength = dimensions.length + newAxes;
+
+    Dimension[] newDimensions = new Dimension[newLength];
+    while (indexIdx < indices.length) {
 
       if (indices[dimIdx].isPoint()) {
         // When an index targets a single point in a given dimension, calculate the offset of this
         // point and cumulate the offset of any subsequent point as well
         long offset = 0;
         do {
-          offset += indices[dimIdx].mapCoordinate(0, dimensions[dimIdx]);
-        } while (++dimIdx < indices.length && indices[dimIdx].isPoint());
+          offset += indices[indexIdx].mapCoordinate(0, dimensions[dimIdx]);
+          dimIdx++;
+        } while (++indexIdx < indices.length && indices[indexIdx].isPoint());
 
         // If this is the first index, then the offset is the position of the whole dimension
         // space within the original one. If not, then we apply the offset to the last vectorial
@@ -65,14 +83,47 @@ public RelativeDimensionalSpace mapTo(Index[] indices) {
           segmentationIdx = newDimIdx - 1;
         }
 
+      } else if(indices[indexIdx].isNewAxis()) {
+        long newSize;
+        if(dimIdx == 0){
+          // includes everything.  Should really include future reduction (at()) but that doesn't seem to cause issues
+          //   elsewhere
+          newSize = dimensions[0].numElements() * dimensions[0].elementSize();
+        } else {
+          newSize = dimensions[dimIdx - 1].elementSize();
+        }
+
+        newDimensions[newDimIdx] = new Axis(1, newSize);
+        segmentationIdx = newDimIdx; // is this correct?
+        ++newDimIdx;
+        ++indexIdx;
+      } else if(indices[indexIdx].isEllipsis()){
+        int remainingDimensions = dimensions.length - dimIdx;
+        int requiredDimensions = 0;
+        for(int i = indexIdx + 1 ; i < indices.length ; i++){
+          if(!indices[i].isNewAxis()){
+            requiredDimensions++;
+          }
+        }
+        // while the number of dimensions left < the number of indices that consume axes
+        while(remainingDimensions > requiredDimensions){
+          Dimension dim = dimensions[dimIdx++];
+          if (dim.isSegmented()) {
+            segmentationIdx = newDimIdx;
+          }
+          newDimensions[newDimIdx++] = dim;
+          remainingDimensions--;
+        }
+        indexIdx++;
       } else {
         // Map any other index to the appropriate dimension of this space
-        Dimension newDimension = indices[dimIdx].apply(dimensions[dimIdx++]);
+        Dimension newDimension = indices[indexIdx].apply(dimensions[dimIdx++]);
         newDimensions[newDimIdx] = newDimension;
         if (newDimension.isSegmented()) {
           segmentationIdx = newDimIdx;
         }
         ++newDimIdx;
+        ++indexIdx;
       }
     }
 

ndarray/src/main/java/org/tensorflow/ndarray/index/All.java

@@ -18,7 +18,7 @@
 
 import org.tensorflow.ndarray.impl.dimension.Dimension;
 
-final class All implements Index {
+final class All implements TensorIndex {
 
   static final All INSTANCE = new All();
 
@@ -39,4 +39,14 @@ public Dimension apply(Dimension dim) {
 
   private All() {
   }
+
+  @Override
+  public boolean beginMask() {
+    return true;
+  }
+
+  @Override
+  public boolean endMask() {
+    return true;
+  }
 }

ndarray/src/main/java/org/tensorflow/ndarray/index/At.java

@@ -18,7 +18,7 @@
 
 import org.tensorflow.ndarray.impl.dimension.Dimension;
 
-final class At implements Index {
+final class At implements TensorIndex {
 
   @Override
   public long numElements(Dimension dim) {
@@ -27,22 +27,45 @@ public long numElements(Dimension dim) {
 
   @Override
   public long mapCoordinate(long coordinate, Dimension dim) {
+    long coord = this.coord > 0 ? this.coord : dim.numElements() - this.coord;
     return dim.positionOf(coord); // TODO validate coordinate is 0?
   }
 
   @Override
   public Dimension apply(Dimension dim) {
-    throw new IllegalStateException(); // FIXME?
+    if(keepDim){
+      return dim.withIndex(this);
+    }
+    else {
+      throw new IllegalStateException(); // FIXME?
+    }
   }
 
   @Override
   public boolean isPoint() {
-    return true;
+    return !keepDim;
   }
 
-  At(long coord) {
+  At(long coord, boolean keepDim) {
     this.coord = coord;
+    this.keepDim = keepDim;
   }
 
   private final long coord;
+  private final boolean keepDim;
+
+  @Override
+  public long begin() {
+    return coord;
+  }
+
+  @Override
+  public long end() {
+    return coord + 1;
+  }
+
+  @Override
+  public boolean shrinkAxisMask() {
+    return !keepDim;
+  }
 }

ndarray/src/main/java/org/tensorflow/ndarray/index/Ellipsis.java

@@ -0,0 +1,48 @@
+/*
+  Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ ==============================================================================
+ */
+package org.tensorflow.ndarray.index;
+
+import org.tensorflow.ndarray.impl.dimension.Dimension;
+
+final class Ellipsis implements TensorIndex{
+
+  static final Ellipsis INSTANCE = new Ellipsis();
+
+  private Ellipsis(){
+
+  }
+
+  @Override
+  public long numElements(Dimension dim) {
+    throw new IllegalStateException();
+  }
+
+  @Override
+  public long mapCoordinate(long coordinate, Dimension dim) {
+    throw new IllegalStateException();
+  }
+
+  @Override
+  public boolean isEllipsis() {
+    return true;
+  }
+
+  @Override
+  public boolean ellipsisMask() {
+    return true;
+  }
+}

ndarray/src/main/java/org/tensorflow/ndarray/index/Index.java

@@ -74,4 +74,19 @@ default Dimension apply(Dimension dim) {
   default boolean isPoint() {
     return false;
   }
+
+  /**
+   * Returns true if this index is a new axis, adding a dimension of size 1
+   */
+  default boolean isNewAxis() {
+    return false;
+  }
+
+  /**
+   * Returns true if this index is an ellipsis, expanding to take as many dimensions as possible
+   * (and applying all() to them)
+   */
+  default boolean isEllipsis() {
+    return false;
+  }
 }

ndarray/src/main/java/org/tensorflow/ndarray/index/Indices.java

@@ -40,8 +40,8 @@ public final class Indices {
    * @param coord coordinate of the element on the indexed axis
    * @return index
    */
-  public static Index at(long coord) {
-    return new At(coord);
+  public static TensorIndex at(long coord) {
+    return new At(coord, false);
   }
 
   /**
@@ -54,11 +54,50 @@ public static Index at(long coord) {
    * @return index
    * @throws IllegalRankException if {@code coord} is not a scalar (rank 0)
    */
-  public static Index at(NdArray<? extends Number> coord) {
+  public static TensorIndex at(NdArray<? extends Number> coord) {
     if (coord.rank() > 0) {
       throw new IllegalRankException("Only scalars are accepted as a value index");
     }
-    return new At(coord.getObject().longValue());
+    return new At(coord.getObject().longValue(), false);
+  }
+
+  /**
+   * A coordinate that selects a specific element on a given dimension.
+   *
+   * <p>When this index is applied to a given dimension, the dimension is resolved as a
+   * single element and therefore, if {@code keepDim} is false, is excluded from the computation of the rank.
+   * If {@code} keepDim is true, the dimension is collapsed down to one element.
+   *
+   * <p>For example, given a 3D matrix on the axis [x, y, z], if
+   * {@code matrix.slice(all(), at(0), at(0)}, then the rank of the returned slice is 1 and its
+   * number of elements is {@code x.numElements()}
+   *
+   * @param coord coordinate of the element on the indexed axis
+   * @param keepDim whether to remove the dimension.
+   * @return index
+   */
+  public static TensorIndex at(long coord, boolean keepDim) {
+    return new At(coord, keepDim);
+  }
+
+  /**
+   * A coordinate that selects a specific element on a given dimension.
+   *
+   * <p>This is equivalent to call {@link #at(long, boolean)} but where the value of the coordinate is
+   * provided by an N-dimensional array.
+   * <p>
+   * If {@code} keepDim is true, the dimension is collapsed down to one element instead of being removed.
+   *
+   * @param coord scalar indicating the coordinate of the element on the indexed axis
+   * @return index
+   * @param keepDim whether to remove the dimension.
+   * @throws IllegalRankException if {@code coord} is not a scalar (rank 0)
+   */
+  public static TensorIndex at(NdArray<? extends Number> coord, boolean keepDim) {
+    if (coord.rank() > 0) {
+      throw new IllegalRankException("Only scalars are accepted as a value index");
+    }
+    return new At(coord.getObject().longValue(), keepDim);
   }
 
   /**
@@ -72,7 +111,7 @@ public static Index at(NdArray<? extends Number> coord) {
    *
    * @return index
    */
-  public static Index all() {
+  public static TensorIndex all() {
     return All.INSTANCE;
   }
 
@@ -216,4 +255,22 @@ public static Index flip() {
   public static Index hyperslab(long start, long stride, long count, long block) {
     return new Hyperslab(start, stride, count, block);
   }
+
+  //TODO comments, tests, remove extra classes in favor of helper methods
+
+  public static TensorIndex newAxis(){
+    return NewAxis.INSTANCE;
+  }
+
+  public static TensorIndex ellipsis(){
+    return Ellipsis.INSTANCE;
+  }
+
+  public static TensorIndex expand(){
+    return ellipsis();
+  }
+
+  public static TensorIndex slice(Long start, Long end, long stride){
+    return new Slice(start, end, stride);
+  }
 }

ndarray/src/main/java/org/tensorflow/ndarray/index/NewAxis.java

@@ -0,0 +1,53 @@
+/*
+  Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ ==============================================================================
+ */
+package org.tensorflow.ndarray.index;
+
+import org.tensorflow.ndarray.impl.dimension.Dimension;
+
+final class NewAxis implements TensorIndex {
+
+  static final NewAxis INSTANCE = new NewAxis();
+
+  private NewAxis(){
+
+  }
+
+  @Override
+  public long numElements(Dimension dim) {
+    return 1;
+  }
+
+  @Override
+  public long mapCoordinate(long coordinate, Dimension dim) {
+    return coordinate;
+  }
+
+  @Override
+  public Dimension apply(Dimension dim) {
+    throw new IllegalStateException();
+  }
+
+  @Override
+  public boolean isNewAxis() {
+    return true;
+  }
+
+  @Override
+  public boolean newAxisMask() {
+    return true;
+  }
+}