wehub-resource-sync
80 lines
1.7 KiB
Python
80 lines
1.7 KiB
Python
import importlib
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
from dgl.backend import *
|
|
from dgl.nn import *
|
|
|
|
from . import backend_unittest
|
|
|
|
mod = importlib.import_module(".%s" % backend_name, __name__)
|
|
thismod = sys.modules[__name__]
|
|
|
|
for api in backend_unittest.__dict__.keys():
|
|
if api.startswith("__"):
|
|
continue
|
|
elif callable(mod.__dict__[api]):
|
|
# Tensor APIs used in unit tests MUST be supported across all backends
|
|
globals()[api] = mod.__dict__[api]
|
|
|
|
# Tensor creation with default dtype and context
|
|
|
|
_zeros = zeros
|
|
_ones = ones
|
|
_randn = randn
|
|
_tensor = tensor
|
|
_arange = arange
|
|
_full = full
|
|
_full_1d = full_1d
|
|
_softmax = softmax
|
|
_default_context_str = os.getenv("DGLTESTDEV", "cpu")
|
|
_context_dict = {
|
|
"cpu": cpu(),
|
|
"gpu": cuda(),
|
|
}
|
|
_default_context = _context_dict[_default_context_str]
|
|
|
|
|
|
def ctx():
|
|
return _default_context
|
|
|
|
|
|
def gpu_ctx():
|
|
return _default_context_str == "gpu"
|
|
|
|
|
|
def zeros(shape, dtype=float32, ctx=_default_context):
|
|
return _zeros(shape, dtype, ctx)
|
|
|
|
|
|
def ones(shape, dtype=float32, ctx=_default_context):
|
|
return _ones(shape, dtype, ctx)
|
|
|
|
|
|
def randn(shape):
|
|
return copy_to(_randn(shape), _default_context)
|
|
|
|
|
|
def tensor(data, dtype=None):
|
|
return copy_to(_tensor(data, dtype), _default_context)
|
|
|
|
|
|
def arange(start, stop, dtype=int64, ctx=None):
|
|
return _arange(
|
|
start, stop, dtype, ctx if ctx is not None else _default_context
|
|
)
|
|
|
|
|
|
def full(shape, fill_value, dtype, ctx=_default_context):
|
|
return _full(shape, fill_value, dtype, ctx)
|
|
|
|
|
|
def full_1d(length, fill_value, dtype, ctx=_default_context):
|
|
return _full_1d(length, fill_value, dtype, ctx)
|
|
|
|
|
|
def softmax(x, dim):
|
|
return _softmax(x, dim)
|