diff --git a/docs/marimo/linalg_snitch.py b/docs/marimo/linalg_snitch.py
index 9c2e2dc3bf..2ba67ce787 100644
--- a/docs/marimo/linalg_snitch.py
+++ b/docs/marimo/linalg_snitch.py
@@ -10,6 +10,12 @@ def _():
     return (mo,)
 
 
+@app.cell
+def _():
+    import xdsl.utils.marimo as xmo
+    return (xmo,)
+
+
 @app.cell
 def _(mo):
     mo.md(
@@ -360,13 +366,13 @@ def _(mo):
 
 
 @app.cell
-def _(asm_html, mo, riscv_asm_module, riscv_code):
+def _(mo, riscv_asm_module, riscv_code, xmo):
     riscv_asm = riscv_code(riscv_asm_module)
 
     mo.md(f"""\
     **RISC-V Assembly:**
 
-    {asm_html(riscv_asm)}
+    {xmo.asm_html(riscv_asm)}
     """
     )
     return (riscv_asm,)
@@ -458,13 +464,13 @@ def _(k, m, mo, n):
 
 
 @app.cell
-def _(asm_html, mo, riscv_code, snitch_asm_module):
+def _(mo, riscv_code, snitch_asm_module, xmo):
     snitch_asm = riscv_code(snitch_asm_module)
 
     mo.md(f"""\
     **Snitch Assembly:**
 
-    {asm_html(snitch_asm)}
+    {xmo.asm_html(snitch_asm)}
     """
     )
     return (snitch_asm,)
@@ -666,19 +672,14 @@ def format_row(key: str, *values: str):
 
 
 @app.cell
-def _(ModuleOp, mo):
+def _(ModuleOp):
     import html as htmllib
 
     def module_html(module: ModuleOp) -> str:
         return f"""\
         <div style="overflow-y: scroll; height:400px;"><small><code style="white-space: pre-wrap;">{htmllib.escape(str(module))}</code></small></div>
         """
-
-    def asm_html(asm: str) -> mo.Html:
-        return mo.ui.code_editor(
-            asm, language="python", disabled=True
-        )
-    return asm_html, htmllib, module_html
+    return htmllib, module_html
 
 
 @app.cell
diff --git a/xdsl/utils/marimo.py b/xdsl/utils/marimo.py
new file mode 100644
index 0000000000..99bd216b20
--- /dev/null
+++ b/xdsl/utils/marimo.py
@@ -0,0 +1,8 @@
+import marimo as mo
+
+
+def asm_html(asm: str) -> mo.Html:
+    """
+    Returns a Marimo-optimised representation of the assembly code passed in.
+    """
+    return mo.ui.code_editor(asm, language="python", disabled=True)