diff --git a/CMakeLists.txt b/CMakeLists.txt
index 2017dcd..dcc9087 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -46,26 +46,26 @@ set(SOURCES
   compression/sfp.h
   compression/sfp-inl.h
   compression/test_util.h
+  backprop/backward.cc
+  backprop/backward.h
+  backprop/backward-inl.h
+  backprop/backward_scalar.h
+  backprop/common_scalar.cc
+  backprop/common_scalar.h
+  backprop/forward.cc
+  backprop/forward.h
+  backprop/forward-inl.h
+  backprop/forward_scalar.h
+  backprop/optimizer.cc
+  backprop/optimizer.h
   gemma/configs.h
   gemma/activations.cc
   gemma/activations.h
-  gemma/backward.cc
-  gemma/backward.h
-  gemma/backward-inl.h
-  gemma/backward_scalar.h
   gemma/common.h
   gemma/common-inl.h
-  gemma/common_scalar.cc
-  gemma/common_scalar.h
-  gemma/forward.cc
-  gemma/forward.h
-  gemma/forward-inl.h
-  gemma/forward_scalar.h
   gemma/gemma.cc
   gemma/gemma.h
   gemma/ops.h
-  gemma/optimizer.cc
-  gemma/optimizer.h
   gemma/weights.cc
   gemma/weights.h
   util/app.h
@@ -122,11 +122,11 @@ enable_testing()
 include(GoogleTest)
 
 set(GEMMA_TEST_FILES
+  backprop/backward_test.cc
+  backprop/backward_scalar_test.cc
+  backprop/optimize_test.cc
   gemma/ops_test.cc
   gemma/gemma_test.cc
-  gemma/backward_test.cc
-  gemma/backward_scalar_test.cc
-  gemma/optimize_test.cc
 )
 
 foreach (TESTFILE IN LISTS GEMMA_TEST_FILES)
diff --git a/gemma/backward-inl.h b/backprop/backward-inl.h
similarity index 99%
rename from gemma/backward-inl.h
rename to backprop/backward-inl.h
index cbeb946..6d11fea 100644
--- a/gemma/backward-inl.h
+++ b/backprop/backward-inl.h
@@ -26,8 +26,8 @@
 #include <array>
 #include <cmath>
 
+#include "backprop/prompt.h"
 #include "gemma/activations.h"
-#include "gemma/prompt.h"
 #include "gemma/weights.h"
 
 #include "hwy/base.h"
diff --git a/gemma/backward.cc b/backprop/backward.cc
similarity index 95%
rename from gemma/backward.cc
rename to backprop/backward.cc
index eb5b511..d6987a4 100644
--- a/gemma/backward.cc
+++ b/backprop/backward.cc
@@ -13,17 +13,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "gemma/backward.h"
+#include "backprop/backward.h"
 
 // Compiles this file for multiple architectures via "foreach_target.h", to
 // which we pass the filename via macro 'argument'.
 #undef HWY_TARGET_INCLUDE
-#define HWY_TARGET_INCLUDE "gemma/backward.cc"  // NOLINT
+#define HWY_TARGET_INCLUDE "backprop/backward.cc"  // NOLINT
 #include "hwy/foreach_target.h"        // IWYU pragma: keep
 
-#include "gemma/backward-inl.h"
-#include "gemma/weights.h"
 #include "hwy/highway.h"
+// After highway.h
+#include "backprop/backward-inl.h"
+#include "gemma/weights.h"
 
 HWY_BEFORE_NAMESPACE();
 namespace gcpp {
diff --git a/gemma/backward.h b/backprop/backward.h
similarity index 97%
rename from gemma/backward.h
rename to backprop/backward.h
index 72281fb..6917f20 100644
--- a/gemma/backward.h
+++ b/backprop/backward.h
@@ -18,8 +18,8 @@
 
 #include <vector>
 
+#include "backprop/prompt.h"
 #include "gemma/common.h"
-#include "gemma/prompt.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
 
 namespace gcpp {
diff --git a/gemma/backward_scalar.h b/backprop/backward_scalar.h
similarity index 99%
rename from gemma/backward_scalar.h
rename to backprop/backward_scalar.h
index 5e5a93b..1c9fbbc 100644
--- a/gemma/backward_scalar.h
+++ b/backprop/backward_scalar.h
@@ -23,9 +23,9 @@
 #include <complex>
 #include <vector>
 
+#include "backprop/common_scalar.h"
+#include "backprop/prompt.h"
 #include "gemma/activations.h"
-#include "gemma/common_scalar.h"
-#include "gemma/prompt.h"
 #include "gemma/weights.h"
 
 namespace gcpp {
diff --git a/gemma/backward_scalar_test.cc b/backprop/backward_scalar_test.cc
similarity index 99%
rename from gemma/backward_scalar_test.cc
rename to backprop/backward_scalar_test.cc
index 713665b..17f1b6c 100644
--- a/gemma/backward_scalar_test.cc
+++ b/backprop/backward_scalar_test.cc
@@ -13,15 +13,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "gemma/backward_scalar.h"
+#include "backprop/backward_scalar.h"
 
 #include <array>
 #include <complex>
 #include <random>
 
-#include "gemma/forward_scalar.h"
-#include "gemma/sampler.h"
-#include "gemma/test_util.h"
+#include "backprop/forward_scalar.h"
+#include "backprop/sampler.h"
+#include "backprop/test_util.h"
 #include "gtest/gtest.h"
 
 namespace gcpp {
diff --git a/gemma/backward_test.cc b/backprop/backward_test.cc
similarity index 97%
rename from gemma/backward_test.cc
rename to backprop/backward_test.cc
index cdeac2e..cf13e6e 100644
--- a/gemma/backward_test.cc
+++ b/backprop/backward_test.cc
@@ -25,27 +25,27 @@
 #include <random>
 #include <vector>
 
+#include "backprop/backward_scalar.h"
+#include "backprop/forward_scalar.h"
+#include "backprop/sampler.h"
+#include "backprop/test_util.h"
 #include "compression/compress.h"
 #include "hwy/aligned_allocator.h"
 #include "hwy/base.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
-#include "gemma/backward_scalar.h"
-#include "gemma/forward_scalar.h"
 #include "gemma/gemma.h"
-#include "gemma/sampler.h"
-#include "gemma/test_util.h"
 #include "gemma/weights.h"
 
 // clang-format off
 #undef HWY_TARGET_INCLUDE
-#define HWY_TARGET_INCLUDE "gemma/backward_test.cc"  //NOLINT
+#define HWY_TARGET_INCLUDE "backprop/backward_test.cc"  //NOLINT
 // clang-format on
 #include "hwy/foreach_target.h"  // IWYU pragma: keep
 #include "hwy/highway.h"
 #include "hwy/tests/test_util-inl.h"
 // After highway.h
-#include "gemma/backward-inl.h"
-#include "gemma/forward-inl.h"
+#include "backprop/backward-inl.h"
+#include "backprop/forward-inl.h"
 #include "gemma/ops.h"
 
 HWY_BEFORE_NAMESPACE();
diff --git a/gemma/common_scalar.cc b/backprop/common_scalar.cc
similarity index 93%
rename from gemma/common_scalar.cc
rename to backprop/common_scalar.cc
index bbf201c..b04dffb 100644
--- a/gemma/common_scalar.cc
+++ b/backprop/common_scalar.cc
@@ -16,10 +16,11 @@
 // Compiles this file for multiple architectures via "foreach_target.h", to
 // which we pass the filename via macro 'argument'.
 #undef HWY_TARGET_INCLUDE
-#define HWY_TARGET_INCLUDE "gemma/common_scalar.cc"  // NOLINT
+#define HWY_TARGET_INCLUDE "backprop/common_scalar.cc"  // NOLINT
 #include "hwy/foreach_target.h"        // IWYU pragma: keep
 
-#include "hwy/highway.h"  // IWYU pragma: keep
+#include "hwy/highway.h"
+// After highway.h
 #include "gemma/ops.h"
 
 HWY_BEFORE_NAMESPACE();
diff --git a/gemma/common_scalar.h b/backprop/common_scalar.h
similarity index 100%
rename from gemma/common_scalar.h
rename to backprop/common_scalar.h
diff --git a/gemma/forward-inl.h b/backprop/forward-inl.h
similarity index 100%
rename from gemma/forward-inl.h
rename to backprop/forward-inl.h
diff --git a/gemma/forward.cc b/backprop/forward.cc
similarity index 94%
rename from gemma/forward.cc
rename to backprop/forward.cc
index 9699a34..fb67c2a 100644
--- a/gemma/forward.cc
+++ b/backprop/forward.cc
@@ -13,16 +13,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "gemma/forward.h"
+#include "backprop/forward.h"
 
 // Compiles this file for multiple architectures via "foreach_target.h", to
 // which we pass the filename via macro 'argument'.
 #undef HWY_TARGET_INCLUDE
-#define HWY_TARGET_INCLUDE "gemma/forward.cc"  // NOLINT
+#define HWY_TARGET_INCLUDE "backprop/forward.cc"  // NOLINT
 #include "hwy/foreach_target.h"        // IWYU pragma: keep
 
-#include "hwy/highway.h"  // IWYU pragma: keep
-#include "gemma/forward-inl.h"
+#include "hwy/highway.h"
+// After highway.h
+#include "backprop/forward-inl.h"
 #include "gemma/weights.h"
 
 HWY_BEFORE_NAMESPACE();
diff --git a/gemma/forward.h b/backprop/forward.h
similarity index 97%
rename from gemma/forward.h
rename to backprop/forward.h
index d8769f2..f17c898 100644
--- a/gemma/forward.h
+++ b/backprop/forward.h
@@ -18,8 +18,8 @@
 
 #include <vector>
 
+#include "backprop/prompt.h"
 #include "gemma/common.h"
-#include "gemma/prompt.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
 
 namespace gcpp {
diff --git a/gemma/forward_scalar.h b/backprop/forward_scalar.h
similarity index 99%
rename from gemma/forward_scalar.h
rename to backprop/forward_scalar.h
index 59c05b6..5643530 100644
--- a/gemma/forward_scalar.h
+++ b/backprop/forward_scalar.h
@@ -23,9 +23,9 @@
 #include <complex>
 #include <vector>
 
+#include "backprop/common_scalar.h"
+#include "backprop/prompt.h"
 #include "gemma/activations.h"
-#include "gemma/common_scalar.h"
-#include "gemma/prompt.h"
 #include "gemma/weights.h"
 
 namespace gcpp {
diff --git a/gemma/optimize_test.cc b/backprop/optimize_test.cc
similarity index 97%
rename from gemma/optimize_test.cc
rename to backprop/optimize_test.cc
index da9fcbb..5c354a9 100644
--- a/gemma/optimize_test.cc
+++ b/backprop/optimize_test.cc
@@ -16,12 +16,12 @@
 #include <iostream>
 #include <string>
 
+#include "backprop/backward.h"
+#include "backprop/forward.h"
+#include "backprop/optimizer.h"
+#include "backprop/sampler.h"
 #include "gemma/activations.h"
-#include "gemma/backward.h"
-#include "gemma/forward.h"
 #include "gemma/gemma.h"
-#include "gemma/optimizer.h"
-#include "gemma/sampler.h"
 #include "gemma/weights.h"
 #include "gtest/gtest.h"
 
diff --git a/gemma/optimizer.cc b/backprop/optimizer.cc
similarity index 99%
rename from gemma/optimizer.cc
rename to backprop/optimizer.cc
index 561b8ee..5b1c61d 100644
--- a/gemma/optimizer.cc
+++ b/backprop/optimizer.cc
@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "gemma/optimizer.h"
+#include "backprop/optimizer.h"
 
 #include <random>
 
diff --git a/gemma/optimizer.h b/backprop/optimizer.h
similarity index 100%
rename from gemma/optimizer.h
rename to backprop/optimizer.h
diff --git a/gemma/prompt.h b/backprop/prompt.h
similarity index 100%
rename from gemma/prompt.h
rename to backprop/prompt.h
diff --git a/gemma/sampler.h b/backprop/sampler.h
similarity index 98%
rename from gemma/sampler.h
rename to backprop/sampler.h
index 34a5ed3..257b993 100644
--- a/gemma/sampler.h
+++ b/backprop/sampler.h
@@ -18,7 +18,7 @@
 
 #include <vector>
 
-#include "gemma/prompt.h"
+#include "backprop/prompt.h"
 
 namespace gcpp {
 
diff --git a/gemma/test_util.h b/backprop/test_util.h
similarity index 100%
rename from gemma/test_util.h
rename to backprop/test_util.h