Python LazyTensor.__matmul__ Examples

Example #1

File: test_lazytensor_gaussian.py Project: dreedxerniitb/keops

def fun(x,y,b,backend):
    if "keops" in backend:
        x = LazyTensor(x)
        y = LazyTensor(y)
    Dxy = ((x-y)**2).sum(dim=2) 
    Kxy = (- Dxy).exp() 
    if backend=="keops_old":
        out = LazyTensor.__matmul__(Kxy,b,optional_flags=['-DENABLE_FINAL_CHUNKS=0'])
    else:
        out = Kxy @ b
    if device_id != 'cpu':
        torch.cuda.synchronize() 
    #print("out:",out.flatten()[:10])
    return out

Example #2

def fun(x, y, b, backend):
    if "keops" in backend:
        x = LazyTensor(x)
        y = LazyTensor(y)
    Dxy = ((x - y) ** 2).sum(dim=2)
    Kxy = (-Dxy).exp()
    if backend == "keops":
        out = LazyTensor.__matmul__(Kxy, b, backend="CPU")
    else:
        out = Kxy @ b
    if device_id != "cpu":
        torch.cuda.synchronize()
    # print("out:",out.flatten()[:10])
    return out

Example #3

File: test_lazytensor_gaussian_batch.py Project: wayne9qiu/keops

def fun(x,y,b,backend):
    if backend=="keops":
        x = LazyTensor(x)
        y = LazyTensor(y)
    elif backend!="torch":
        raise ValueError("wrong backend")
    Dxy = ((x-y)**2).sum(dim=4) 
    Kxy = (- Dxy).exp() 
    if backend=="keops":
        out = LazyTensor.__matmul__(Kxy,b,optional_flags=['-DENABLE_FINAL_CHUNKS=1','-DDIMFINALCHUNK=64'])
    else:
        out = Kxy @ b
    if device_id != 'cpu':
        torch.cuda.synchronize() 
    #print("out:",out.flatten()[:10])
    return out