public static void updateForNewAxes ( INDArray arr , INDArrayIndex ... indexes ) { int numNewAxes = NDArrayIndex . numNewAxis ( indexes ) ; if ( numNewAxes >= 1 && ( indexes [ 0 ] . length ( ) > 1 || indexes [ 0 ] instanceof NDArrayIndexAll ) ) { List < Long > newShape = new ArrayList <> ( ) ; List < Long > newStrides = new ArrayList <> ( ) ; int currDimension = 0 ; for ( int i = 0 ; i < indexes . length ; i ++ ) { if ( indexes [ i ] instanceof NewAxis ) { newShape . add ( 1L ) ; newStrides . add ( 0L ) ; } else { newShape . add ( arr . size ( currDimension ) ) ; newStrides . add ( arr . size ( currDimension ) ) ; currDimension ++ ; } } while ( currDimension < arr . rank ( ) ) { newShape . add ( ( long ) currDimension ) ; newStrides . add ( ( long ) currDimension ) ; currDimension ++ ; } long [ ] newShapeArr = Longs . toArray ( newShape ) ; long [ ] newStrideArr = Longs . toArray ( newStrides ) ; // FIXME: this is wrong, it breaks shapeInfo immutability arr . setShape ( newShapeArr ) ; arr . setStride ( newStrideArr ) ; } else { if ( numNewAxes > 0 ) { long [ ] newShape = Longs . concat ( ArrayUtil . toLongArray ( ArrayUtil . nTimes ( numNewAxes , 1 ) ) , arr . shape ( ) ) ; long [ ] newStrides = Longs . concat ( new long [ numNewAxes ] , arr . stride ( ) ) ; arr . setShape ( newShape ) ; arr . setStride ( newStrides ) ; } } }
Set the shape and stride for new axes based dimensions