From 5f3797f6e1f96a421128573057bdf4225a17816a Mon Sep 17 00:00:00 2001
From: Biruk Mammo <birukw@google.com>
Date: Mon, 16 Jun 2025 10:18:35 -0700
Subject: [PATCH] Allow creating empty `AttentionActivations` for experimental
 code.

PiperOrigin-RevId: 772077675
---
 gemma/activations.h | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/gemma/activations.h b/gemma/activations.h
index 6f14900..60986fe 100644
--- a/gemma/activations.h
+++ b/gemma/activations.h
@@ -98,7 +98,10 @@ struct AttentionActivations {
 
         div_seq_len(static_cast<uint32_t>(seq_len)),
         query_scale(ChooseQueryScale(config)) {
-    HWY_ASSERT(batch_size != 0);
+    if (batch_size == 0) {
+      HWY_WARN("Creating mostly empty activations with a batch_size of 0.");
+      return;
+    }
 
     // For MatMul outputs, precompute their row pointers.
     // If we forget any MatMul outputs here, debug builds print a warning but