feat: Add item method to Extend Template (#104)

synnada-ai · Dec 24, 2024 · f65eafe · f65eafe
1 parent 980e226
commit f65eafe
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 3 deletions.
diff --git a/mithril/framework/common.py b/mithril/framework/common.py
@@ -870,7 +870,7 @@ def __getitem__(
             start, stop, step = key.start, key.stop, key.step
             return ExtendTemplate(connections=[self, start, stop, step], model="slice")
         elif isinstance(key, int | tuple):
-            return ExtendTemplate(connections=[self, key], model="item")
+            return ExtendTemplate(connections=[self, key], model="get_item")
         else:
             raise TypeError(f"Unsupported key type: {type(key)}")
 
@@ -1072,6 +1072,9 @@ def transpose(self, axes: tuple[int, ...] | TemplateBase | None = None):
     def split(self, split_size: int, axis: int):
         return ExtendTemplate(connections=[self, split_size, axis], model="split")
 
+    def item(self):
+        return ExtendTemplate(connections=[self], model="item")
+
 
 class ExtendTemplate(TemplateBase):
     output_connection: ConnectionData | None

diff --git a/mithril/framework/logical/model.py b/mithril/framework/logical/model.py
@@ -56,6 +56,7 @@
     FloorDivide,
     Greater,
     GreaterEqual,
+    Item,
     Length,
     Less,
     LessEqual,
@@ -110,6 +111,7 @@
     "size": Size,
     "tensor": ToTensor,
     "list": TensorToList,
+    "item": Item,
     "mean": Mean,
     "sqrt": Sqrt,
     "exp": Exponential,
@@ -138,8 +140,8 @@
 
 
 coercion_table: dict[tuple[str, type[Tensor] | type[Scalar]], type[PrimitiveModel]] = {
-    ("item", Tensor): TensorItem,
-    ("item", Scalar): ScalarItem,
+    ("get_item", Tensor): TensorItem,
+    ("get_item", Scalar): ScalarItem,
     ("slice", Tensor): TensorSlice,
     ("slice", Scalar): PrimitiveSlice,
 }

diff --git a/tests/scripts/test_extend_template.py b/tests/scripts/test_extend_template.py
@@ -32,6 +32,7 @@
     Greater,
     GreaterEqual,
     IOKey,
+    Item,
     Less,
     LessEqual,
     Linear,
@@ -1676,3 +1677,23 @@ def test_immediate_values_with_extend_template_and_regular_case():
         == big_model_1.conns.latent_input_keys
         == {"$1"}
     )
+
+
+def test_item():
+    model1 = Model(enforce_jit=False)
+
+    buffer_model_1 = Buffer()
+    item_model = Item()
+    totensor = ToTensor()
+
+    model1 += buffer_model_1(input="input")
+    model1 += item_model(input=buffer_model_1.output)
+    model1 += totensor(input=item_model.output, output=IOKey("output"))
+
+    model2 = Model(enforce_jit=False)
+    buffer_model_1 = Buffer()
+    model2 += buffer_model_1(input="input")
+    conn = buffer_model_1.output.item()
+    model2 += ToTensor()(input=conn, output=IOKey("output"))
+
+    check_logical_models(model1, model2)