Wrong behavior of top_k_data() with cuda

[manabu0522]

I found a weird behavior when using top_k_data() function with cuda.
Here is a test code.
nnabla v1.35.1
nnabla-ext-cuda v1.35.1
cuda 11.6
cudnn 8.4

import numpy as np
import nnabla as nn
import nnabla.functions as F
import nnabla.ext_utils

ctx = nnabla.ext_utils.get_extension_context('cudnn', device_id=0)
nn.set_default_context(ctx)

topk = 3
data_size = 5
data = nn.Variable(shape=(data_size,))

# case 0
data.d = np.arange(data_size)
with nn.auto_forward():
    topk_data, topk_index = F.top_k_data(data, topk, reduce=True, with_index=True, base_axis=0)

print('Case 0')
print(f'Data: {data.d}')
print(f'TopK: {topk_data.d} (index: {topk_index.d})')

# case 1
data.d[:data_size-3] = data_size-3
with nn.auto_forward():
    topk_data, topk_index = F.top_k_data(data, topk, reduce=True, with_index=True, base_axis=0)

print('Case 1')
print(f'Data: {data.d}')
print(f'TopK: {topk_data.d} (index: {topk_index.d})')

The results are as follows. The case 1 is wrong. It should be TopK [4. 3. 2.].

Case 0
Data: [0 1 2 3 4]
TopK: [4. 3. 2.] (index: [4 3 2])
Case 1
Data: [2 2 2 3 4]
TopK: [2. 2. 2.] (index: [2 2 2])

With CPU, seems no problem.

Case 0
Data: [0 1 2 3 4]
TopK: [4. 3. 2.] (index: [4 3 2])
Case 1
Data: [2 2 2 3 4]
TopK: [4. 3. 2.] (index: [4 3 0])

[TomonobuTsujikawa]

Thank you for creating issue report.
We are trying to identify where the problem has, please wait a moment.

[TomonobuTsujikawa]

We have prepared a fix here. It will be released in the next version.
sony/nnabla-ext-cuda#473

[manabu0522]

Thanks. It works.