dmlc--dgl
Minjie Wang
3c129ad71f
* cython nogil * move APIs to internal and add unit test * fix lint * disable callback array test
40 行
1.3 KiB
Python
40 行
1.3 KiB
Python
import dgl
|
|
import numpy as np
|
|
import backend as F
|
|
import unittest, pytest
|
|
import os
|
|
|
|
@unittest.skipIf(os.name == 'nt', reason='Cython only works on linux')
|
|
def test_cython():
|
|
import dgl._ffi._cy3.core
|
|
|
|
@pytest.mark.parametrize('arg', [1, 2.3])
|
|
def test_callback(arg):
|
|
def cb(x):
|
|
return x + 1
|
|
ret = dgl._api_internal._TestPythonCallback(cb, arg)
|
|
assert ret == arg + 1
|
|
|
|
@pytest.mark.parametrize('dtype', [F.float32, F.float64, F.int32, F.int64])
|
|
def _test_callback_array(dtype):
|
|
def cb(x):
|
|
return F.to_dgl_nd(F.from_dgl_nd(x) + 1)
|
|
arg = F.copy_to(F.tensor([1, 2, 3], dtype=dtype), F.ctx())
|
|
ret = F.from_dgl_nd(dgl._api_internal._TestPythonCallback(cb, F.to_dgl_nd(arg)))
|
|
assert np.allclose(F.asnumpy(ret), F.asnumpy(arg) + 1)
|
|
|
|
@pytest.mark.parametrize('arg', [1, 2.3])
|
|
def test_callback_thread(arg):
|
|
def cb(x):
|
|
return x + 1
|
|
ret = dgl._api_internal._TestPythonCallbackThread(cb, arg)
|
|
assert ret == arg + 1
|
|
|
|
@pytest.mark.parametrize('dtype', [F.float32, F.float64, F.int32, F.int64])
|
|
def _test_callback_array_thread(dtype):
|
|
def cb(x):
|
|
return F.to_dgl_nd(F.from_dgl_nd(x) + 1)
|
|
arg = F.copy_to(F.tensor([1, 2, 3], dtype=dtype), F.ctx())
|
|
ret = F.from_dgl_nd(dgl._api_internal._TestPythonCallbackThread(cb, F.to_dgl_nd(arg)))
|
|
assert np.allclose(F.asnumpy(ret), F.asnumpy(arg) + 1)
|