docs: INDArray bucketData(INDArray query){
          INDArray mask = bucket(query);
          int nRes = mask.sum(0).getInt(0);
          INDArray res = Nd4j.create(new int[] {nRes, inDimension});
          int j = 0;
          for (int i = 0; i < nRes; i++){
              while (mask.getInt(j) == 0 && j < mask.length() - 1) {
                  j += 1;
              }
              if (mask.getInt(j) == 1) res.putRow(i, indexData.getRow(j));
              j += 1;
          }
          return res;
      }