Add list(string) (#20617)

davidriazati · facebook-github-bot · commit c819d767893b · 2019-05-17T14:57:20.000-07:00
Summary: Stack from [ghstack](https://github.com/ezyang/ghstack): * **#20617 [jit] Add list(string)** Pull Request resolved: #20617 Pulled By: driazati Differential Revision: D15397739 fbshipit-source-id: 55a1e54f0ed2d9fd3ce44a8bb64bb4ba63181c9a
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -12824,6 +12824,13 @@ def test_pickle_checkpoint(self):
         self._test_pickle_checkpoint('cpu')
         self._test_pickle_checkpoint_views('cpu')
 
+    def test_string_list(self):
+        def fn(string):
+            # type: (str) -> List[str]
+            return list(string)
+
+        self.checkScript(fn, ("abcdefgh",))
+
     def test_split(self):
         def split_two(tensor):
             a, b, c = torch.split(tensor, 2, dim=1)
diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
@@ -1717,6 +1717,19 @@ RegisterOperators reg2({
           push(stack, t.sizes()[0]);
           return 0;
         }),
+    Operator(
+        "aten::list(str t) -> str[]",
+        [](Stack& stack) {
+          auto str = pop(stack).toStringRef();
+          std::vector<IValue> chars;
+          chars.reserve(str.size());
+          for (auto c : str) {
+            chars.push_back(std::string(1, c));
+          }
+          push(stack, chars);
+          return 0;
+        }
+    ),
 // Mutable ops for lists containing mutable types.
 #define CREATE_MUTABLE_LIST_OPS(decl_type, c_type)                          \
   Operator(                                                                 \
