diff --git a/tensorflow_quantum/core/ops/noise/BUILD b/tensorflow_quantum/core/ops/noise/BUILD
index 8548b791f..326f34f99 100644
--- a/tensorflow_quantum/core/ops/noise/BUILD
+++ b/tensorflow_quantum/core/ops/noise/BUILD
@@ -60,6 +60,7 @@ cc_binary(
         "//tensorflow_quantum/core/ops:parse_context",
         "//tensorflow_quantum/core/ops:tfq_simulate_utils",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
+        "//tensorflow_quantum/core/src:util_balance_trajectory",
         "//tensorflow_quantum/core/src:util_qsim",
         "@qsim//lib:qsim_lib",
         # tensorflow core framework
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
index 5cbf0b50d..dabe6ceac 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
index 89263b56a..66a0e168c 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
index 341c87910..78c633f13 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
@@ -40,6 +40,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
 #include "tensorflow_quantum/core/src/circuit_parser_qsim.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {
diff --git a/tensorflow_quantum/core/src/BUILD b/tensorflow_quantum/core/src/BUILD
index 9b9419b5b..df1e0e0e9 100644
--- a/tensorflow_quantum/core/src/BUILD
+++ b/tensorflow_quantum/core/src/BUILD
@@ -93,12 +93,12 @@ cc_test(
 
 cc_library(
     name = "util_qsim",
-    srcs = [],
     hdrs = ["util_qsim.h"],
     deps = [
         ":circuit_parser_qsim",
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/container:inlined_vector",  # unclear why needed.
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -126,6 +126,25 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "util_balance_trajectory",
+    srcs = ["util_balance_trajectory.cc"],
+    hdrs = ["util_balance_trajectory.h"],
+    deps = [
+    ],
+)
+
+cc_test(
+    name = "util_balance_trajectory_test",
+    size = "small",
+    srcs = ["util_balance_trajectory_test.cc"],
+    linkstatic = 0,
+    deps = [
+        ":util_balance_trajectory",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "program_resolution",
     srcs = ["program_resolution.cc"],
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.cc b/tensorflow_quantum/core/src/util_balance_trajectory.cc
new file mode 100644
index 000000000..6230e747a
--- /dev/null
+++ b/tensorflow_quantum/core/src/util_balance_trajectory.cc
@@ -0,0 +1,109 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
+
+namespace tfq {
+
+// Balance the number of trajectory computations done between
+// threads. num_samples is a 2d vector containing the number of reps
+// requested for each pauli_sum[i,j]. After running thread_offsets
+// contains 0/-1 values that will offset the work for each thread.
+// to make it as close to uniform as possible. **Assumes circuits
+// have roughly equal simulation cost**
+void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
+                       const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets) {
+  std::vector<int> rep_limits(num_samples.size(), -1);
+  std::vector<int> height(num_threads, 0);
+
+  for (int i = 0; i < num_samples.size(); i++) {
+    for (int j = 0; j < num_samples[i].size(); j++) {
+      rep_limits[i] = std::max(rep_limits[i], num_samples[i][j]);
+    }
+  }
+  int prev_max_height = -1;
+  for (int j = 0; j < num_samples.size(); j++) {
+    int run_ceiling = ((rep_limits[j] + num_threads - 1) / num_threads);
+    int num_lo = num_threads * run_ceiling - rep_limits[j];
+    int num_hi = num_threads - num_lo;
+    int cur_max = prev_max_height;
+    for (int i = 0; i < num_threads; i++) {
+      if (height[i] == cur_max && num_lo) {
+        // previously had extra work on this thread and
+        // have remaining low budget to give.
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      } else if (height[i] == cur_max - 1 && num_hi) {
+        // previously had less work on this thread and
+        // remaining high budget to give.
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else if (num_hi) {
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else {
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      }
+      prev_max_height = std::max(height[i], prev_max_height);
+    }
+  }
+}
+
+// Simpler case of TrajectoryBalance where num_samples is fixed
+// across all circuits.
+void BalanceTrajectory(const int& num_samples, const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets) {
+  std::vector<int> height(num_threads, 0);
+
+  int prev_max_height = -1;
+  for (int j = 0; j < (*thread_offsets)[0].size(); j++) {
+    int run_ceiling = ((num_samples + num_threads - 1) / num_threads);
+    int num_lo = num_threads * run_ceiling - num_samples;
+    int num_hi = num_threads - num_lo;
+    int cur_max = prev_max_height;
+    for (int i = 0; i < num_threads; i++) {
+      if (height[i] == cur_max && num_lo) {
+        // previously had extra work on this thread and
+        // have remaining low budget to give.
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      } else if (height[i] == cur_max - 1 && num_hi) {
+        // previously had less work on this thread and
+        // remaining high budget to give.
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else if (num_hi) {
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else {
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      }
+      prev_max_height = std::max(height[i], prev_max_height);
+    }
+  }
+}
+
+}  // namespace tfq
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.h b/tensorflow_quantum/core/src/util_balance_trajectory.h
new file mode 100644
index 000000000..fca6624d7
--- /dev/null
+++ b/tensorflow_quantum/core/src/util_balance_trajectory.h
@@ -0,0 +1,34 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef UTIL_BALANCE_TRAJECTORY_H_
+#define UTIL_BALANCE_TRAJECTORY_H_
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+namespace tfq {
+
+void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
+                       const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets);
+
+void BalanceTrajectory(const int& num_samples, const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets);
+
+}  // namespace tfq
+
+#endif  // UTIL_BALANCE_TRAJECTORY_H_
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory_test.cc b/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
new file mode 100644
index 000000000..1656a9acf
--- /dev/null
+++ b/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
@@ -0,0 +1,144 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
+
+#include "gtest/gtest.h"
+
+namespace tfq {
+namespace {
+
+static void AssertWellBalanced(const std::vector<std::vector<int>>& n_reps,
+                               const int& num_threads,
+                               const std::vector<std::vector<int>>& offsets) {
+  auto max_work = std::vector<int>(n_reps.size(), -1);
+  for (int i = 0; i < n_reps.size(); i++) {
+    for (int j = 0; j < n_reps[0].size(); j++) {
+      max_work[i] = std::max(max_work[i], n_reps[i][j]);
+    }
+  }
+
+  for (int i = 0; i < n_reps.size(); i++) {
+    int sum = 0;
+    int prev_local_work = 0;
+    for (int k = 0; k < num_threads; k++) {
+      int local_work = (max_work[i] + num_threads - 1) / num_threads;
+      local_work += offsets[k][i];
+      sum += local_work;
+      if (k > 0) {
+        EXPECT_LT(abs(local_work - prev_local_work), 2);
+      }
+      prev_local_work = local_work;
+    }
+    EXPECT_EQ(sum, max_work[i]);
+  }
+}
+
+TEST(UtilQsimTest, BalanceTrajectorySimple) {
+  std::vector<std::vector<int>> n_reps = {{1, 3, 5, 10, 15},
+                                          {1, 10, 20, 30, 40},
+                                          {50, 70, 100, 100, 100},
+                                          {100, 200, 200, 200, 200}};
+  const int num_threads = 3;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectoryPreventIdle) {
+  std::vector<std::vector<int>> n_reps = {{1, 1, 1, 1, 11},
+                                          {1, 1, 1, 11, 1},
+                                          {1, 1, 11, 1, 1},
+                                          {1, 11, 1, 1, 1},
+                                          {11, 1, 1, 1, 1}};
+  const int num_threads = 10;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectoryLowRep) {
+  std::vector<std::vector<int>> n_reps = {
+      {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1},
+      {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}};
+  const int num_threads = 5;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectoryFewHigh) {
+  std::vector<std::vector<int>> n_reps = {
+      {1, 100, 1, 1, 1}, {1, 1, 1, 1, 1000}, {1, 1, 1, 1, 1},   {1, 1, 1, 1, 1},
+      {1, 1, 1, 1, 1},   {1, 10, 1, 1, 1},   {1, 1, 1, 1, 1000}};
+  const int num_threads = 5;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectory1D) {
+  const int n_reps = 100;
+  const int num_threads = 5;
+  // [num_threads, batch_size]
+  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0}};
+
+  std::vector<std::vector<int>> tmp(offsets[0].size(),
+                                    std::vector<int>(2, n_reps));
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(tmp, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectory1D_2) {
+  const int n_reps = 11;
+  const int num_threads = 10;
+  // [num_threads, batch_size]
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
+
+  std::vector<std::vector<int>> tmp(offsets[0].size(),
+                                    std::vector<int>(2, n_reps));
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(tmp, num_threads, offsets);
+}
+
+}  // namespace
+}  // namespace tfq
diff --git a/tensorflow_quantum/core/src/util_qsim.h b/tensorflow_quantum/core/src/util_qsim.h
index f08715343..fc2461acf 100644
--- a/tensorflow_quantum/core/src/util_qsim.h
+++ b/tensorflow_quantum/core/src/util_qsim.h
@@ -441,95 +441,6 @@ tensorflow::Status AccumulateFusedCircuits(
   return status;
 }
 
-// Balance the number of trajectory computations done between
-// threads. num_samples is a 2d vector containing the number of reps
-// requested for each pauli_sum[i,j]. After running thread_offsets
-// contains 0/-1 values that will offset the work for each thread.
-// to make it as close to uniform as possible. **Assumes circuits
-// have rouhgly equal simulation cost**
-static void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
-                              const int& num_threads,
-                              std::vector<std::vector<int>>* thread_offsets) {
-  std::vector<int> rep_limits(num_samples.size(), -1);
-  std::vector<int> height(num_threads, 0);
-
-  for (int i = 0; i < num_samples.size(); i++) {
-    for (int j = 0; j < num_samples[i].size(); j++) {
-      rep_limits[i] = std::max(rep_limits[i], num_samples[i][j]);
-    }
-  }
-  int prev_max_height = -1;
-  for (int j = 0; j < num_samples.size(); j++) {
-    int run_ceiling = ((rep_limits[j] + num_threads - 1) / num_threads);
-    int num_lo = num_threads * run_ceiling - rep_limits[j];
-    int num_hi = num_threads - num_lo;
-    int cur_max = prev_max_height;
-    for (int i = 0; i < num_threads; i++) {
-      if (height[i] == cur_max && num_lo) {
-        // previously had extra work on this thread and
-        // have remaining low budget to give.
-        height[i]++;
-        (*thread_offsets)[i][j] = -1;
-        num_lo--;
-      } else if (height[i] == cur_max - 1 && num_hi) {
-        // previously had less work on this thread and
-        // remaining high budget to give.
-        height[i] += 2;
-        (*thread_offsets)[i][j] = 0;
-        num_hi--;
-      } else if (num_hi) {
-        height[i] += 2;
-        (*thread_offsets)[i][j] = 0;
-        num_hi--;
-      } else {
-        height[i]++;
-        (*thread_offsets)[i][j] = -1;
-        num_lo--;
-      }
-      prev_max_height = std::max(height[i], prev_max_height);
-    }
-  }
-}
-
-// Simpler case of TrajectoryBalance where num_samples is fixed
-// across all circuits.
-static void BalanceTrajectory(const int& num_samples, const int& num_threads,
-                              std::vector<std::vector<int>>* thread_offsets) {
-  std::vector<int> height(num_threads, 0);
-
-  int prev_max_height = -1;
-  for (int j = 0; j < (*thread_offsets)[0].size(); j++) {
-    int run_ceiling = ((num_samples + num_threads - 1) / num_threads);
-    int num_lo = num_threads * run_ceiling - num_samples;
-    int num_hi = num_threads - num_lo;
-    int cur_max = prev_max_height;
-    for (int i = 0; i < num_threads; i++) {
-      if (height[i] == cur_max && num_lo) {
-        // previously had extra work on this thread and
-        // have remaining low budget to give.
-        height[i]++;
-        (*thread_offsets)[i][j] = -1;
-        num_lo--;
-      } else if (height[i] == cur_max - 1 && num_hi) {
-        // previously had less work on this thread and
-        // remaining high budget to give.
-        height[i] += 2;
-        (*thread_offsets)[i][j] = 0;
-        num_hi--;
-      } else if (num_hi) {
-        height[i] += 2;
-        (*thread_offsets)[i][j] = 0;
-        num_hi--;
-      } else {
-        height[i]++;
-        (*thread_offsets)[i][j] = -1;
-        num_lo--;
-      }
-      prev_max_height = std::max(height[i], prev_max_height);
-    }
-  }
-}
-
 }  // namespace tfq
 
 #endif  // UTIL_QSIM_H_
diff --git a/tensorflow_quantum/core/src/util_qsim_test.cc b/tensorflow_quantum/core/src/util_qsim_test.cc
index 766a9209b..e6d450182 100644
--- a/tensorflow_quantum/core/src/util_qsim_test.cc
+++ b/tensorflow_quantum/core/src/util_qsim_test.cc
@@ -642,125 +642,5 @@ TEST(UtilQsimTest, AccumulateFusedCircuitsEmpty) {
   EXPECT_NEAR(ss.GetAmpl(dest, 3).real(), 0.0, 1e-5);
 }
 
-static void AssertWellBalanced(const std::vector<std::vector<int>>& n_reps,
-                               const int& num_threads,
-                               const std::vector<std::vector<int>>& offsets) {
-  auto max_work = std::vector<int>(n_reps.size(), -1);
-  for (int i = 0; i < n_reps.size(); i++) {
-    for (int j = 0; j < n_reps[0].size(); j++) {
-      max_work[i] = std::max(max_work[i], n_reps[i][j]);
-    }
-  }
-
-  for (int i = 0; i < n_reps.size(); i++) {
-    int sum = 0;
-    int prev_local_work = 0;
-    for (int k = 0; k < num_threads; k++) {
-      int local_work = (max_work[i] + num_threads - 1) / num_threads;
-      local_work += offsets[k][i];
-      sum += local_work;
-      if (k > 0) {
-        EXPECT_LT(abs(local_work - prev_local_work), 2);
-      }
-      prev_local_work = local_work;
-    }
-    EXPECT_EQ(sum, max_work[i]);
-  }
-}
-
-TEST(UtilQsimTest, BalanceTrajectorySimple) {
-  std::vector<std::vector<int>> n_reps = {{1, 3, 5, 10, 15},
-                                          {1, 10, 20, 30, 40},
-                                          {50, 70, 100, 100, 100},
-                                          {100, 200, 200, 200, 200}};
-  const int num_threads = 3;
-  // [num_threads, n_reps.size()]
-  std::vector<std::vector<int>> offsets = {
-      {0, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}};
-
-  BalanceTrajectory(n_reps, num_threads, &offsets);
-  AssertWellBalanced(n_reps, num_threads, offsets);
-}
-
-TEST(UtilQsimTest, BalanceTrajectoryPreventIdle) {
-  std::vector<std::vector<int>> n_reps = {{1, 1, 1, 1, 11},
-                                          {1, 1, 1, 11, 1},
-                                          {1, 1, 11, 1, 1},
-                                          {1, 11, 1, 1, 1},
-                                          {11, 1, 1, 1, 1}};
-  const int num_threads = 10;
-  // [num_threads, n_reps.size()]
-  std::vector<std::vector<int>> offsets = {
-      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
-      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
-      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
-
-  BalanceTrajectory(n_reps, num_threads, &offsets);
-  AssertWellBalanced(n_reps, num_threads, offsets);
-}
-
-TEST(UtilQsimTest, BalanceTrajectoryLowRep) {
-  std::vector<std::vector<int>> n_reps = {
-      {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1},
-      {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}};
-  const int num_threads = 5;
-  // [num_threads, n_reps.size()]
-  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0}};
-
-  BalanceTrajectory(n_reps, num_threads, &offsets);
-  AssertWellBalanced(n_reps, num_threads, offsets);
-}
-
-TEST(UtilQsimTest, BalanceTrajectoryFewHigh) {
-  std::vector<std::vector<int>> n_reps = {
-      {1, 100, 1, 1, 1}, {1, 1, 1, 1, 1000}, {1, 1, 1, 1, 1},   {1, 1, 1, 1, 1},
-      {1, 1, 1, 1, 1},   {1, 10, 1, 1, 1},   {1, 1, 1, 1, 1000}};
-  const int num_threads = 5;
-  // [num_threads, n_reps.size()]
-  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0}};
-
-  BalanceTrajectory(n_reps, num_threads, &offsets);
-  AssertWellBalanced(n_reps, num_threads, offsets);
-}
-
-TEST(UtilQsimTest, BalanceTrajectory1D) {
-  const int n_reps = 100;
-  const int num_threads = 5;
-  // [num_threads, batch_size]
-  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0},
-                                           {0, 0, 0, 0, 0, 0, 0}};
-
-  std::vector<std::vector<int>> tmp(offsets[0].size(),
-                                    std::vector<int>(2, n_reps));
-  BalanceTrajectory(n_reps, num_threads, &offsets);
-  AssertWellBalanced(tmp, num_threads, offsets);
-}
-
-TEST(UtilQsimTest, BalanceTrajectory1D_2) {
-  const int n_reps = 11;
-  const int num_threads = 10;
-  // [num_threads, batch_size]
-  std::vector<std::vector<int>> offsets = {
-      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
-      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
-      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
-
-  std::vector<std::vector<int>> tmp(offsets[0].size(),
-                                    std::vector<int>(2, n_reps));
-  BalanceTrajectory(n_reps, num_threads, &offsets);
-  AssertWellBalanced(tmp, num_threads, offsets);
-}
-
 }  // namespace
 }  // namespace tfq