Pass allocated-unallocated tests

lib/kernels/src/legion_dim.cc

@@ -13,10 +13,16 @@
                                       ff_dim.value.unwrap_nonnegative() - 1}};
 }
 
 ff_dim_t legion_dim_from_ff_dim(legion_dim_t legion_dim,
                                 nonnegative_int num_dimensions) {
   return ff_dim_t{nonnegative_int{num_dimensions.unwrap_nonnegative() -
                                   legion_dim.value.unwrap_nonnegative() - 1}};
 }
 
+ff_dim_t ff_dim_from_legion_dim(legion_dim_t legion_dim,
+                                nonnegative_int num_dimensions) {
+  return ff_dim_t{nonnegative_int{num_dimensions.unwrap_nonnegative() -
+                                  legion_dim.value.unwrap_nonnegative() - 1}};
+}
+
 } // namespace FlexFlow

lib/local-execution/include/local-execution/gradient_tensor_source.h

@@ -11,6 +11,8 @@ struct GradientTensorSource {
 
   gradient_tensor_t new_gradient_tensor();
 
+  void reset();
+
 private:
   static size_t next_available_gradient_tensor_id;
 };

lib/local-execution/include/local-execution/optimizer_tensor_source.h

@@ -11,6 +11,8 @@ struct OptimizerTensorSource {
 
   optimizer_tensor_t new_optimizer_tensor();
 
+  void reset();
+
 private:
   static size_t next_available_optimizer_tensor_id;
 };

lib/local-execution/src/gradient_tensor_source.cc

@@ -11,4 +11,8 @@ gradient_tensor_t GradientTensorSource::new_gradient_tensor() {
       GradientTensorSource::next_available_gradient_tensor_id++};
 }
 
+void GradientTensorSource::reset() {
+  GradientTensorSource::next_available_gradient_tensor_id = 0;
+}
+
 } // namespace FlexFlow

lib/local-execution/src/local_tensor_backing.cc

