Fixed minor things and added comments.

BUILD.bazel

@@ -176,4 +176,4 @@ cc_binary(
         "@hwy//:thread_pool",
         "@nlohmann_json//:json",
     ],
-)
+)

debug_prompt.cc

@@ -5,9 +5,7 @@
 
 #include "gemma/gemma.h"
 #include "nlohmann/json.hpp"
-// copybara:import_next_line:gemma_cpp
 #include "util/app.h"
-// copybara:import_next_line:gemma_cpp
 #include "util/args.h"
 
 using json = nlohmann::json;
@@ -67,7 +65,7 @@ class OutputJsonLogger {
   json json_output;
 
   gcpp::LayersOutputT layers_output_log_f =
-      [this](int pos, std::string key, const float* values, size_t values_len) {
+      [this](int pos, const std::string& key, const float* values, size_t values_len) {
         std::vector<float> v{values, values + values_len};
         json_output[std::to_string(pos)][key] = v;
       };
@@ -132,4 +130,4 @@ int main(int argc, char** argv) {
   }
 
   return EXIT_SUCCESS;
-}
+}

gemma/gemma.cc

@@ -1180,7 +1180,7 @@ float ComputeCrossEntropyImpl(GemmaImpl<TConfig>& gemma, size_t max_tokens,
              total_entropy / std::log(2.0) / (pos + 1));
     }
     Transformer(token, pos, weights, activations, kv_cache, pool, inner_pool,
-                nullptr);
+                /*layers_output=*/nullptr);
     MatVec<kVocabSize, kModelDim>(weights.embedder_input_embedding, 0,
                                   activations.x.data(),
                                   activations.logits.data(), pool);
@@ -1609,7 +1609,7 @@ void GenerateGemma(Gemma& gemma, RuntimeConfig runtime_config,
       gemma, runtime_config.max_tokens, runtime_config.max_generated_tokens,
       runtime_config.temperature, prompt, start_pos, kv_cache, pool, inner_pool,
       stream_token, [](int) { return true; }, gen, runtime_config.verbosity,
-      nullptr);
+      /*layers_output=*/nullptr);
 }
 
 void CompressWeights(gcpp::Model model, const Path& weights,

gemma/gemma.h

@@ -38,7 +38,7 @@ using EmbedderInputT = hwy::bfloat16_t;
 // - ponter to the data array
 // - size of the data array
 using LayersOutputT =
-    std::function<void(int, std::string, const float*, size_t)>;
+    std::function<void(int, const std::string&, const float*, size_t)>;
 constexpr size_t kPrefillBatchSize = 16;
 constexpr bool kSystemPrompt = false;