diff --git a/compression/python/compression_clif_aux.cc b/compression/python/compression_clif_aux.cc
index 3568ad3..5721499 100644
--- a/compression/python/compression_clif_aux.cc
+++ b/compression/python/compression_clif_aux.cc
@@ -29,19 +29,20 @@
 
 #include "compression/compress.h"  // ScaleWeights
 #include "gemma/configs.h"         // ModelConfig
+#include "gemma/io/blob_store.h"   // BlobWriter
+#include "gemma/io/io.h"           // Path
 #include "gemma/model_store.h"     // ModelStore
 #include "gemma/tensor_info.h"     // TensorInfo
 #include "gemma/tokenizer.h"
-#include "io/blob_store.h"  // BlobWriter
-#include "io/io.h"          // Path
 #include "util/basics.h"
 #include "util/mat.h"
 #include "util/threading_context.h"
 
+
 #undef HWY_TARGET_INCLUDE
 #define HWY_TARGET_INCLUDE \
   "compression/python/compression_clif_aux.cc"  // NOLINT
-#include "hwy/foreach_target.h"  // IWYU pragma: keep
+#include "hwy/foreach_target.h"                 // IWYU pragma: keep
 #include "hwy/highway.h"
 // After highway.h
 #include "compression/compress-inl.h"
diff --git a/compression/python/compression_clif_aux.h b/compression/python/compression_clif_aux.h
index 6979865..13b5818 100644
--- a/compression/python/compression_clif_aux.h
+++ b/compression/python/compression_clif_aux.h
@@ -23,11 +23,12 @@
 
 #include "compression/types.h"  // Type
 #include "gemma/configs.h"
+#include "gemma/io/blob_store.h"
 #include "gemma/model_store.h"
 #include "gemma/tensor_info.h"
-#include "io/blob_store.h"
-#include "util/mat.h"
 #include "hwy/aligned_allocator.h"  // Span
