Use pybind_utils for type inference

davidriazati · davidriazati · commit 11300f242068 · 2019-06-12T17:25:16.000-07:00
diff --git a/torch/csrc/jit/init.cpp b/torch/csrc/jit/init.cpp
@@ -338,6 +338,15 @@ void initJITBindings(PyObject* module) {
       .def(
           "_jit_set_first_class_mode",
           [](bool enabled) { script::getFirstClassMode() = enabled; })
+      .def(
+          "_jit_try_infer_type",
+          [](py::object obj) -> TypePtr {
+            auto match = tryToInferType(obj);
+            if (match.type) {
+              return *match.type;
+            }
+            return nullptr;
+          })
       .def(
           "_jit_fuser_get_fused_kernel_code",
           [](Graph& g, std::vector<at::Tensor> inps) {
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -1646,7 +1646,7 @@ def __init__(self, original, stubs):
                 if name in class_annotations:
                     the_type = torch.jit.annotations.ann_to_type(class_annotations[name])
                 else:
-                    the_type = torch.jit.annotations.try_to_infer_type(item)
+                    the_type = torch._C._jit_try_infer_type(item)
                 if the_type is not None:
                     self._c._register_attribute(name, the_type, item)
 
diff --git a/torch/jit/annotations.py b/torch/jit/annotations.py
@@ -196,48 +196,6 @@ def as_ann(ann):
     return_type = ann_to_type(as_ann(sig.return_annotation))
     return arg_types, return_type
 
-def try_to_infer_type(value):
-    if type(value) is float:
-        return FloatType.get()
-    elif type(value) is int:
-        a = IntType.get()
-        return a
-    elif type(value) is str:
-        return StringType.get()
-    elif type(value) is bool:
-        return BoolType.get()
-    elif type(value) is torch.Tensor:
-        return TensorType.get()
-    elif type(value) is list:
-        if len(value) == 0:
-            return None
-        element_type = None
-        for item in value:
-            item_element_type = try_to_infer_type(item)
-            if element_type is None:
-                element_type = item_element_type
-        if element_type is None:
-            return None
-        return ListType(element_type)
-    elif type(value) is tuple:
-        if len(value) == 0:
-            return None
-        types = [try_to_infer_type(item) for item in value]
-        if any(elem is None for elem in types):
-            return None
-        return TupleType(types)
-    elif type(value) is dict:
-        if len(value) == 0:
-            return None
-        an_entry = next(iter(value.items()))
-        # TODO: use value.itervalues().next() for PY2 so it's lazily evaluated
-        key_type = try_to_infer_type(an_entry[0])
-        value_type = try_to_infer_type(an_entry[1])
-        if key_type is None or value_type is None:
-            return None
-        return DictType(key_type, value_type)
-    return None
-
 
 def ann_to_type(ann):
     if ann is None:
