code: public static long sliceOffsetForTensor(int index, INDArray arr, int[] tensorShape) {
          long tensorLength = ArrayUtil.prodLong(tensorShape);
          long lengthPerSlice = NDArrayMath.lengthPerSlice(arr);
          long offset = index * tensorLength / lengthPerSlice;
          return offset;
      }
docs: calculates the offset for a tensor
  @param index
  @param arr
  @param tensorShape
  @return