pytorch_unique_timing

Code to test:

x = torch.randint(2, (1000, 1000))

torch_times = {}
torch_tmp = []
for dim in range(2):
    for _ in range(10):
        t0 = time.perf_counter()
        tu, ti = torch.unique(x, return_inverse=True, dim=dim)
        t1 = time.perf_counter()
        torch_tmp.append(t1 - t0)
    torch_times['dim' + str(dim)] = np.mean(torch_tmp)

numpy_times = {}
numpy_tmp = []
n = x.numpy()
for dim in range(2):
    for _ in range(10):
        t0 = time.perf_counter()
        nu, ni = np.unique(n, return_inverse=True, axis=dim)
        t1 = time.perf_counter()
        numpy_tmp.append(t1 - t0)
    numpy_times['dim' + str(dim)] = np.mean(numpy_tmp)
    
print("PyTorch:\n", torch_times)
print("Numpy:\n", numpy_times)

======================================================================================
Timings (in seconds)

Full code:
PyTorch:
 {'dim0': 0.420559009705903, 'dim1': 0.5335184658004437}
Numpy:
 {'dim0': 0.03593983739847317, 'dim1': 0.03847570090438239}

--------------------------------------------------------------------------------------

Without inverse_indices_dim calculation (return empty tensor instead):
PyTorch:
 {'dim0': 0.24947218560846524, 'dim1': 0.24243437854747754}
Numpy:
 {'dim0': 0.03328046520473436, 'dim1': 0.03735014664998744}

--------------------------------------------------------------------------------------

Without lambda sort function (return false everytime in sort):
PyTorch:
 {'dim0': 0.2105050868995022, 'dim1': 0.31288547125295735}
Numpy:
 {'dim0': 0.03443476958782412, 'dim1': 0.03836798309639562}

--------------------------------------------------------------------------------------

Without inverse_indices_dim & lambda sort function:
PyTorch:
 {'dim0': 0.0034200293943285943, 'dim1': 0.004693126596976072}
Numpy:
 {'dim0': 0.03603356489911676, 'dim1': 0.04012758959725034}
 
 --------------------------------------------------------------------------------------

Improved version full code:

PyTorch:
 {'dim0': 0.015360455599147827, 'dim1': 0.01664820614969358}
Numpy:
 {'dim0': 0.037852845003362744, 'dim1': 0.04158545115205925}

======================================================================================

CUDA

Full code:

PyTorch:
 {'dim0': 0.1933254515985027, 'dim1': 0.18684556494408752}
Numpy:
 {'dim0': 0.03505929970997386, 'dim1': 0.03857367945602164}

--------------------------------------------------------------------------------------

Without inverse_indices_dim calculation:

PyTorch:
 {'dim0': 0.08246875299955718, 'dim1': 0.08229753980122041}
Numpy:
 {'dim0': 0.044091779983136806, 'dim1': 0.04337026579014491}

--------------------------------------------------------------------------------------

Without reshape back:

PyTorch:
 {'dim0': 0.18773792859283275, 'dim1': 0.18284180923947133}
Numpy:
 {'dim0': 0.04005314850364812, 'dim1': 0.040591494156979026}

--------------------------------------------------------------------------------------

Without lambda sort function:

PyTorch:
 {'dim0': 0.1631539837049786, 'dim1': 0.1583032851020107}
Numpy:
 {'dim0': 0.03965735248639248, 'dim1': 0.040714282492990604}
 
 --------------------------------------------------------------------------------------

Without thrust::sort call:

PyTorch:
 {'dim0': 0.17885396589408628, 'dim1': 0.17443836564780213}
Numpy:
 {'dim0': 0.03556413790211081, 'dim1': 0.03808894795365632}

--------------------------------------------------------------------------------------

Without thrust::unique:

PyTorch:
 {'dim0': 0.18713290240266361, 'dim1': 0.19582855790504256}
Numpy:
 {'dim0': 0.03598183499416337, 'dim1': 0.04059765480051283}
 
 --------------------------------------------------------------------------------------

Without copy for sorting:

PyTorch:
 {'dim0': 0.14137323559843934, 'dim1': 0.14095361944346224}
Numpy:
 {'dim0': 0.03910874019493349, 'dim1': 0.04177295434928965}