add index and count to list (pytorch#17446)

Summary:
see pytorch#16662
Pull Request resolved: pytorch#17446

Differential Revision: D14461293

Pulled By: Krovatkin

fbshipit-source-id: 03572467cdf85efc909c1864c0558a93085c8ff3

Author: TheCodez

test/test_jit.py

@@ -4315,6 +4315,76 @@ def test_list_remove():
             return a == [1, 2, 4]
         self.checkScript(test_list_remove, ())
 
+    def test_list_index_not_existing(self):
+        @torch.jit.script
+        def list_index_not_existing():
+            a = [4, 1, 3, 2]
+            i = a.index(5)
+
+            return i
+
+        with self.assertRaisesRegex(RuntimeError, "'5' is not in list"):
+            list_index_not_existing()
+
+    def test_list_index(self):
+        def list_index():
+            a = [4, 1, 3, 2]
+            i = a.index(3)
+
+            return i == 2
+        self.checkScript(list_index, ())
+
+    def test_tensor_list_index(self):
+        def tensor_list_index():
+            a = [torch.tensor(4), torch.tensor(1), torch.tensor(3), torch.tensor(2)]
+            i = a.index(torch.tensor(3))
+
+            return i == 2
+        self.checkScript(tensor_list_index, ())
+
+    def test_tensor_list_index_not_existing(self):
+        @torch.jit.script
+        def tensor_list_index_not_existing():
+            a = [torch.tensor(4), torch.tensor(1), torch.tensor(3), torch.tensor(2)]
+            i = a.index(torch.tensor(5))
+
+            return i
+
+        with self.assertRaisesRegex(RuntimeError, "is not in list"):
+            tensor_list_index_not_existing()
+
+    def test_list_count(self):
+        def list_count():
+            a = [4, 1, 4, 2, 4]
+            i = a.count(4)
+
+            return i == 3
+        self.checkScript(list_count, ())
+
+    def test_list_count_not_existing(self):
+        def list_count_not_existing():
+            a = [4, 1, 4, 2, 4]
+            i = a.count(5)
+
+            return i == 0
+        self.checkScript(list_count_not_existing, ())
+
+    def test_tensor_list_count(self):
+        def tensor_list_count():
+            a = [torch.tensor(4), torch.tensor(1), torch.tensor(4), torch.tensor(4)]
+            i = a.count(torch.tensor(4))
+
+            return i == 3
+        self.checkScript(tensor_list_count, ())
+
+    def test_tensor_list_count_not_existing(self):
+        def tensor_list_count_not_existing():
+            a = [torch.tensor(4), torch.tensor(1), torch.tensor(4), torch.tensor(4)]
+            i = a.count(torch.tensor(5))
+
+            return i == 0
+        self.checkScript(tensor_list_count_not_existing, ())
+
     def test_mutable_list_remove_tensor(self):
         def test_list_remove_tensor():
             a = [torch.ones(1), torch.zeros(1), torch.ones(2)]

torch/csrc/jit/register_prim_ops.cpp

@@ -54,8 +54,7 @@ void checkImplicitTensorToNum(at::Tensor t, bool toInt) {
         "Cannot input a tensor of dimension other than 0 as a scalar argument");
   }
   if (toInt &&
-      !isIntegralType(
-          autograd::as_variable_ref(t).data().scalar_type())) {
+      !isIntegralType(autograd::as_variable_ref(t).data().scalar_type())) {
     std::stringstream ss;
     ss << "Cannot input a tensor of type " << t.scalar_type()
        << " as an integral argument";
@@ -1114,6 +1113,76 @@ int listRemove<Shared<TensorList>, at::Tensor>(Stack& stack) {
   return 0;
 }
 
+template <typename TList, typename TElement>
+int listIndex(Stack& stack) {
+  TList list;
+  TElement elem;
+  pop(stack, list, elem);
+
+  auto& elements = list->elements();
+  auto pos = std::find(elements.begin(), elements.end(), elem);
+
+  if (pos != elements.end()) {
+    push(stack, static_cast<int64_t>(std::distance(elements.begin(), pos)));
+  } else {
+    AT_ERROR("'", elem, "' is not in list");
+  }
+
+  return 0;
+}
+
+template <>
+int listIndex<Shared<TensorList>, at::Tensor>(Stack& stack) {
+  Shared<TensorList> list;
+  at::Tensor elem;
+  pop(stack, list, elem);
+
+  auto& elements = list->elements();
+  auto pos = std::find_if(
+      elements.begin(), elements.end(), [elem](const at::Tensor& b) {
+        const auto cmp_result = elem.eq(b);
+        return cmp_result.is_nonzero();
+      });
+
+  if (pos != elements.end()) {
+    push(stack, static_cast<int64_t>(std::distance(elements.begin(), pos)));
+  } else {
+    AT_ERROR("'", elem, "' is not in list");
+  }
+
+  return 0;
+}
+
+template <typename TList, typename TElement>
+int listCount(Stack& stack) {
+  TList list;
+  TElement elem;
+  pop(stack, list, elem);
+
+  auto& elements = list->elements();
+  const int64_t count = std::count(elements.begin(), elements.end(), elem);
+  push(stack, count);
+
+  return 0;
+}
+
+template <>
+int listCount<Shared<TensorList>, at::Tensor>(Stack& stack) {
+  Shared<TensorList> list;
+  at::Tensor elem;
+  pop(stack, list, elem);
+
+  auto& elements = list->elements();
+  const int64_t count = std::count_if(
+      elements.begin(), elements.end(), [elem](const at::Tensor& b) {
+        const auto cmp_result = elem.eq(b);
+        return cmp_result.is_nonzero();
+      });
+  push(stack, count);
+
+  return 0;
+}
+
 template <typename TList>
 Operation listExtend(const Node* node) {
   return [](Stack& stack) {
@@ -1485,6 +1554,12 @@ RegisterOperators reg2({
     Operator(
         "aten::remove(Tensor[](a!) self, Tensor el) -> ()",
         listRemove<Shared<TensorList>, at::Tensor>),
+    Operator(
+        "aten::index(Tensor[] self, Tensor el) -> int",
+        listIndex<Shared<TensorList>, at::Tensor>),
+    Operator(
+        "aten::count(Tensor[] self, Tensor el) -> int",
+        listCount<Shared<TensorList>, at::Tensor>),
 
 // Mutable ops for lists containing immutable types.
 #define CREATE_IMMUTABLE_LIST_OPS(decl_type, c_type)                   \
@@ -1524,6 +1599,16 @@ RegisterOperators reg2({
           "[](a!) self,                      \
           " decl_type " el) -> ()",                                    \
           listRemove<Shared<c_type>, c_type::ElemType>),               \
+      Operator(                                                        \
+          "aten::index(" decl_type                                     \
+          "[] self,                           \
+          " decl_type " el) -> int",                                   \
+          listIndex<Shared<c_type>, c_type::ElemType>),                \
+      Operator(                                                        \
+          "aten::count(" decl_type                                     \
+          "[] self,                           \
+          " decl_type " el) -> int",                                   \
+          listCount<Shared<c_type>, c_type::ElemType>),                \
       Operator(                                                        \
           "aten::pop(" decl_type                                       \
           "[](a!) self, int idx=-1)             \