code: public static long vectorsPerSlice(INDArray arr, int... rank) {
          if (arr.rank() > 2) {
              return arr.size(-2) * arr.size(-1);
          }
  
          return arr.size(-1);
  
      }
docs: The number of vectors
  in each slice of an ndarray.
  @param arr the array to
  get the number
  of vectors per slice for
  @param rank the dimensions to get the number of vectors per slice for
  @return the number of vectors per slice