Make rpc and tensorlist part of PBG

Reviewed By: adamlerer

Differential Revision: D17742505

fbshipit-source-id: 6791a0a57127989f5af3ef0b63cee7b05f1b111f

Author: lw

test/test_entitylist.py

@@ -10,9 +10,9 @@
 from unittest import TestCase, main
 
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.entitylist import EntityList
+from torchbiggraph.tensorlist import TensorList
 
 
 def tensor_list_from_lists(lists: Sequence[Sequence[int]]) -> TensorList:

test/test_graph_storages.py

@@ -12,9 +12,9 @@
 import h5py
 import numpy as np
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.graph_storages import FileEdgeAppender
+from torchbiggraph.tensorlist import TensorList
 
 
 class TestFileEdgeAppender(TestCase):

test/test_model.py

@@ -14,7 +14,6 @@
 from unittest import TestCase, main
 
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.entitylist import EntityList
 from torchbiggraph.model import (
@@ -39,6 +38,7 @@
     TranslationOperator,
     match_shape,
 )
+from torchbiggraph.tensorlist import TensorList
 
 
 class TensorTestCase(TestCase):

torch_extensions/README.md

@@ -11,10 +11,9 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, NamedTuple, Optional, Set, Tuple
 
-from torch_extensions.rpc.rpc import Client, Server
-
 from torchbiggraph.config import BucketOrder
 from torchbiggraph.distributed import Startable
+from torchbiggraph.rpc import Client, Server
 from torchbiggraph.stats import Stats
 from torchbiggraph.types import Bucket, EntityName, Partition, Rank, Side
 

torch_extensions/__init__.py

@@ -9,8 +9,8 @@
 from typing import Any, Sequence, Union
 
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
+from torchbiggraph.tensorlist import TensorList
 from torchbiggraph.types import LongTensorType
 
 

torch_extensions/rpc/__init__.py

@@ -18,11 +18,11 @@
 import h5py
 import numpy as np
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.edgelist import EdgeList
 from torchbiggraph.entitylist import EntityList
 from torchbiggraph.plugin import URLPluginRegistry
+from torchbiggraph.tensorlist import TensorList
 from torchbiggraph.util import CouldNotLoadData