cuML bug report

repro.py

import cudf
from cuml.neighbors import NearestNeighbors
from cuml.datasets import make_blobs

X, _ = make_blobs(n_samples=10, centers=5,
                  n_features=10, random_state=42)

# build a cudf Dataframe
df_numeric = cudf.DataFrame(X)

mid_idx = len(df_numeric) // 2


# fit model
model = NearestNeighbors(n_neighbors=3, algorithm="brute")
model.fit(X)
NearestNeighbors()

distances, indices = model.kneighbors(df_numeric, two_pass_precision=True)
#distances, indices = model.kneighbors(df_numeric)

print("one large batch")
print(indices)
print("NOTE: 8  8  1  7 is the right order")
print(distances) 


df1 = df_numeric.iloc[:mid_idx]
df2 = df_numeric.iloc[mid_idx:]

distances1, indices1 = model.kneighbors(df1, two_pass_precision=True)
distances2, indices2 = model.kneighbors(df2, two_pass_precision=True)
#distances1, indices1 = model.kneighbors(df1)
#distances2, indices2 = model.kneighbors(df2)

distances = cudf.concat([distances1, distances2], ignore_index=True)
indices = cudf.concat([indices1, indices2], ignore_index=True)

print("two smaller batches")
print(indices)
print("NOTE: 8  8  1  7 is the right order")
print("here 1 and 7 are flipped!")
print(distances) 
print("it seems like the distances got shifted? 588.513428 is the distance to 7 which is right but the distance to 1 should be 26.002613")


df9th = df_numeric.iloc[8:9]

distances, indices = model.kneighbors(df9th, two_pass_precision=True)
#distances, indices = model.kneighbors(df9th)

print("one at a time")
print(indices)
print(distances) 
print("and when I look up one at a time the distances are 0?")