From 3866cbfbe8622f41b4f9fa17227aaa7a8de13890 Mon Sep 17 00:00:00 2001
From: Rasmus Munk Larsen <rmlarsen@google.com>
Date: Tue, 25 Mar 2025 23:01:46 +0000
Subject: [PATCH] Fix test for TensorRef of trace.

---
 unsupported/test/cxx11_tensor_ref.cpp | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/unsupported/test/cxx11_tensor_ref.cpp b/unsupported/test/cxx11_tensor_ref.cpp
index 5e796b47b..8d6821ac5 100644
--- a/unsupported/test/cxx11_tensor_ref.cpp
+++ b/unsupported/test/cxx11_tensor_ref.cpp
@@ -41,9 +41,6 @@ static void test_simple_lvalue_ref() {
   for (int i = 0; i < 6; ++i) {
     VERIFY_IS_EQUAL(input(i), -i * 2);
   }
-
-  TensorRef<const Tensor<int, 1>> ref5(input.trace());
-  VERIFY_IS_EQUAL(ref5[0], input[0]);
 }
 
 static void test_simple_rvalue_ref() {
@@ -111,6 +108,17 @@ static void test_slice() {
   VERIFY_IS_EQUAL(slice.data(), tensor.data());
 }
 
+static void test_ref_of_trace() {
+  Tensor<int, 2> input(6, 6);
+  input.setRandom();
+  int trace = 0;
+  for (int i = 0; i < 6; ++i) {
+    trace += input(i, i);
+  }
+  TensorRef<const Tensor<int, 0>> ref(input.trace());
+  VERIFY_IS_EQUAL(ref.coeff(0), trace);
+}
+
 static void test_ref_of_ref() {
   Tensor<float, 3> input(3, 5, 7);
   input.setRandom();
@@ -227,6 +235,7 @@ EIGEN_DECLARE_TEST(cxx11_tensor_ref) {
   CALL_SUBTEST(test_simple_rvalue_ref());
   CALL_SUBTEST(test_multiple_dims());
   CALL_SUBTEST(test_slice());
+  CALL_SUBTEST(test_ref_of_trace());
   CALL_SUBTEST(test_ref_of_ref());
   CALL_SUBTEST(test_ref_in_expr());
   CALL_SUBTEST(test_coeff_ref());