From 78f589e794081abbc8cb001552444aec712d1f89 Mon Sep 17 00:00:00 2001 From: David Riazati Date: Tue, 16 Apr 2019 15:03:47 -0700 Subject: [PATCH] Add len() for strings (#19320) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/19320 ghimport-source-id: 62131cb24e9bf65f0ef3e60001cb36509a1f4163 Reviewed By: bethebunny Differential Revision: D14961078 Pulled By: driazati fbshipit-source-id: 08b9a4b10e4a47ea09ebf55a4743defa40c74698 --- test/test_jit.py | 9 +++++++++ torch/csrc/jit/register_prim_ops.cpp | 6 +++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/test/test_jit.py b/test/test_jit.py index 429cb26..bc47572 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -11308,6 +11308,15 @@ a") imported_m = self.getExportImportCopy(m) self.assertEqual(m(), imported_m()) + def test_string_len(self): + def fn(x): + # type: (str) -> int + return len(x) + + self.checkScript(fn, ("",)) + self.checkScript(fn, ("h",)) + self.checkScript(fn, ("hello",)) + @unittest.skipIf(IS_WINDOWS or IS_SANDCASTLE, "NYI: TemporaryFileName support for Windows or Sandcastle") def test_attribute_unpickling(self): class M(torch.jit.ScriptModule): diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp index 5ce746d..bf0bec2 100644 --- a/torch/csrc/jit/register_prim_ops.cpp +++ b/torch/csrc/jit/register_prim_ops.cpp @@ -1600,7 +1600,11 @@ RegisterOperators reg2({ DEFINE_STRING_OP(aten::ne, a != b, bool), DEFINE_STRING_OP(aten::add, a + b, str), #undef DEFINE_STRING_OP - + Operator("aten::len(str s) -> int", [](Stack& stack) { + auto string = pop(stack).toStringRef(); + push(stack, static_cast(string.size())); + return 0; + }), // tensor length op (size of 1st dimension) Operator( "aten::len(Tensor t) -> int", -- 2.7.4