diff --git a/lib/compiler/include/compiler/cost_estimator/op_cost_metrics.struct.toml b/lib/compiler/include/compiler/cost_estimator/op_cost_metrics.struct.toml
index f137935a4d..d2ff3f42e7 100644
--- a/lib/compiler/include/compiler/cost_estimator/op_cost_metrics.struct.toml
+++ b/lib/compiler/include/compiler/cost_estimator/op_cost_metrics.struct.toml
@@ -7,6 +7,7 @@ features = [
 ]
 
 includes = [
+  "utils/nonnegative_int/nonnegative_int.h"
 ]
 
 [[fields]]
@@ -15,4 +16,4 @@ type = "float"
 
 [[fields]]
 name = "memory"
-type = "size_t"
+type = "::FlexFlow::nonnegative_int"
diff --git a/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc b/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc
index f5d5a5ee1b..ac180cd079 100644
--- a/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc
+++ b/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc
@@ -146,13 +146,13 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     auto map1 = std::unordered_map<OpCostEstimateKey, OpCostMetrics>{{
         {map_unmapped_op_cost_estimate_key(k1, mv1),
-         OpCostMetrics{/*runtime=*/1.0, /*memory=*/0}},
+         OpCostMetrics{/*runtime=*/1.0, /*memory=*/nonnegative_int{0}}},
         {map_unmapped_op_cost_estimate_key(k2, mv1),
-         OpCostMetrics{/*runtime=*/2.0, /*memory=*/0}},
+         OpCostMetrics{/*runtime=*/2.0, /*memory=*/nonnegative_int{0}}},
         {map_unmapped_op_cost_estimate_key(k1, mv2),
-         OpCostMetrics{/*runtime=*/1.5, /*memory=*/0}},
+         OpCostMetrics{/*runtime=*/1.5, /*memory=*/nonnegative_int{0}}},
         {map_unmapped_op_cost_estimate_key(k2, mv2),
-         OpCostMetrics{/*runtime=*/2.5, /*memory=*/0}},
+         OpCostMetrics{/*runtime=*/2.5, /*memory=*/nonnegative_int{0}}},
     }};
 
     CostEstimator cost_estimator = make_fake_cost_estimator(
diff --git a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc
index 8761116be2..9706f1c75f 100644
--- a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc
+++ b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc
@@ -146,10 +146,14 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     CostEstimator cost_estimator = make_fake_cost_estimator(
         std::unordered_map<OpCostEstimateKey, OpCostMetrics>{{
-            {map_unmapped_op_cost_estimate_key(k1, mv1), OpCostMetrics{1.0, 2}},
-            {map_unmapped_op_cost_estimate_key(k2, mv1), OpCostMetrics{2.0, 3}},
-            {map_unmapped_op_cost_estimate_key(k1, mv2), OpCostMetrics{1.5, 1}},
-            {map_unmapped_op_cost_estimate_key(k2, mv2), OpCostMetrics{2.5, 2}},
+            {map_unmapped_op_cost_estimate_key(k1, mv1),
+             OpCostMetrics{1.0, nonnegative_int{2}}},
+            {map_unmapped_op_cost_estimate_key(k2, mv1),
+             OpCostMetrics{2.0, nonnegative_int{3}}},
+            {map_unmapped_op_cost_estimate_key(k1, mv2),
+             OpCostMetrics{1.5, nonnegative_int{1}}},
+            {map_unmapped_op_cost_estimate_key(k2, mv2),
+             OpCostMetrics{2.5, nonnegative_int{2}}},
         }},
         std::unordered_map<TensorSetMovement, float>{{
             {TensorSetMovement{{}}, 0.0},
@@ -183,13 +187,13 @@ TEST_SUITE(FF_TEST_SUITE) {
               cache, context, problem_tree, full_machine_spec, constraints);
       MachineMappingWithMemoryResult correct = MachineMappingWithMemoryResult{{
           MachineMappingForSingleLayer{
-              OpCostMetrics{1.0, 2},
+              OpCostMetrics{1.0, nonnegative_int{2}},
               ParallelLayerGuidObliviousMachineMapping{{
                   {binary_tree_root_path(), mv1},
               }},
           },
           MachineMappingForSingleLayer{
-              OpCostMetrics{1.5, 1},
+              OpCostMetrics{1.5, nonnegative_int{1}},
               ParallelLayerGuidObliviousMachineMapping{{
                   {binary_tree_root_path(), mv2},
               }},
@@ -214,7 +218,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           MachineMappingForSingleLayer{
               OpCostMetrics{
                   /*runtime=*/1.0 + 2.0 + 0.1,
-                  /*memory=*/2 + 3,
+                  /*memory=*/nonnegative_int{2 + 3},
               },
               ParallelLayerGuidObliviousMachineMapping{{
                   {
@@ -232,7 +236,7 @@ TEST_SUITE(FF_TEST_SUITE) {
               }},
           },
           MachineMappingForSingleLayer{
-              OpCostMetrics{1.5 + 2.5 + 0.1, 1 + 2},
+              OpCostMetrics{1.5 + 2.5 + 0.1, nonnegative_int{1 + 2}},
               ParallelLayerGuidObliviousMachineMapping{{
                   {
                       BinaryTreePath{{
@@ -266,7 +270,7 @@ TEST_SUITE(FF_TEST_SUITE) {
               cache, context, problem_tree, full_machine_spec, constraints);
       MachineMappingWithMemoryResult correct =
           MachineMappingWithMemoryResult{{MachineMappingForSingleLayer{
-              OpCostMetrics{2.5, 2},
+              OpCostMetrics{2.5, nonnegative_int{2}},
               ParallelLayerGuidObliviousMachineMapping{{
                   {
                       BinaryTreePath{{
diff --git a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.cc b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.cc
index a47d8713e9..ecfb7cfeb3 100644
--- a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.cc
+++ b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.cc
@@ -53,15 +53,15 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     OpCostMetrics cost1 = OpCostMetrics{
         /*runtime=*/2.0,
-        /*memory=*/2,
+        /*memory=*/nonnegative_int{2},
     };
     OpCostMetrics cost2 = OpCostMetrics{
         /*runtime=*/4.0,
-        /*memory=*/1,
+        /*memory=*/nonnegative_int{1},
     };
     OpCostMetrics cost3 = OpCostMetrics{
         /*runtime=*/2.0,
-        /*memory=*/3,
+        /*memory=*/nonnegative_int{3},
     };
 
     MachineMappingForSingleLayer mm1 = MachineMappingForSingleLayer{
@@ -183,7 +183,7 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     OpCostMetrics pre_cost = OpCostMetrics{
         /*runtime=*/2.0,
-        /*memory=*/2,
+        /*memory=*/nonnegative_int{2},
     };
     MachineMappingWithMemoryResult pre = MachineMappingWithMemoryResult{{
         MachineMappingForSingleLayer{
@@ -209,7 +209,7 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     OpCostMetrics post_cost = OpCostMetrics{
         /*runtime=*/4.0,
-        /*memory=*/1,
+        /*memory=*/nonnegative_int{1},
     };
 
     MachineMappingWithMemoryResult post = MachineMappingWithMemoryResult{{
@@ -378,7 +378,7 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     OpCostMetrics lhs_cost = OpCostMetrics{
         /*runtime=*/2.0,
-        /*memory=*/2,
+        /*memory=*/nonnegative_int{2},
     };
     MachineMappingWithMemoryResult lhs = MachineMappingWithMemoryResult{{
         MachineMappingForSingleLayer{
@@ -404,7 +404,7 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     OpCostMetrics rhs_cost = OpCostMetrics{
         /*runtime=*/4.0,
-        /*memory=*/1,
+        /*memory=*/nonnegative_int{1},
     };
     MachineMappingWithMemoryResult rhs = MachineMappingWithMemoryResult{{
         MachineMappingForSingleLayer{
@@ -519,15 +519,15 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     OpCostMetrics cost1 = OpCostMetrics{
         /*runtime=*/2.0,
-        /*memory=*/2,
+        /*memory=*/nonnegative_int{2},
     };
     OpCostMetrics cost2 = OpCostMetrics{
         /*runtime=*/4.0,
-        /*memory=*/1,
+        /*memory=*/nonnegative_int{1},
     };
     OpCostMetrics cost3 = OpCostMetrics{
         /*runtime=*/2.0,
-        /*memory=*/3,
+        /*memory=*/nonnegative_int{3},
     };
 
     MachineMappingForSingleLayer mm1 = MachineMappingForSingleLayer{
diff --git a/lib/utils/include/utils/nonnegative_int/nonnegative_int.h b/lib/utils/include/utils/nonnegative_int/nonnegative_int.h
index 01bee29f63..0749497c56 100644
--- a/lib/utils/include/utils/nonnegative_int/nonnegative_int.h
+++ b/lib/utils/include/utils/nonnegative_int/nonnegative_int.h
@@ -38,6 +38,8 @@ class nonnegative_int {
   friend bool operator!=(int const &lhs, nonnegative_int const &rhs);
   friend bool operator>=(int const &lhs, nonnegative_int const &rhs);
 
+  nonnegative_int operator+(nonnegative_int const &other) const;
+
   friend std::ostream &operator<<(std::ostream &os, nonnegative_int const &n);
 
   friend int format_as(nonnegative_int const &);
diff --git a/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc b/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc
index 2c317dce86..9088cc4bf9 100644
--- a/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc
+++ b/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc
@@ -71,6 +71,10 @@ bool operator>=(int const &lhs, nonnegative_int const &rhs) {
   return lhs >= rhs.value_;
 }
 
+nonnegative_int nonnegative_int::operator+(nonnegative_int const &other) const {
+  return nonnegative_int{this->value_ + other.value_};
+}
+
 std::ostream &operator<<(std::ostream &os, nonnegative_int const &n) {
   os << n.value_;
   return os;
diff --git a/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc b/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc
index a7b813d1f1..73d382d830 100644
--- a/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc
+++ b/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc
@@ -198,6 +198,15 @@ TEST_SUITE(FF_TEST_SUITE) {
     }
   }
 
+  TEST_CASE("nonnegative_int + operation") {
+    nonnegative_int nn_int_1a = nonnegative_int{1};
+    nonnegative_int nn_int_1b = nonnegative_int{1};
+    nonnegative_int nn_int_2 = nonnegative_int{2};
+    SUBCASE("LHS: nonnegative_int, RHS: nonnegative_int") {
+      CHECK(nn_int_1a + nn_int_1b == nn_int_2);
+    }
+  }
+
   TEST_CASE("adl_serializer<nonnegative_int>") {
     SUBCASE("to_json") {
       nonnegative_int input = nonnegative_int{5};