+#include "util/mat.h"
+
 
 namespace gcpp {
 
diff --git a/evals/benchmark.cc b/evals/benchmark.cc
index 69cd644..95b6b4f 100644
--- a/evals/benchmark.cc
+++ b/evals/benchmark.cc
@@ -11,11 +11,12 @@
 #include "evals/benchmark_helper.h"
 #include "evals/cross_entropy.h"
 #include "gemma/gemma.h"
-#include "io/io.h"  // Path
-#include "util/args.h"
+#include "gemma/io/io.h"  // Path
 #include "hwy/base.h"
 #include "hwy/timer.h"
 #include "nlohmann/json.hpp"
+#include "util/args.h"
+
 
 namespace gcpp {
 
@@ -85,8 +86,8 @@ int BenchmarkCrossEntropy(GemmaEnv& env, const Path& text,
     LogSpeedStats(time_start, pos + num_tokens);
     std::string text_slice = env.StringFromTokens(prompt_slice);
     total_input_len += text_slice.size();
-    printf("Total cross entropy: %f [cumulative: %f]\n",
-           entropy, total_entropy);
+    printf("Total cross entropy: %f [cumulative: %f]\n", entropy,
+           total_entropy);
     printf("Cross entropy per byte: %f [cumulative: %f]\n",
            entropy / text_slice.size(), total_entropy / total_input_len);
   }
diff --git a/evals/debug_prompt.cc b/evals/debug_prompt.cc
index a6cf8c4..60be008 100644
--- a/evals/debug_prompt.cc
+++ b/evals/debug_prompt.cc
@@ -20,10 +20,11 @@
 
 #include "evals/benchmark_helper.h"
 #include "gemma/gemma.h"  // LayersOutputFunc
-#include "io/io.h"
-#include "util/args.h"
+#include "gemma/io/io.h"
 #include "hwy/base.h"
 #include "nlohmann/json.hpp"
+#include "util/args.h"
+
 
 using json = nlohmann::json;
 
diff --git a/evals/run_mmlu.cc b/evals/run_mmlu.cc
index c6ce972..7476ce3 100644
--- a/evals/run_mmlu.cc
+++ b/evals/run_mmlu.cc
@@ -21,12 +21,13 @@
 
 #include "evals/benchmark_helper.h"
 #include "gemma/gemma.h"  // Gemma
-#include "io/io.h"        // Path
-#include "util/args.h"
+#include "gemma/io/io.h"  // Path
 #include "hwy/base.h"
 #include "hwy/highway.h"
 #include "hwy/profiler.h"
 #include "nlohmann/json.hpp"
+#include "util/args.h"
+
 
 namespace gcpp {
 
diff --git a/gemma/io/blob_compare.cc b/gemma/io/blob_compare.cc
index 9bb860e..85217e4 100644
--- a/gemma/io/blob_compare.cc
+++ b/gemma/io/blob_compare.cc
@@ -21,14 +21,15 @@
 #include <string>
 #include <vector>
 
-#include "io/blob_store.h"
-#include "io/io.h"        // Path
-#include "util/basics.h"  // IndexRange
-#include "util/threading.h"
-#include "util/threading_context.h"
+#include "gemma/io/blob_store.h"
+#include "gemma/io/io.h"            // Path
 #include "hwy/aligned_allocator.h"  // Span
 #include "hwy/base.h"
 #include "hwy/timer.h"
+#include "util/basics.h"  // IndexRange
+#include "util/threading.h"
+#include "util/threading_context.h"
+
 
 namespace gcpp {
 
@@ -106,8 +107,8 @@ void ReadBlobs(BlobReader& reader, const RangeVec& ranges, BlobVec& blobs,
                ThreadingContext& ctx, size_t cluster_idx) {
   HWY_ASSERT(reader.Keys().size() == blobs.size());
   HWY_ASSERT(ranges.size() == blobs.size());
-  ParallelFor(Parallelism::kWithinCluster, blobs.size(), ctx,
-              cluster_idx, Callers::kTest, [&](size_t i, size_t /*thread*/) {
+  ParallelFor(Parallelism::kWithinCluster, blobs.size(), ctx, cluster_idx,
+              Callers::kTest, [&](size_t i, size_t /*thread*/) {
                 HWY_ASSERT(ranges[i].bytes == blobs[i].size());
                 reader.file().Read(ranges[i].offset, ranges[i].bytes,
                                    blobs[i].data());
@@ -189,8 +190,8 @@ void CompareBlobs(const KeyVec& keys, BlobVec& blobs1, BlobVec& blobs2,
   const double t0 = hwy::platform::Now();
   std::atomic<size_t> blobs_equal{};
   std::atomic<size_t> blobs_diff{};
-  ParallelFor(Parallelism::kHierarchical, keys.size(), ctx, 0,
-              Callers::kTest, [&](size_t i, size_t /*thread*/) {
+  ParallelFor(Parallelism::kHierarchical, keys.size(), ctx, 0, Callers::kTest,
+              [&](size_t i, size_t /*thread*/) {
                 const size_t mismatches =
                     BlobDifferences(blobs1[i], blobs2[i], keys[i]);
                 if (mismatches != 0) {
diff --git a/gemma/io/blob_store.cc b/gemma/io/blob_store.cc
index 8346e4b..02b3c3e 100644
--- a/gemma/io/blob_store.cc
+++ b/gemma/io/blob_store.cc
@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "io/blob_store.h"
+#include "gemma/io/blob_store.h"
 
 #include <stddef.h>
 #include <stdint.h>
@@ -24,12 +24,13 @@
 #include <utility>  // std::move
 #include <vector>
 
-#include "io/io.h"
-#include "util/threading_context.h"
+#include "gemma/io/io.h"
 #include "hwy/aligned_allocator.h"  // Span
 #include "hwy/base.h"
 #include "hwy/detect_compiler_arch.h"
 #include "hwy/profiler.h"
+#include "util/threading_context.h"
+
 
 namespace gcpp {
 
@@ -490,25 +491,24 @@ void BlobWriter::Add(const std::string& key, const void* data, size_t bytes) {
 
   const Parallelism parallelism =
       file_->IsAppendOnly() ? Parallelism::kNone : Parallelism::kFlat;
-  ParallelFor(
-      parallelism, writes.size(), ctx_,
-      /*cluster_idx=*/0, Callers::kBlobWriter,
-      [this, &writes](uint64_t i, size_t /*thread*/) {
-        const BlobRange& range = writes[i].range;
-        if (!file_->Write(writes[i].data, range.bytes, range.offset)) {
-          const std::string& key = StringFromKey(keys_[range.key_idx]);
-          HWY_ABORT("Write failed for %s from %zu, %zu bytes to %p.",
-                    key.c_str(), static_cast<size_t>(range.offset), range.bytes,
-                    writes[i].data);
-        }
-      });
+  ParallelFor(parallelism, writes.size(), ctx_,
+              /*cluster_idx=*/0, Callers::kBlobWriter,
+              [this, &writes](uint64_t i, size_t /*thread*/) {
+                const BlobRange& range = writes[i].range;
+                if (!file_->Write(writes[i].data, range.bytes, range.offset)) {
+                  const std::string& key = StringFromKey(keys_[range.key_idx]);
+                  HWY_ABORT("Write failed for %s from %zu, %zu bytes to %p.",
+                            key.c_str(), static_cast<size_t>(range.offset),
+                            range.bytes, writes[i].data);
+                }
+              });
   curr_offset_ = writes.back().range.End();
 }
 
 void BlobWriter::Finalize() {
   if (!file_->IsAppendOnly() && curr_offset_ != file_->FileSize()) {
-    HWY_WARN("Computed offset %zu does not match file size %zu.",
-             curr_offset_, file_->FileSize());
+    HWY_WARN("Computed offset %zu does not match file size %zu.", curr_offset_,
+             file_->FileSize());
   }
   const BlobStore bs = BlobStore(keys_, blob_sizes_);
 
diff --git a/gemma/io/blob_store.h b/gemma/io/blob_store.h
index 82c2357..e5461d3 100644
--- a/gemma/io/blob_store.h
+++ b/gemma/io/blob_store.h
@@ -26,11 +26,12 @@
 #include <unordered_map>
 #include <vector>
 
-#include "io/io.h"                // File, Path, MapPtr
-#include "util/basics.h"          // Tristate
-#include "util/threading_context.h"
+#include "gemma/io/io.h"            // File, Path, MapPtr
 #include "hwy/aligned_allocator.h"  // Span
 #include "hwy/base.h"               // HWY_ASSERT
+#include "util/basics.h"            // Tristate
+#include "util/threading_context.h"
+
 
 namespace gcpp {
 
diff --git a/gemma/io/blob_store_test.cc b/gemma/io/blob_store_test.cc
index cf96684..e1f2995 100644
--- a/gemma/io/blob_store_test.cc
+++ b/gemma/io/blob_store_test.cc
@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "io/blob_store.h"
+#include "gemma/io/blob_store.h"
 
 #include <stdio.h>
 
@@ -22,11 +22,12 @@
 #include <string>
 #include <vector>
 
-#include "io/io.h"
-#include "util/basics.h"
-#include "util/threading_context.h"
+#include "gemma/io/io.h"
 #include "hwy/tests/hwy_gtest.h"
 #include "hwy/tests/test_util-inl.h"  // HWY_ASSERT_EQ
+#include "util/basics.h"
+#include "util/threading_context.h"
+
 
 namespace gcpp {
 namespace {
@@ -130,8 +131,8 @@ TEST(BlobStoreTest, TestNumBlobs) {
     HWY_ASSERT_EQ(reader.Keys().size(), num_blobs);
 
     ParallelFor(
-        Parallelism::kFlat, num_blobs, ctx, /*cluster_idx=*/0,
-        Callers::kTest, [&](uint64_t i, size_t /*thread*/) {
+        Parallelism::kFlat, num_blobs, ctx, /*cluster_idx=*/0, Callers::kTest,
+        [&](uint64_t i, size_t /*thread*/) {
           HWY_ASSERT_STRING_EQ(reader.Keys()[i].c_str(),
                                std::to_string(i).c_str());
           const BlobRange* range = reader.Find(keys[i]);
diff --git a/gemma/io/fields.cc b/gemma/io/fields.cc
index 4516e89..dce67e6 100644
--- a/gemma/io/fields.cc
+++ b/gemma/io/fields.cc
@@ -13,8 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "io/fields.h"
-
+#include "gemma/io/fields.h"
 #include <stdarg.h>
 #include <stddef.h>
 #include <stdint.h>
diff --git a/gemma/io/fields_test.cc b/gemma/io/fields_test.cc
index f720c15..fe4645c 100644
--- a/gemma/io/fields_test.cc
+++ b/gemma/io/fields_test.cc
@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "io/fields.h"
+#include "gemma/io/fields.h"
 
 #include <stddef.h>
 #include <stdint.h>
@@ -21,11 +21,12 @@
 
 #include <limits>
 #include <string>
-#include <vector>
 #include <type_traits>
+#include <vector>
 
 #include "hwy/tests/hwy_gtest.h"
 
+
 namespace gcpp {
 namespace {
 
diff --git a/gemma/io/io.cc b/gemma/io/io.cc
index 2f479b2..e39671e 100644
--- a/gemma/io/io.cc
+++ b/gemma/io/io.cc
@@ -31,7 +31,7 @@
 #include <memory>
 #include <string>
 
-#include "io/io.h"
+#include "gemma/io/io.h"
 #include "hwy/base.h"  // HWY_ASSERT
 
 #if (HWY_OS_LINUX || HWY_OS_FREEBSD) && \
diff --git a/gemma/io/io_win.cc b/gemma/io/io_win.cc
index 34773d3..7fa287c 100644
--- a/gemma/io/io_win.cc
+++ b/gemma/io/io_win.cc
@@ -21,9 +21,10 @@
 #include <stddef.h>
 #include <stdint.h>
 
-#include "io/io.h"
-#include "util/allocator.h"
+#include "gemma/io/io.h"
 #include "hwy/base.h"  // HWY_ASSERT
+#include "util/allocator.h"
+
 #ifndef WIN32_LEAN_AND_MEAN
 #define WIN32_LEAN_AND_MEAN
 #endif
diff --git a/paligemma/image.cc b/paligemma/image.cc
index d8b0cfc..aa68e2a 100644
--- a/paligemma/image.cc
+++ b/paligemma/image.cc
@@ -30,7 +30,7 @@
 #include <utility>
 #include <vector>
 
-#include "io/io.h"
+#include "gemma/io/io.h"
 #include "hwy/aligned_allocator.h"  // hwy::Span
 #include "hwy/base.h"
 #include "hwy/profiler.h"
diff --git a/util/args.h b/util/args.h
index 8c6423b..d7eb7dc 100644
--- a/util/args.h
+++ b/util/args.h
@@ -24,9 +24,10 @@
 #include <string>
 #include <vector>
 
-#include "io/io.h"        // Path
+#include "gemma/io/io.h"  // Path
+#include "hwy/base.h"     // HWY_ABORT
 #include "util/basics.h"  // Tristate
-#include "hwy/base.h"  // HWY_ABORT
+
 
 namespace gcpp {
 
diff --git a/util/mat.h b/util/mat.h
index 0830046..b68ef16 100644
--- a/util/mat.h
+++ b/util/mat.h
@@ -24,10 +24,11 @@
 
 // IWYU pragma: begin_exports
 #include "compression/types.h"  // Type
+#include "gemma/io/fields.h"
 #include "gemma/tensor_info.h"
-#include "io/fields.h"
 #include "util/allocator.h"  // AlignedPtr
-#include "util/basics.h"  // Extents2D
+#include "util/basics.h"     // Extents2D
+
 // IWYU pragma: end_exports
 #include "hwy/base.h"
 
@@ -457,7 +458,7 @@ decltype(auto) CallUpcastedActivation(const MatPtr* base, const Func& func,
 // Like CallUpcasted, but only for kv_cache types: kBF16 and kF32.
 template <class Func, typename... Args>
 decltype(auto) CallUpcastedKV(const MatPtr* base, const Func& func,
-                                      Args&&... args) {
+                              Args&&... args) {
   if (base->GetType() == Type::kF32) {
     const MatPtrT<float> mat(*base);
     return func(&mat, std::forward<Args>(args)...);
diff --git a/util/threading_context.h b/util/threading_context.h
index 7e595ba..1f8452d 100644
--- a/util/threading_context.h
+++ b/util/threading_context.h
@@ -23,14 +23,15 @@
 #include <stdint.h>
 
 // IWYU pragma: begin_exports
-#include "io/io.h"  // Path
+#include "gemma/io/io.h"  // Path
+#include "hwy/profiler.h"
 #include "util/allocator.h"
 #include "util/args.h"
 #include "util/basics.h"  // Tristate
 #include "util/threading.h"
 #include "util/topology.h"
 #include "util/zones.h"
-#include "hwy/profiler.h"
+
 // IWYU pragma: end_exports
 
 namespace gcpp {