diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py
index 98f812f1..4683b1e6 100644
--- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py
+++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py
@@ -667,6 +667,8 @@ def clone(self) -> "TaxBenefitSystem":
                 "_parameters_at_instant_cache",
                 "variables",
                 "entities",
+                "person_entity",
+                "group_entities"
             ):
                 new_dict[key] = value
 
@@ -676,10 +678,19 @@ def clone(self) -> "TaxBenefitSystem":
             variable_name: variable.clone()
             for variable_name, variable in self.variables.items()
         }
+
+        # Apply shallow copies to all relevant entities
         new_dict["entities"] = [copy.copy(entity) for entity in self.entities]
+        new_dict["person_entity"] = copy.copy(self.person_entity)
+        new_dict["group_entities"] = [copy.copy(entity) for entity in self.group_entities]
 
+        # For all shallow-copied entities, set entity._tax_benefit_system to the new system
         for entity in new_dict["entities"]:
             entity.set_tax_benefit_system(new)
+        for entity in new_dict["group_entities"]:
+            entity.set_tax_benefit_system(new)
+        new_dict["person_entity"].set_tax_benefit_system(new)
+        
         return new
 
     def entities_plural(self) -> dict: