Removed Asserts in Python Bindings (#1545)

* Removed Asserts in Python Bindings * Changed explorer use of asserted binding
tenstorrent · Dec 10, 2024 · f0e03af · f0e03af
1 parent f2f2e97
commit f0e03af
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 18 deletions.
diff --git a/python/TTModule.cpp b/python/TTModule.cpp
@@ -65,14 +65,17 @@ void populateTTModule(py::module &m) {
                 .withElementType(unwrap(ctx), unwrap(elementType));
           })
       .def("getLayout",
-           [](MlirType &type) {
-             assert(isa<RankedTensorType>(
-                 unwrap(type))); // Make sure that this is operating on a
-                                 // RankedTensorType object
+           [](MlirType &type) -> std::variant<tt::MetalLayoutAttr, py::object> {
+             // Make sure that this is operating on a RankedTensorType object
+             if (not isa<RankedTensorType>(unwrap(type))) {
+               return py::none();
+             }
              RankedTensorType tensor =
                  mlir::cast<RankedTensorType>(unwrap(type));
-             assert(tensor.getEncoding()); // Make sure that this Tensor has an
-                                           // encoding value
+             // Make sure that this Tensor has an encoding value
+             if (not tensor.getEncoding()) {
+               return py::none();
+             }
              tt::MetalLayoutAttr layout =
                  mlir::cast<tt::MetalLayoutAttr>(tensor.getEncoding());
              return layout;

diff --git a/python/TTNNModule.cpp b/python/TTNNModule.cpp
@@ -156,12 +156,14 @@ void populateTTNNModule(py::module &m) {
       .def_property_readonly(
           "memref",
           [](tt::ttnn::TTNNLayoutAttr self) { return wrap(self.getMemref()); })
-      .def_property_readonly(
-          "memory_layout_as_int", [](tt::ttnn::TTNNLayoutAttr self) {
-            if (!self.getMemLayout()) {
-              assert(false && "Memory layout is not set");
-            }
-            return static_cast<uint32_t>(self.getMemLayout().getValue());
-          });
+      .def_property_readonly("memory_layout_as_int",
+                             [](tt::ttnn::TTNNLayoutAttr self)
+                                 -> std::variant<uint32_t, py::object> {
+                               if (!self.getMemLayout()) {
+                                 return py::none();
+                               }
+                               return static_cast<uint32_t>(
+                                   self.getMemLayout().getValue());
+                             });
 }
 } // namespace mlir::ttmlir::python
diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py
@@ -398,12 +398,14 @@ def parse_ttnn_ttnn_layout(attr):
     layout = ttnn.ir.TTNNLayoutAttr.maybe_downcast(attr)
     result = []
     result.append(graph_builder.KeyValue(key="linear", value=str(layout.linear)))
-    result.append(
-        graph_builder.KeyValue(
-            key="memory_layout",
-            value=str(ttnn.TensorMemoryLayout(layout.memory_layout_as_int)),
+    memory_layout = layout.memory_layout_as_int
+    if memory_layout is not None:
+        result.append(
+            graph_builder.KeyValue(
+                key="memory_layout",
+                value=str(ttnn.TensorMemoryLayout(memory_layout)),
+            )
         )
-    )
     result.append(
         graph_builder.KeyValue(
             key="grid_shape", value="x".join(map(str, layout.grid_attr.shape))