Fix len() for tensors (#13398)

David Riazati · facebook-github-bot · commit f9c0a08eed72 · 2018-10-31T13:13:21.000-07:00
Summary: Fixes #13376, `len(tensor)` was converting tensor to a 1 element list and returning 1 every time. Pull Request resolved: #13398 Differential Revision: D12867630 Pulled By: driazati fbshipit-source-id: 28f3580a072d763df0980b3149c49d1894842ec9
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -3511,6 +3511,12 @@ def bad_negative_index():
         self.checkScriptRaisesRegex(bad_negative_index, (), IndexError,
                                     "list index out of range")
 
+    def test_tensor_len(self):
+        def func(x):
+            return len(x)
+
+        self.checkScript(func, [torch.ones(4, 5, 6)])
+
     def test_list_len(self):
         def func():
             a = [1, 2, 3]
diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
@@ -768,7 +768,7 @@ RegisterOperators reg2({
 
 #define DEFINE_STRING_OP(op_name, string_op, result)                           \
 Operator(                                                                      \
-    #op_name "(str a, str b) ->" #result,                                \
+    #op_name "(str a, str b) ->" #result,                                      \
     [](Node* node) {                                                           \
       return [=](Stack& stack) {                                               \
         auto b = pop(stack).toStringRef();                                     \
@@ -781,8 +781,20 @@ Operator(                                                                      \
   DEFINE_STRING_OP(aten::eq, a == b, bool)
   DEFINE_STRING_OP(aten::ne, a != b, bool)
   DEFINE_STRING_OP(aten::add, a + b, str)
+#undef DEFINE_STRING_OP
 
-
+    // tensor length op (size of 1st dimension)
+    Operator(
+      "aten::len(Tensor t) -> int",
+      [](Stack& stack) {
+        at::Tensor t = pop(stack).toTensor();
+        if (t.dim() == 0) {
+          AT_ERROR("len() of a 0-d tensor");
+        }
+        push(stack, t.sizes()[0]);
+        return 0;
+      }
+    ),
 #define CREATE_LIST_OPS(decl_type, c_type) \
     Operator("aten::select(" decl_type "[] a, int b) -> " decl_type, listSelect<Shared<c_type>>), \
     Operator("aten::len(" decl_type "[] a) -> int", listLen<Shared<c_type>>), \
@@ -799,6 +811,7 @@ Operator(                                                                      \
     CREATE_LIST_OPS("float", DoubleList)
     CREATE_LIST_OPS("Tensor", TensorList)
     CREATE_LIST_OPS("t", GenericList)
+#undef CREATE_LIST_OPS
 
 
     Operator("aten::eq(int[] a, int[] b) -> int", listEq<Shared<IntList>>),
