From f9114f22769b3efc643434ce5fb76bc9d329f93b Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Fri, 31 Jan 2025 11:48:04 -0800
Subject: [PATCH] [executorch][flat_tensor] update flat tensor header test

after changes in D68578075

Differential Revision: [D68971340](https://our.internmc.facebook.com/intern/diff/D68971340/)

[ghstack-poisoned]
---
 extension/flat_tensor/serialize/flat_tensor_header.cpp | 3 ++-
 extension/flat_tensor/serialize/flat_tensor_header.h   | 6 ++++++
 extension/flat_tensor/test/flat_tensor_header_test.cpp | 2 +-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/extension/flat_tensor/serialize/flat_tensor_header.cpp b/extension/flat_tensor/serialize/flat_tensor_header.cpp
index fe1db318e97..b329015e4ce 100644
--- a/extension/flat_tensor/serialize/flat_tensor_header.cpp
+++ b/extension/flat_tensor/serialize/flat_tensor_header.cpp
@@ -76,7 +76,8 @@ uint64_t GetUInt64LE(const uint8_t* data) {
   if (size < FlatTensorHeader::kNumHeadBytes) {
     return Error::InvalidArgument;
   }
-  const uint8_t* header = reinterpret_cast<const uint8_t*>(data);
+  const uint8_t* header =
+      reinterpret_cast<const uint8_t*>(data) + kHeaderOffset;
 
   // Check magic bytes.
   if (std::memcmp(
diff --git a/extension/flat_tensor/serialize/flat_tensor_header.h b/extension/flat_tensor/serialize/flat_tensor_header.h
index 909d85e63ab..e722cd749bb 100644
--- a/extension/flat_tensor/serialize/flat_tensor_header.h
+++ b/extension/flat_tensor/serialize/flat_tensor_header.h
@@ -23,6 +23,12 @@ struct FlatTensorHeader {
    */
   static constexpr size_t kNumHeadBytes = 64;
 
+  /**
+   * The offset into the serialized FlatTensor data where the FlatTensor
+   * header should begin.
+   */
+  static constexpr size_t kHeaderOffset = 8;
+
   /**
    * The magic bytes that identify the header. This should be in sync with
    * the magic in executorch/extension/flat_tensor/serialize/serialize.py
diff --git a/extension/flat_tensor/test/flat_tensor_header_test.cpp b/extension/flat_tensor/test/flat_tensor_header_test.cpp
index d8d95686515..4a0341691e6 100644
--- a/extension/flat_tensor/test/flat_tensor_header_test.cpp
+++ b/extension/flat_tensor/test/flat_tensor_header_test.cpp
@@ -68,7 +68,7 @@ std::vector<uint8_t> CreateExampleFlatTensorHeader() {
   memset(ret.data(), 0x55, ret.size());
   // Copy the example header into the right offset.
   memcpy(
-      ret.data(),
+      ret.data() + FlatTensorHeader::kHeaderOffset,
       kExampleHeaderData,
       sizeof(kExampleHeaderData));
   return ret;