[Feature] Import subpackge based on backend of dgl.nn (dmlc#1305)

* automatically select nn backend * rename load_backend to private * switch to backend_name Co-authored-by: Jinjing Zhou <[email protected]>
shouwangbuqi · Mar 5, 2020 · e2db791 · e2db791
1 parent 93ac29c
commit e2db791
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/python/dgl/__init__.py b/python/dgl/__init__.py
@@ -8,7 +8,6 @@
 from .backend import load_backend
 
 from . import function
-from . import nn
 from . import contrib
 from . import container
 from . import random

diff --git a/python/dgl/nn/__init__.py b/python/dgl/nn/__init__.py
@@ -1 +1,12 @@
 """Package for neural network common components."""
+import importlib
+import sys
+from ..backend import backend_name
+
+def _load_backend(mod_name):
+    mod = importlib.import_module('.%s' % mod_name, __name__)
+    thismod = sys.modules[__name__]
+    for api, obj in mod.__dict__.items():
+        setattr(thismod, api, obj)
+
+_load_backend(backend_name)