diff --git a/src/bonsai/bonsai/tool/drawing.py b/src/bonsai/bonsai/tool/drawing.py
index ff17b307d0..6f12faf5b5 100644
--- a/src/bonsai/bonsai/tool/drawing.py
+++ b/src/bonsai/bonsai/tool/drawing.py
@@ -1342,17 +1342,32 @@ class Drawing(bonsai.core.tool.Drawing):
     ) -> Union[ifcopenshell.entity_instance, bool, None]:
         if drawing == reference_element:
             return True
+        if reference_element.is_a("IfcGridAxis"):
+            # We cannot associate IfcGridAxis directly, so we establish a convention:
+            # IfcRelAssignsToProduct.RelatingProduct = IfcGrid
+            # IfcRelAssignsToProduct.Name = IfcGridAxis.AxisTag
+            grid = None
+            for attribute in ("PartOfW", "PartOfV", "PartOfU"):
+                if getattr(reference_element, attribute, None):
+                    grid = getattr(reference_element, attribute)[0]
+                    break
+
+            for element in cls.get_group_elements(cls.get_drawing_group(drawing)):
+                if not element.is_a("IfcAnnotation"):
+                    continue
+                for rel in element.HasAssignments:
+                    if (
+                        rel.is_a("IfcRelAssignsToProduct")
+                        and rel.RelatingProduct == grid
+                        and rel.Name == reference_element.AxisTag
+                    ):
+                        return element
+            return
         for element in cls.get_group_elements(cls.get_drawing_group(drawing)):
             if element.is_a("IfcAnnotation"):
                 for rel in element.HasAssignments:
-                    if rel.is_a("IfcRelAssignsToProduct"):
-                        if rel.RelatingProduct == reference_element:
-                            return element
-                        # We cannot associate IfcGridAxis directly, so we establish a convention:
-                        # IfcRelAssignsToProduct.RelatingProduct = IfcGrid
-                        # IfcRelAssignsToProduct.Name = IfcGridAxis.AxisTag
-                        elif reference_element.is_a("IfcGridAxis") and rel.Name == reference_element.AxisTag:
-                            return element
+                    if rel.is_a("IfcRelAssignsToProduct") and rel.RelatingProduct == reference_element:
+                        return element
 
     @classmethod
     def generate_reference_annotation(
diff --git a/src/ifcopenshell-python/ifcopenshell/api/drawing/assign_product.py b/src/ifcopenshell-python/ifcopenshell/api/drawing/assign_product.py
index 87f786d807..e8b22c1024 100644
--- a/src/ifcopenshell-python/ifcopenshell/api/drawing/assign_product.py
+++ b/src/ifcopenshell-python/ifcopenshell/api/drawing/assign_product.py
@@ -59,49 +59,47 @@ def assign_product(
         ifcopenshell.api.drawing.assign_product(model,
             relating_product=furniture, related_object=annotation)
     """
-    is_grid_axis = relating_product.is_a("IfcGridAxis")
-
-    if is_grid_axis:
-        if related_object.HasAssignments:
-            for rel in related_object.HasAssignments:
-                if rel.is_a("IfcRelAssignsToProduct") and rel.Name == relating_product.AxisTag:
-                    return
-    elif related_object.HasAssignments:
-        for rel in related_object.HasAssignments:
-            if rel.is_a("IfcRelAssignsToProduct") and rel.RelatingProduct == relating_product:
-                return
-
-    referenced_by = None
-
-    if is_grid_axis:
+    if relating_product.is_a("IfcGridAxis"):
         axis = relating_product
         grid = None
         for attribute in ("PartOfW", "PartOfV", "PartOfU"):
             if getattr(axis, attribute, None):
                 grid = getattr(axis, attribute)[0]
-        for rel in grid.ReferencedBy:
-            if rel.Name == axis.AxisTag:
-                referenced_by = rel
                 break
-    elif relating_product.ReferencedBy:
+        for rel in grid.ReferencedBy or []:
+            if rel.Name == axis.AxisTag and related_object in rel.RelatedObjects:
+                return
+
+        return file.create_entity(
+            "IfcRelAssignsToProduct",
+            GlobalId=ifcopenshell.guid.new(),
+            Name=axis.AxisTag,
+            OwnerHistory=ifcopenshell.api.owner.create_owner_history(file),
+            RelatedObjects=[related_object],
+            RelatingProduct=grid,
+        )
+
+    if related_object.HasAssignments:
+        for rel in related_object.HasAssignments:
+            if rel.is_a("IfcRelAssignsToProduct") and rel.RelatingProduct == relating_product:
+                return
+
+    referenced_by = None
+    if relating_product.ReferencedBy:
         referenced_by = relating_product.ReferencedBy[0]
 
     if referenced_by:
         related_objects = list(referenced_by.RelatedObjects)
         related_objects.append(related_object)
         referenced_by.RelatedObjects = related_objects
-        ifcopenshell.api.owner.update_owner_history(file, **{"element": referenced_by})
+        ifcopenshell.api.owner.update_owner_history(file, element=referenced_by)
     else:
         referenced_by = file.create_entity(
             "IfcRelAssignsToProduct",
-            **{
-                "GlobalId": ifcopenshell.guid.new(),
-                "OwnerHistory": ifcopenshell.api.owner.create_owner_history(file),
-                "RelatedObjects": [related_object],
-                "RelatingProduct": relating_product,
-            },
+            GlobalId=ifcopenshell.guid.new(),
+            OwnerHistory=ifcopenshell.api.owner.create_owner_history(file),
+            RelatedObjects=[related_object],
+            RelatingProduct=relating_product,
         )
 
-    if is_grid_axis:
-        referenced_by.Name = axis.AxisTag
     return referenced_by
diff --git a/src/ifcopenshell-python/ifcopenshell/api/drawing/unassign_product.py b/src/ifcopenshell-python/ifcopenshell/api/drawing/unassign_product.py
index f95a98a1fe..7b7245e6d6 100644
--- a/src/ifcopenshell-python/ifcopenshell/api/drawing/unassign_product.py
+++ b/src/ifcopenshell-python/ifcopenshell/api/drawing/unassign_product.py
@@ -51,6 +51,14 @@ def unassign_product(
         ifcopenshell.api.drawing.unassign_product(model,
             relating_product=furniture, related_object=annotation)
     """
+    if relating_product.is_a("IfcGridAxis"):
+        grid = None
+        for attribute in ("PartOfW", "PartOfV", "PartOfU"):
+            if getattr(relating_product, attribute, None):
+                grid = getattr(relating_product, attribute)[0]
+                break
+        relating_product = grid
+
     for rel in related_object.HasAssignments or []:
         if not rel.is_a("IfcRelAssignsToProduct") or rel.RelatingProduct != relating_product:
             continue
diff --git a/src/ifcopenshell-python/test/api/drawing/test_assign_product.py b/src/ifcopenshell-python/test/api/drawing/test_assign_product.py
index 2a0bcf4421..c2464203c0 100644
--- a/src/ifcopenshell-python/test/api/drawing/test_assign_product.py
+++ b/src/ifcopenshell-python/test/api/drawing/test_assign_product.py
@@ -38,6 +38,17 @@ class TestAssignProduct(test.bootstrap.IFC4):
         assert len(wall.ReferencedBy) == 1
         assert wall.ReferencedBy[0].RelatedObjects == (label,)
 
+    def test_assigning_a_grid_axis(self):
+        axis = self.file.createIfcGridAxis(AxisTag="A")
+        grid = self.file.createIfcGrid(UAxes=[axis])
+        line = self.file.createIfcAnnotation()
+        ifcopenshell.api.drawing.assign_product(self.file, relating_product=axis, related_object=line)
+        assert grid.ReferencedBy[0].RelatedObjects == (line,)
+        assert grid.ReferencedBy[0].Name == "A"
+        assert len(self.file.by_type("IfcRelAssignsToProduct")) == 1
+        ifcopenshell.api.drawing.assign_product(self.file, relating_product=axis, related_object=line)
+        assert len(self.file.by_type("IfcRelAssignsToProduct")) == 1
+
 
 class TestAssignProductIFC2X3(test.bootstrap.IFC2X3, TestAssignProduct):
     pass
diff --git a/src/ifcopenshell-python/test/api/drawing/test_unassign_product.py b/src/ifcopenshell-python/test/api/drawing/test_unassign_product.py
index b8b049ca51..f80a9b2ecc 100644
--- a/src/ifcopenshell-python/test/api/drawing/test_unassign_product.py
+++ b/src/ifcopenshell-python/test/api/drawing/test_unassign_product.py
@@ -28,6 +28,14 @@ class TestUnassignProduct(test.bootstrap.IFC4):
         ifcopenshell.api.drawing.unassign_product(self.file, relating_product=wall, related_object=label)
         assert len(self.file.by_type("IfcRelAssignsToProduct")) == 0
 
+    def test_unassigning_a_grid_axis(self):
+        axis = self.file.createIfcGridAxis(AxisTag="A")
+        grid = self.file.createIfcGrid(UAxes=[axis])
+        line = self.file.createIfcAnnotation()
+        ifcopenshell.api.drawing.assign_product(self.file, relating_product=axis, related_object=line)
+        ifcopenshell.api.drawing.unassign_product(self.file, relating_product=axis, related_object=line)
+        assert len(self.file.by_type("IfcRelAssignsToProduct")) == 0
+
 
 class TestUnassignProductIFC2X3(test.bootstrap.IFC2X3, TestUnassignProduct):
     pass