@@ -11,104 +11,104 @@
 
 namespace FlexFlow {
 
 LocalTensorBacking::LocalTensorBacking(
     AllocatedTensors const &allocated_tensors,
     UnallocatedTensors const &unallocated_tensors,
     Allocator const &allocator)
     : tensor_gradient_mapping(allocated_tensors.gradient_mapping),
       tensor_optimizer_mapping(allocated_tensors.optimizer_mapping),
       allocator(allocator) {
 
   // handle already-allocated tensors
   for (std::pair<TensorTypeVariant, GenericTensorAccessorW> const
            &tensor_type_backing : allocated_tensors.tensor_type_backings) {
     lowered_tensor_t lowered_tensor =
         this->insert_tensor(tensor_type_backing.first);
     this->tensor_backings.insert({lowered_tensor, tensor_type_backing.second});
   }
 
   // allocate new tensors
   this->tensor_gradient_mapping.insert(
       unallocated_tensors.gradient_mapping.begin(),
       unallocated_tensors.gradient_mapping.end());
 
   for (std::pair<tensor_guid_t, std::vector<optimizer_tensor_t>> const
            &unallocated_optimizer_tensors :
        unallocated_tensors.optimizer_mapping) {
     if (this->tensor_optimizer_mapping.count(
             unallocated_optimizer_tensors.first)) {
       for (optimizer_tensor_t const &optimizer_tensor :
            unallocated_optimizer_tensors.second) {
         this->tensor_optimizer_mapping[unallocated_optimizer_tensors.first]
             .push_back(optimizer_tensor);
       }
     } else {
       this->tensor_optimizer_mapping.insert({unallocated_optimizer_tensors});
     }
   }
 
   for (std::pair<TensorTypeVariant, TensorShape> const &tensor_type_shape :
        unallocated_tensors.tensor_type_shapes) {
     lowered_tensor_t lowered_tensor =
         this->insert_tensor(tensor_type_shape.first);
     GenericTensorAccessorW tensor_backing =
         this->allocator.allocate_tensor(tensor_type_shape.second);
     this->tensor_backings.insert({lowered_tensor, tensor_backing});
   }
 };
 
 lowered_tensor_t
     LocalTensorBacking::insert_tensor(TensorTypeVariant const &tensor_type) {
   lowered_tensor_t lowered_tensor =
       this->lowered_tensor_source.new_lowered_tensor();
   tensor_type.visit<std::nullopt_t>(overload{
       [&](tensor_guid_t const &tensor_guid) {
         this->tensor_lowering_mapping.insert({tensor_guid, lowered_tensor});
         return std::nullopt;
       },
       [&](gradient_tensor_t const &gradient_tensor) {
         this->gradient_tensor_lowering_mapping.insert(
             {gradient_tensor, lowered_tensor});
         return std::nullopt;
       },
       [&](optimizer_tensor_t const &optimizer_tensor) {
         this->optimizer_tensor_lowering_mapping.insert(
             {optimizer_tensor, lowered_tensor});
         return std::nullopt;
       },
       [&](loss_tensor_t const &loss_tensor) {
         this->loss_tensor_lowering_mapping.insert(
             {loss_tensor, lowered_tensor});
         return std::nullopt;
       },
       [&](auto const &any_tensor) {
         throw mk_runtime_error(
             fmt::format("Unhandled tensor type {}", any_tensor));
       }});
   return lowered_tensor;
 }
 
 GenericTensorAccessorW
     LocalTensorBacking::get_tensor(TensorTypeVariant const &tensor_type) const {
   lowered_tensor_t lowered_tensor =
       tensor_type.visit<lowered_tensor_t>(overload{
           [&](tensor_guid_t const &tensor_guid) {
             return this->tensor_lowering_mapping.at(tensor_guid);
           },
           [&](gradient_tensor_t const &gradient_tensor) {
             return this->gradient_tensor_lowering_mapping.at(gradient_tensor);
           },
           [&](optimizer_tensor_t const &optimizer_tensor) {
             return this->optimizer_tensor_lowering_mapping.at(optimizer_tensor);
           },
           [&](loss_tensor_t const &loss_tensor) {
             return this->loss_tensor_lowering_mapping.at(loss_tensor);
           },
           [&](auto const &any_tensor) {
             throw mk_runtime_error(
                 fmt::format("Unhandled tensor type {}", any_tensor));
           }});
   return this->tensor_backings.at(lowered_tensor);
 }
 
 UnallocatedTensors generate_unallocated_tensors(
@@ -135,7 +135,7 @@
       gradient_tensor_t gradient_tensor =
           gradient_tensor_source.new_gradient_tensor();
       tensor_type_shapes.insert(
-          {TensorTypeVariant{tensor_guid}, tensor_attrs.shape});
+          {TensorTypeVariant{gradient_tensor}, tensor_attrs.shape});
       gradient_mapping.insert({tensor_guid, gradient_tensor});
     }
   }
@@ -168,8 +168,7 @@
        tensor_attrs_mapping) {
     tensor_guid_t tensor_guid = tensor_guid_attrs.first;
     TensorAttrs tensor_attrs = tensor_guid_attrs.second;
-    if (tensor_attrs.create_gradients == CreateGrad::YES &&
-        !allocated_tensors.optimizer_mapping.count(tensor_guid)) {
+    if (tensor_attrs.create_gradients == CreateGrad::YES) {
       std::vector<optimizer_tensor_t> optimizer_tensors;
 
       int num_optimizer_tensors_to_allocate =
@@ -178,6 +177,7 @@
         num_optimizer_tensors_to_allocate -=
             allocated_tensors.optimizer_mapping.at(tensor_guid).size();
       }
+      std::cout << num_optimizer_tensors_to_allocate;
 
       for (int i = 0; i < num_optimizer_tensors_to_allocate; ++i) {
         optimizer_tensor_t optimizer_tensor =
@@ -186,7 +186,10 @@
         tensor_type_shapes.insert(
             {TensorTypeVariant{optimizer_tensor}, tensor_attrs.shape});
       }
-      optimizer_mapping.insert({tensor_guid, optimizer_tensors});
+
+      if (num_optimizer_tensors_to_allocate > 0) {
+        optimizer_mapping.insert({tensor_guid, optimizer_tensors});
+      }
     }
   }
 
