log cache_update_index

Author: DavidLandup0

keras_hub/src/models/smollm3/smollm3_backbone.py

@@ -124,10 +124,9 @@ def __init__(
         )
 
         hidden_states = self.token_embedding(token_id_input)
-        
+        position_embeddings = self.rotary_embedding(hidden_states, start_index=start_index)
 
         for decoder_layer in self.transformer_layers[:num_layers]:
-            position_embeddings = self.rotary_embedding(hidden_states, start_index=start_index)
             hidden_states = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,

keras_hub/src/models/smollm3/smollm3_layers.py

@@ -601,6 +601,7 @@ def call(
         seq_len = ops.shape(x)[1]
         positions = ops.arange(seq_len, dtype="float32")
         positions = positions + ops.cast(start_index, dtype="float32")
+        print(start_index)
 
         inv_freq_expanded = ops.broadcast_to(
             inv_freq_expanded, (batch_size, ops.shape(self.inv_freq)[0], 1)