Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
ffi.py 370 Bytes
Newer Older
rusty1s's avatar
clean up  
rusty1s committed
1
2
3
4
5
6
7
8 
from ..._ext import ffi


def get_func(name, tensor):
    cuda = '_cuda' if tensor.is_cuda else ''
    return getattr(ffi, 'cluster_{}{}'.format(name, cuda))
rusty1s's avatar
rename  
rusty1s committed
9 
def get_typed_func(name, tensor):
rusty1s's avatar
clean up  
rusty1s committed
10
11
12 
    typename = type(tensor).__name__.replace('Tensor', '')
    cuda = 'cuda_' if tensor.is_cuda else ''
    return getattr(ffi, 'cluster_{}_{}{}'.format(name, cuda, typename))