@@ -194,18 +197,18 @@
       tensor_type_shapes, gradient_mapping, optimizer_mapping};
 }
 
 TensorSlotsBacking construct_tensor_slots_backing(
     LocalTensorBacking const &local_tensor_backing,
     TaskBinding const &binding) {
   TensorSlotsBacking mapping;
 
   for (std::pair<SlotTensorTypeId, TensorTypeVariant> const &tensor_binding :
        binding.get_tensor_bindings()) {
     mapping.insert({tensor_binding.first,
                     local_tensor_backing.get_tensor(tensor_binding.second)});
   }
 
   return mapping;
 }
 
 } // namespace FlexFlow

lib/local-execution/src/local_training_backing.cc

@@ -19,227 +19,226 @@
     AllocatedTensors const &allocated_tensors,
     ComputationGraph const &computation_graph,
     RuntimeArgConfig const &runtime_arg_config)
     : computation_graph(computation_graph),
       task_registry(construct_task_registry(computation_graph)),
       local_tensor_backing(allocated_tensors,
                            generate_unallocated_tensors(
                                allocated_tensors,
                                get_all_tensor_attrs(this->computation_graph),
                                this->gradient_tensor_source),
                            allocator),
       local_args_backing(initialize_args_backing(this->task_registry,
                                                  this->computation_graph,
                                                  runtime_arg_config,
                                                  this->local_tensor_backing)){};
 
 LocalTrainingBacking::LocalTrainingBacking(
     Allocator const &allocator,
     AllocatedTensors const &allocated_tensors,
     ComputationGraph const &computation_graph,
     RuntimeArgConfig const &runtime_arg_config,
     OptimizerAttrs const &optimizer_attrs)
     : computation_graph(computation_graph),
       task_registry(construct_task_registry(computation_graph)),
       local_tensor_backing(allocated_tensors,
                            generate_unallocated_tensors_with_optimizer(
                                allocated_tensors,
                                get_all_tensor_attrs(this->computation_graph),
                                this->gradient_tensor_source,
                                this->optimizer_tensor_source,
                                optimizer_attrs),
                            allocator),
       local_args_backing(initialize_args_backing(this->task_registry,
                                                  this->computation_graph,
                                                  runtime_arg_config,
                                                  this->local_tensor_backing)){};
 
 LocalArgsBacking
     initialize_args_backing(TaskRegistry const &task_registry,
                             ComputationGraph const &cg,
                             RuntimeArgConfig const &runtime_arg_config,
                             LocalTensorBacking const &local_tensor_backing) {
   std::unordered_map<layer_guid_t, DeviceSpecificDeviceStates>
       per_device_op_states;
   for (layer_guid_t const &node : topological_ordering(cg)) {
     if (registry_contains_task_for_layer(
             task_registry, node, OpTaskType::INIT)) {
       ComputationGraphOpAttrs attrs = get_layer_attrs(cg, node).attrs;
 
       TaskInvocation invocation =
           lower_to_task_invocation(init(attrs),
                                    node,
                                    get_incoming_inputs(cg, node),
                                    get_incoming_input_shapes(cg, node),
                                    get_outgoing_tensors(cg, node),
                                    get_incoming_weights(cg, node),
                                    local_tensor_backing.tensor_gradient_mapping,
                                    std::nullopt);
       TaskArgumentAccessor accessor = get_task_arg_accessor(
           local_tensor_backing,
           make_args_backing_with_empty_device_states(runtime_arg_config),
           invocation);
       TaskSignatureAndImpl task_sig_impl =
           task_registry.task_mapping.at(invocation.task_id);
       auto fn = task_sig_impl.impl_function.get<InitOpTaskImplFunction>()
                     .function_ptr;
       DeviceSpecificDeviceStates device_state = fn(accessor);
       per_device_op_states.insert({node, device_state});
     }
   }
 
   return LocalArgsBacking{runtime_arg_config, per_device_op_states};
 }
 
 std::optional<float> call_task_impl(TaskRegistry const &task_registry,
                                     task_id_t const &task_id,
                                     TaskArgumentAccessor const &acc) {
   TaskSignatureAndImpl task_sig_impl = task_registry.task_mapping.at(task_id);
   auto fn =
       task_sig_impl.impl_function.get<FwdBwdOpTaskImplFunction>().function_ptr;
   return fn(acc);
 }
 
 std::optional<float>
     execute_forward(LocalTrainingBacking const &local_training_backing,
-                    layer_guid_t const &operator_node,
-                    Allocator &allocator) {
+                    layer_guid_t const &operator_node) {
   if (registry_contains_task_for_layer(local_training_backing.task_registry,
                                        operator_node,
                                        OpTaskType::FWD)) {
     ComputationGraphOpAttrs attrs =
         get_layer_attrs(local_training_backing.computation_graph, operator_node)
             .attrs;
 
     std::optional<DeviceSpecificDeviceStates> device_state =
         get_per_device_op_state_if_exists(
             local_training_backing.local_args_backing, operator_node);
     TaskInvocation invocation = lower_to_task_invocation(
         forward(attrs),
         operator_node,
         get_incoming_inputs(local_training_backing.computation_graph,
                             operator_node),
         get_incoming_input_shapes(local_training_backing.computation_graph,
                                   operator_node),
         get_outgoing_tensors(local_training_backing.computation_graph,
                              operator_node),
         get_incoming_weights(local_training_backing.computation_graph,
                              operator_node),
         local_training_backing.local_tensor_backing.tensor_gradient_mapping,
         device_state);
     TaskArgumentAccessor accessor =
         get_task_arg_accessor(local_training_backing.local_tensor_backing,
                               local_training_backing.local_args_backing,
                               invocation);
     return call_task_impl(
         local_training_backing.task_registry, invocation.task_id, accessor);
   } else {
     return std::nullopt;
   }
 }
 
 void compute_loss(LocalTrainingBacking const &local_training_backing,
                   LossAttrs const &loss_attrs,
                   tensor_guid_t const &logit_tensor,
                   loss_tensor_t const &label_tensor) {
   TaskInvocation loss_invocation = backward(
       loss_attrs,
       logit_tensor,
       local_training_backing.local_tensor_backing.tensor_gradient_mapping.at(
           logit_tensor),
       label_tensor);
   // TODO: https://github.com/flexflow/flexflow-train/issues/1442
   // assert(is_invocation_valid(get_loss_bwd_signature(), loss_invocation));
   TaskArgumentAccessor loss_accessor =
       get_task_arg_accessor(local_training_backing.local_tensor_backing,
                             local_training_backing.local_args_backing,
                             loss_invocation);
   TaskImplFunction loss_impl_fn = get_loss_bwd_task_impl();
   loss_impl_fn.get<GenericTaskImplFunction>().function_ptr(loss_accessor);
 }
 
 std::optional<float>
     execute_backward(LocalTrainingBacking const &local_training_backing,
                      layer_guid_t const &operator_node) {
   if (registry_contains_task_for_layer(local_training_backing.task_registry,
                                        operator_node,
                                        OpTaskType::BWD)) {
     ComputationGraphOpAttrs attrs =
         get_layer_attrs(local_training_backing.computation_graph, operator_node)
             .attrs;
 
     std::optional<DeviceSpecificDeviceStates> device_state =
         get_per_device_op_state_if_exists(
             local_training_backing.local_args_backing, operator_node);
     TaskInvocation invocation = lower_to_task_invocation(
         backward(attrs),
         operator_node,
         get_incoming_inputs(local_training_backing.computation_graph,
                             operator_node),
         get_incoming_input_shapes(local_training_backing.computation_graph,
                                   operator_node),
         get_outgoing_tensors(local_training_backing.computation_graph,
                              operator_node),
         get_incoming_weights(local_training_backing.computation_graph,
                              operator_node),
         local_training_backing.local_tensor_backing.tensor_gradient_mapping,
         device_state);
     TaskArgumentAccessor accessor =
         get_task_arg_accessor(local_training_backing.local_tensor_backing,
                               local_training_backing.local_args_backing,
                               invocation);
     return call_task_impl(
         local_training_backing.task_registry, invocation.task_id, accessor);
   } else {
     return std::nullopt;
   }
 }
 
 void execute_update(LocalTrainingBacking const &local_training_backing,
                     layer_guid_t const &node,
                     OptimizerAttrs const &optimizer_attrs) {
   LayerAttrs layer_attrs =
       get_layer_attrs(local_training_backing.computation_graph, node);
   if (layer_attrs.attrs.has<WeightAttrs>()) {
     // get tensors
     tensor_guid_t weight_tensor = get_only(
         get_outgoing_tensors(local_training_backing.computation_graph, node));
 
     gradient_tensor_t weight_grad_tensor =
         local_training_backing.local_tensor_backing.tensor_gradient_mapping.at(
             weight_tensor);
     std::vector<optimizer_tensor_t> optimizer_buffer_tensors =
         local_training_backing.local_tensor_backing.tensor_optimizer_mapping.at(
             weight_tensor);
 
     // get invocation
     TaskInvocation invocation = get_update_invocation(optimizer_attrs,
                                                       weight_tensor,
                                                       weight_grad_tensor,
                                                       optimizer_buffer_tensors);
 
     // TODO: https://github.com/flexflow/flexflow-train/issues/1442
     // assert(is_invocation_valid(get_update_signature(attrs), invocation));
 
     // execute update
     TaskArgumentAccessor accessor =
         get_task_arg_accessor(local_training_backing.local_tensor_backing,
                               local_training_backing.local_args_backing,
                               invocation);
     TaskImplFunction update_impl_fn = get_update_task_impl(optimizer_attrs);
     update_impl_fn.get<GenericTaskImplFunction>().function_ptr(accessor);
   }
 }
 
 TaskArgumentAccessor
     get_task_arg_accessor(LocalTensorBacking const &local_tensor_backing,
                           LocalArgsBacking const &local_args_backing,
                           TaskInvocation const &invocation) {
   TensorSlotsBacking tensor_slots_backing =
       construct_tensor_slots_backing(local_tensor_backing, invocation.binding);
   ArgSlotsBacking arg_slots_backing = construct_arg_slots_backing(
       invocation.binding, local_args_backing.runtime_arg_config);
   return TaskArgumentAccessor::create<LocalTaskArgumentAccessor>(
       local_tensor_backing.allocator, tensor_slots_backing, arg_slots_backing);
 }
 
 } // namespace FlexFlow

