From 9726651d1e89f8f1e34a038c81a1cc579981f225 Mon Sep 17 00:00:00 2001
From: Duc Ngo <duc@fb.com>
Date: Thu, 13 Dec 2018 20:43:00 -0800
Subject: [PATCH] caffe2 - easy - test utils for tensor assertion (#15020)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15020

Add test utils for assertion of a tensor (sizes and values)

Reviewed By: salexspb

Differential Revision: D13401146

fbshipit-source-id: bc385df074043e03ea884940b5631b96de4a607e
---
 caffe2/core/test_utils.cc |  8 ++++++++
 caffe2/core/test_utils.h  | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/caffe2/core/test_utils.cc b/caffe2/core/test_utils.cc
index 6b58240..070ec34 100644
--- a/caffe2/core/test_utils.cc
+++ b/caffe2/core/test_utils.cc
@@ -18,6 +18,13 @@ void assertTensorEqualsWithType(
 namespace caffe2 {
 namespace testing {
 
+// Asserts that two float values are close within epsilon.
+void assertNear(float value1, float value2, float epsilon) {
+  // These two enforces will give good debug messages.
+  CAFFE_ENFORCE_LE(value1, value2 + epsilon);
+  CAFFE_ENFORCE_GE(value1, value2 - epsilon);
+}
+
 void assertTensorEquals(const TensorCPU& tensor1, const TensorCPU& tensor2) {
   CAFFE_ENFORCE_EQ(tensor1.sizes(), tensor2.sizes());
   if (tensor1.IsType<float>()) {
@@ -49,6 +56,7 @@ void assertTensorListEquals(
 const caffe2::Tensor& getTensor(
     const caffe2::Workspace& workspace,
     const std::string& name) {
+  CAFFE_ENFORCE(workspace.HasBlob(name));
   return workspace.GetBlob(name)->Get<caffe2::Tensor>();
 }
 
diff --git a/caffe2/core/test_utils.h b/caffe2/core/test_utils.h
index 3b3af64..6754684 100644
--- a/caffe2/core/test_utils.h
+++ b/caffe2/core/test_utils.h
@@ -4,6 +4,9 @@
 #include "caffe2/core/tensor.h"
 #include "caffe2/core/workspace.h"
 
+#include <cmath>
+#include <vector>
+
 // Utilities that make it easier to write caffe2 C++ unit tests.
 // These utils are designed to be concise and easy to use. They may sacrifice
 // performance and should only be used in tests/non production code.
@@ -13,6 +16,37 @@ namespace testing {
 // Asserts that the values of two tensors are the same.
 void assertTensorEquals(const TensorCPU& tensor1, const TensorCPU& tensor2);
 
+// Asserts that two float values are close within epsilon.
+void assertNear(float value1, float value2, float epsilon);
+
+// Asserts that the numeric values of a tensor is equal to a data vector.
+template <typename T>
+void assertTensorEquals(
+    const TensorCPU& tensor,
+    const std::vector<T>& data,
+    float epsilon = 0.1f) {
+  CAFFE_ENFORCE(tensor.IsType<T>());
+  CAFFE_ENFORCE_EQ(tensor.numel(), data.size());
+  for (auto idx = 0; idx < tensor.numel(); ++idx) {
+    if (tensor.IsType<float>()) {
+      assertNear(tensor.data<T>()[idx], data[idx], epsilon);
+    } else {
+      CAFFE_ENFORCE_EQ(tensor.data<T>()[idx], data[idx]);
+    }
+  }
+}
+
+// Assertion for tensor sizes and values.
+template <typename T>
+void assertTensor(
+    const TensorCPU& tensor,
+    const std::vector<int64_t>& sizes,
+    const std::vector<T>& data,
+    float epsilon = 0.1f) {
+  CAFFE_ENFORCE_EQ(tensor.sizes(), sizes);
+  assertTensorEquals(tensor, data, epsilon);
+}
+
 // Asserts a list of tensors presented in two workspaces are equal.
 void assertTensorListEquals(
     const std::vector<std::string>& tensorNames,
@@ -121,7 +155,7 @@ class WorkspaceMutator {
   template <typename T>
   WorkspaceMutator& newTensorConst(
       const string& name,
-      const vector<int64_t>& shape,
+      const std::vector<int64_t>& shape,
       const T& data) {
     createTensorAndConstantFill<T>(name, shape, data, workspace_);
     return *this;
-- 
2.7.4