[resubmit] try to infer rref type from python (pytorch#33992)

Summary:
Pull Request resolved: pytorch#33992

resubmit of pytorch#33369 with tweaks on when the rref type being created to ensure ivalue->type() hold the correct RRef type inside of inner element type.

Test Plan: Imported from OSS

Differential Revision: D20175043

Pulled By: wanchaol

fbshipit-source-id: a08b178e989c995632374e6c868d23c5a85526ae

Author: wanchaol

aten/src/ATen/core/ivalue.cpp

@@ -48,7 +48,7 @@ TypePtr IValue::type() const {
     case Tag::Future:
       return toFuture()->type();
     case Tag::RRef:
-      return toRRef()->type();
+      return RRefType::create(toRRef()->type());
     case Tag::Device:
       return DeviceObjType::get();
     case Tag::Object:

torch/csrc/jit/python/pybind_utils.h

@@ -177,6 +177,11 @@ inline InferredType tryToInferType(py::handle input) {
   if (py::isinstance<script::Object>(input)) {
     auto object = py::cast<script::Object>(input);
     return InferredType(object.type());
+#ifdef USE_DISTRIBUTED
+  } else if (py::isinstance<torch::distributed::rpc::PyRRef>(input)) {
+    auto rref_ivalue = input.cast<torch::distributed::rpc::PyRRef>().toIValue();
+    return InferredType(rref_ivalue.type());
+#endif
   }
 
   // Try container types

torch/testing/_internal/distributed/rpc/jit/rpc_test.py

@@ -17,6 +17,23 @@ def one_arg(value):
     return value + 1
 
 
+class MyScriptModuleWithRRefs(torch.jit.ScriptModule):
+    def __init__(self, dst_worker):
+        super().__init__()
+        self.rrefs = []
+        for _ in range(4):
+            self.rrefs.append(rpc_return_rref(dst_worker))
+
+    @torch.jit.script_method
+    def forward(self):
+        # type: () -> Tensor
+        res_tensor = torch.ones(2, 2)
+        for rref in self.rrefs:
+            res_tensor += rref.to_here()
+
+        return res_tensor
+
+
 @torch.jit.script
 class MyScriptClass:
     def __init__(self):
@@ -205,6 +222,15 @@ def rref_tensor_is_owner(rref_var):
         res = rref_tensor_is_owner(rref_var)
         self.assertEqual(res, False)
 
+    @dist_init
+    def test_my_script_module_with_rrefs(self):
+        n = self.rank + 1
+        dst_rank = n % self.world_size
+
+        module_with_rrefs = MyScriptModuleWithRRefs("worker{}".format(dst_rank))
+        res = module_with_rrefs()
+        self.assertEqual(res, torch.ones(2, 2) * 9)
+
     @dist_init
     def test_rref_python_annotation(self):
         n = self.rank + 1