dist_batch_numba.py

@jit(nopython=True, nogil=True)
def dist_batch(points,  n_begin, n_end):
    #expd=np.expand_dims(points,2)
    #tiled=np.tile(expd, n_end-n_begin) #we tile up only (n_end-n_begin) times
    tiled = np.zeros((points.shape[0], points.shape[1], n_end-n_begin))
    for dim in range(n_end-n_begin):
        tiled[:,:,dim]=points
        
    selected=points[n_begin:n_end] # we select only part (a batch) of the whole dataset
    trans=np.transpose(selected) 
    diff=trans-tiled
    num=np.sum(np.square(diff), axis=1)

    den_sq_norm=1-np.sum(np.square(points),1)
    den_selected=den_sq_norm[n_begin:n_end] # we select only part (a batch) of the whole dataset
    den_expd=np.expand_dims(den_sq_norm,1)
    den=den_expd * den_selected.T

return np.arccosh(1+2*num/ den)