lib/local-execution/src/optimizer_tensor_source.cc

@@ -11,4 +11,8 @@ optimizer_tensor_t OptimizerTensorSource::new_optimizer_tensor() {
       OptimizerTensorSource::next_available_optimizer_tensor_id++};
 }
 
+void OptimizerTensorSource::reset() {
+  OptimizerTensorSource::next_available_optimizer_tensor_id = 0;
+}
+
 } // namespace FlexFlow

lib/local-execution/test/src/test_unallocated_tensors.cc

@@ -17,15 +17,26 @@ using namespace ::FlexFlow;
 TEST_SUITE(FF_TEST_SUITE) {
   TEST_CASE("UnallocatedTensors") {
     MockTensorGuidSource tensor_guid_source;
+    GradientTensorSource gradient_tensor_source;
     OptimizerTensorSource optimizer_tensor_source;
 
+    gradient_tensor_source.reset();
+    optimizer_tensor_source.reset();
+
     Allocator allocator = create_local_cpu_memory_allocator();
 
     tensor_guid_t mock_tensor_1 = tensor_guid_source.new_mock_tensor_guid();
     tensor_guid_t mock_tensor_2 = tensor_guid_source.new_mock_tensor_guid();
     tensor_guid_t mock_tensor_3_with_grad =
         tensor_guid_source.new_mock_tensor_guid();
 
+    gradient_tensor_t grad_tensor =
+        gradient_tensor_source.new_gradient_tensor();
+    optimizer_tensor_t optimizer_tensor_1 =
+        optimizer_tensor_source.new_optimizer_tensor();
+    optimizer_tensor_t optimizer_tensor_2 =
+        optimizer_tensor_source.new_optimizer_tensor();
+
     TensorAttrs tensor_attrs_1_no_grad = TensorAttrs{
         TensorShape{TensorDims{FFOrdered<nonnegative_int>{16_n, 10_n}},
                     DataType::FLOAT},
@@ -61,13 +72,10 @@ TEST_SUITE(FF_TEST_SUITE) {
     SUBCASE("Without optimizer") {
       SUBCASE("AllocatedTensors is empty") {
         AllocatedTensors empty = AllocatedTensors{{}, {}, {}};
-        GradientTensorSource gradient_tensor_source;
+        gradient_tensor_source.reset();
         UnallocatedTensors result = generate_unallocated_tensors(
             empty, tensor_attrs_mapping, gradient_tensor_source);
 
-        GradientTensorSource mock_gradient_tensor_source;
-        gradient_tensor_t grad_tensor =
-            mock_gradient_tensor_source.new_gradient_tensor();
         std::unordered_map<TensorTypeVariant, TensorShape>
             correct_tensor_type_shapes = {
                 {TensorTypeVariant{mock_tensor_1},
@@ -93,15 +101,12 @@ TEST_SUITE(FF_TEST_SUITE) {
             },
             {},
             {}};
-        GradientTensorSource gradient_tensor_source;
+
+        gradient_tensor_source.reset();
         UnallocatedTensors result =
             generate_unallocated_tensors(allocated_forward_tensors,
                                          tensor_attrs_mapping,
                                          gradient_tensor_source);
-
-        GradientTensorSource mock_gradient_tensor_source;
-        gradient_tensor_t grad_tensor =
-            mock_gradient_tensor_source.new_gradient_tensor();
         std::unordered_map<TensorTypeVariant, TensorShape>
             correct_tensor_type_shapes = {
                 {TensorTypeVariant{mock_tensor_2},
@@ -127,15 +132,13 @@ TEST_SUITE(FF_TEST_SUITE) {
             },
             {},
             {}};
-        GradientTensorSource gradient_tensor_source;
+
+        gradient_tensor_source.reset();
         UnallocatedTensors result =
             generate_unallocated_tensors(allocated_forward_tensors,
                                          tensor_attrs_mapping,
                                          gradient_tensor_source);
 
-        GradientTensorSource mock_gradient_tensor_source;
-        gradient_tensor_t grad_tensor =
-            mock_gradient_tensor_source.new_gradient_tensor();
         std::unordered_map<TensorTypeVariant, TensorShape>
             correct_tensor_type_shapes = {
                 {TensorTypeVariant{grad_tensor},
@@ -149,9 +152,7 @@ TEST_SUITE(FF_TEST_SUITE) {
       }
 
       SUBCASE("AllocatedTensors contains only gradient tensor") {
-        GradientTensorSource gradient_tensor_source;
-        gradient_tensor_t grad_tensor =
-            gradient_tensor_source.new_gradient_tensor();
+
         AllocatedTensors allocated_forward_tensors = AllocatedTensors{
             {
                 {TensorTypeVariant{grad_tensor}, tensor_backing_3},
@@ -178,9 +179,7 @@ TEST_SUITE(FF_TEST_SUITE) {
       }
 
       SUBCASE("AllocatedTensors contains mixture") {
-        GradientTensorSource gradient_tensor_source;
-        gradient_tensor_t grad_tensor =
-            gradient_tensor_source.new_gradient_tensor();
+
         AllocatedTensors allocated_forward_tensors = AllocatedTensors{
             {
                 {TensorTypeVariant{mock_tensor_1}, tensor_backing_1},
@@ -206,9 +205,7 @@ TEST_SUITE(FF_TEST_SUITE) {
       }
 
       SUBCASE("Fully AllocatedTensors") {
-        GradientTensorSource gradient_tensor_source;
-        gradient_tensor_t grad_tensor =
-            gradient_tensor_source.new_gradient_tensor();
+
         AllocatedTensors allocated_forward_tensors = AllocatedTensors{
             {
                 {TensorTypeVariant{mock_tensor_1}, tensor_backing_1},
@@ -235,8 +232,8 @@ TEST_SUITE(FF_TEST_SUITE) {
           OptimizerAttrs attrs =
               OptimizerAttrs{SGDOptimizerAttrs{0.0, momentum, false, 0.0}};
           AllocatedTensors empty = AllocatedTensors{{}, {}, {}};
-          GradientTensorSource gradient_tensor_source;
-          OptimizerTensorSource optimizer_tensour_source;
+
+          gradient_tensor_source.reset();
           UnallocatedTensors result =
               generate_unallocated_tensors_with_optimizer(
                   empty,
@@ -245,9 +242,9 @@ TEST_SUITE(FF_TEST_SUITE) {
                   optimizer_tensor_source,
                   attrs);
 
-          GradientTensorSource mock_gradient_tensor_source;
+          gradient_tensor_source.reset();
           UnallocatedTensors correct = generate_unallocated_tensors(
-              empty, tensor_attrs_mapping, mock_gradient_tensor_source);
+              empty, tensor_attrs_mapping, gradient_tensor_source);
           CHECK(result == correct);
         }
         SUBCASE("with momentum") {
@@ -257,8 +254,9 @@ TEST_SUITE(FF_TEST_SUITE) {
 
           SUBCASE("unallocated") {
             AllocatedTensors empty = AllocatedTensors{{}, {}, {}};
-            GradientTensorSource gradient_tensor_source;
-            OptimizerTensorSource optimizer_tensour_source;
+
+            gradient_tensor_source.reset();
+            optimizer_tensor_source.reset();
             UnallocatedTensors result =
                 generate_unallocated_tensors_with_optimizer(
                     empty,
@@ -267,13 +265,6 @@ TEST_SUITE(FF_TEST_SUITE) {
                     optimizer_tensor_source,
                     attrs);
 
-            GradientTensorSource mock_gradient_tensor_source;
-            gradient_tensor_t grad_tensor =
-                mock_gradient_tensor_source.new_gradient_tensor();
-            OptimizerTensorSource mock_optimizer_tensour_source;
-            optimizer_tensor_t optimizer_tensor =
-                mock_optimizer_tensour_source.new_optimizer_tensor();
-
             std::unordered_map<TensorTypeVariant, TensorShape>
                 correct_tensor_type_shapes = {
                     {TensorTypeVariant{mock_tensor_1},
@@ -284,26 +275,25 @@ TEST_SUITE(FF_TEST_SUITE) {
                      tensor_attrs_3_with_grad.shape},
                     {TensorTypeVariant{grad_tensor},
                      tensor_attrs_3_with_grad.shape},
-                    {TensorTypeVariant{optimizer_tensor},
+                    {TensorTypeVariant{optimizer_tensor_1},
                      tensor_attrs_3_with_grad.shape},
                 };
             UnallocatedTensors correct = UnallocatedTensors{
                 correct_tensor_type_shapes,
                 {{mock_tensor_3_with_grad, grad_tensor}},
-                {{mock_tensor_3_with_grad, {optimizer_tensor}}}};
+                {{mock_tensor_3_with_grad, {optimizer_tensor_1}}}};
 
             CHECK(result == correct);
           }
 
           SUBCASE("allocated") {
-            OptimizerTensorSource optimizer_tensour_source;
-            optimizer_tensor_t optimizer_tensor =
-                optimizer_tensour_source.new_optimizer_tensor();
+
             AllocatedTensors allocated_optimizer_tensor = AllocatedTensors{
-                {{TensorTypeVariant{optimizer_tensor}, tensor_backing_3}},
+                {{TensorTypeVariant{optimizer_tensor_1}, tensor_backing_3}},
                 {},
-                {{mock_tensor_3_with_grad, {optimizer_tensor}}}};
-            GradientTensorSource gradient_tensor_source;
+                {{mock_tensor_3_with_grad, {optimizer_tensor_1}}}};
+
+            gradient_tensor_source.reset();
             UnallocatedTensors result =
                 generate_unallocated_tensors_with_optimizer(
                     allocated_optimizer_tensor,
@@ -312,10 +302,6 @@ TEST_SUITE(FF_TEST_SUITE) {
                     optimizer_tensor_source,
                     attrs);
 
-            GradientTensorSource mock_gradient_tensor_source;
-            gradient_tensor_t grad_tensor =
-                mock_gradient_tensor_source.new_gradient_tensor();
-
             std::unordered_map<TensorTypeVariant, TensorShape>
                 correct_tensor_type_shapes = {
                     {TensorTypeVariant{mock_tensor_1},
@@ -348,8 +334,9 @@ TEST_SUITE(FF_TEST_SUITE) {
                                               /*epsilon=*/1e-8}};
         SUBCASE("Empty") {
           AllocatedTensors empty = AllocatedTensors{{}, {}, {}};
-          GradientTensorSource gradient_tensor_source;
-          OptimizerTensorSource optimizer_tensour_source;
+
+          gradient_tensor_source.reset();
+          optimizer_tensor_source.reset();
           UnallocatedTensors result =
               generate_unallocated_tensors_with_optimizer(
                   empty,
@@ -358,15 +345,6 @@ TEST_SUITE(FF_TEST_SUITE) {
                   optimizer_tensor_source,
                   attrs);
 
-          GradientTensorSource mock_gradient_tensor_source;
-          gradient_tensor_t grad_tensor =
-              mock_gradient_tensor_source.new_gradient_tensor();
-          OptimizerTensorSource mock_optimizer_tensour_source;
-          optimizer_tensor_t optimizer_tensor_1 =
-              mock_optimizer_tensour_source.new_optimizer_tensor();
-          optimizer_tensor_t optimizer_tensor_2 =
-              mock_optimizer_tensour_source.new_optimizer_tensor();
-
           std::unordered_map<TensorTypeVariant, TensorShape>
               correct_tensor_type_shapes = {
                   {TensorTypeVariant{mock_tensor_1},
@@ -391,14 +369,16 @@ TEST_SUITE(FF_TEST_SUITE) {
           CHECK(result == correct);
         }
         SUBCASE("Partially allocated") {
-          OptimizerTensorSource optimizer_tensour_source;
-          optimizer_tensor_t optimizer_tensor_1 =
-              optimizer_tensour_source.new_optimizer_tensor();
+          gradient_tensor_source.reset();
+          optimizer_tensor_source.reset();
+          optimizer_tensor_t optimizer_tensor_pre_allocated =
+              optimizer_tensor_source.new_optimizer_tensor();
           AllocatedTensors allocated_optimizer_tensor = AllocatedTensors{
-              {{TensorTypeVariant{optimizer_tensor_1}, tensor_backing_3}},
+              {{TensorTypeVariant{optimizer_tensor_pre_allocated},
+                tensor_backing_3}},
               {},
-              {{mock_tensor_3_with_grad, {optimizer_tensor_1}}}};
-          GradientTensorSource gradient_tensor_source;
+              {{mock_tensor_3_with_grad, {optimizer_tensor_pre_allocated}}}};
+
           UnallocatedTensors result =
               generate_unallocated_tensors_with_optimizer(
                   allocated_optimizer_tensor,
@@ -407,14 +387,6 @@ TEST_SUITE(FF_TEST_SUITE) {
                   optimizer_tensor_source,
                   attrs);
 
-          GradientTensorSource mock_gradient_tensor_source;
-          gradient_tensor_t grad_tensor =
-              mock_gradient_tensor_source.new_gradient_tensor();
-          OptimizerTensorSource mock_optimizer_tensour_source;
-          optimizer_tensor_source.new_optimizer_tensor();
-          optimizer_tensor_t optimizer_tensor_2 =
-              optimizer_tensour_source.new_optimizer_tensor();
-
           std::unordered_map<TensorTypeVariant, TensorShape>
               correct_tensor_type_shapes = {
                   {TensorTypeVariant{mock_tensor_1},
@@ -437,18 +409,14 @@ TEST_SUITE(FF_TEST_SUITE) {
         }
 
         SUBCASE("Fully allocated") {
-          OptimizerTensorSource optimizer_tensour_source;
-          optimizer_tensor_t optimizer_tensor_1 =
-              optimizer_tensour_source.new_optimizer_tensor();
-          optimizer_tensor_t optimizer_tensor_2 =
-              optimizer_tensour_source.new_optimizer_tensor();
           AllocatedTensors allocated_optimizer_tensor = AllocatedTensors{
               {{TensorTypeVariant{optimizer_tensor_1}, tensor_backing_3},
                {TensorTypeVariant{optimizer_tensor_2}, tensor_backing_3}},
               {},
               {{mock_tensor_3_with_grad,
                 {optimizer_tensor_1, optimizer_tensor_2}}}};
-          GradientTensorSource gradient_tensor_source;
+
+          gradient_tensor_source.reset();
           UnallocatedTensors result =
               generate_unallocated_tensors_with_optimizer(
                   allocated_optimizer_tensor,
@@ -457,10 +425,6 @@ TEST_SUITE(FF_TEST_SUITE) {
                   optimizer_tensor_source,
                   attrs);
 
-          GradientTensorSource mock_gradient_tensor_source;
-          gradient_tensor_t grad_tensor =
-              mock_gradient_tensor_source.new_gradient_tensor();
-
           std::unordered_map<TensorTypeVariant, TensorShape>
               correct_tensor_type_shapes = {
                   {TensorTypeVariant{mock_tensor_1},