From 32ae4eefcd6a912aff19cb2e7a20a7922201e037 Mon Sep 17 00:00:00 2001
From: "The gemma.cpp Authors" <no-reply@google.com>
Date: Tue, 23 Jul 2024 11:44:51 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 655246378
---
 compression/python/compression_clif_aux.cc | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/compression/python/compression_clif_aux.cc b/compression/python/compression_clif_aux.cc
index 6c938d95..334530a8 100644
--- a/compression/python/compression_clif_aux.cc
+++ b/compression/python/compression_clif_aux.cc
@@ -72,7 +72,9 @@ class SbsWriterImpl : public WriterInterface {
   }
 
   void AddScales(const std::vector<float>& scales) override {
-    compressor_.AddScales(scales.data(), scales.size());
+    HWY_ASSERT(scales_.empty());
+    scales_ = scales;
+    compressor_.AddScales(scales_.data(), scales_.size());
   }
 
   void Write(std::string path) override {
@@ -85,6 +87,7 @@ class SbsWriterImpl : public WriterInterface {
   std::vector<std::vector<SfpStream>> sfp_streams_;
   std::vector<std::vector<NuqStream>> nuq_streams_;
   std::vector<std::vector<hwy::bfloat16_t>> bf16_streams_;
+  std::vector<float> scales_;
 };
 
 WriterInterface* NewSbsWriter() { return new SbsWriterImpl(); }