diff --git a/docs/_build/scripts/macro.py b/docs/_build/scripts/macro.py
index 5dad68348415..15fbf9e1ddc2 100644
--- a/docs/_build/scripts/macro.py
+++ b/docs/_build/scripts/macro.py
@@ -131,6 +131,23 @@ def code_tab(
 
 
 def define_env(env):
+    @env.macro
+    def code_header(language: str, section: str = [], api_functions: List[str] = []) -> str:
+        language_info = LANGUAGES[language]
+
+        language = language_info["code_name"]
+
+        # Create feature flags
+        feature_flags_links = create_feature_flag_links(language, api_functions)
+
+        # Create API Links if they are defined in the YAML
+        api_functions = [
+            link for f in api_functions if (link := create_api_function_link(language, f))
+        ]
+        language_headers = " ·".join(api_functions + feature_flags_links)
+        return f"""=== \":fontawesome-brands-{language_info['icon_name']}: {language_info['display_name']}\"
+    {language_headers}"""
+
     @env.macro
     def code_block(
         path: str, section: str = None, api_functions: List[str] = None
diff --git a/docs/src/python/user-guide/lazy/gpu.py b/docs/src/python/user-guide/lazy/gpu.py
index 14a7de1fcc7a..21d246d09d76 100644
--- a/docs/src/python/user-guide/lazy/gpu.py
+++ b/docs/src/python/user-guide/lazy/gpu.py
@@ -1,5 +1,4 @@
 # --8<-- [start:setup]
-# --8<-- [start:simple]
 import polars as pl
 
 df = pl.LazyFrame({"a": [1.242, 1.535]})
@@ -7,10 +6,6 @@
 q = df.select(pl.col("a").round(1))
 # --8<-- [end:setup]
 
-result = q.collect(engine="gpu")
-print(result)
-# --8<-- [end:simple]
-
 # Avoiding requiring the GPU engine for doc build output
 # --8<-- [start:simple-result]
 result = q.collect()
@@ -18,21 +13,16 @@
 # --8<-- [end:simple-result]
 
 
-# --8<-- [start:engine]
 # --8<-- [start:engine-setup]
 q = df.select((pl.col("a") ** 4))
 
 # --8<-- [end:engine-setup]
-result = q.collect(engine=pl.GPUEngine(device=1))
-print(result)
-# --8<-- [end:engine]
 
 # --8<-- [start:engine-result]
 result = q.collect()
 print(result)
 # --8<-- [end:engine-result]
 
-# --8<-- [start:fallback-warning]
 # --8<-- [start:fallback-setup]
 df = pl.LazyFrame(
     {
@@ -45,12 +35,6 @@
 
 # --8<-- [end:fallback-setup]
 
-with pl.Config() as cfg:
-    cfg.set_verbose(True)
-    result = q.collect(engine="gpu")
-
-print(result)
-# --8<-- [end:fallback-warning]
 # --8<-- [start:fallback-result]
 print(
     "PerformanceWarning: Query execution with GPU not supported, reason: \n"
diff --git a/docs/user-guide/gpu-support.md b/docs/user-guide/gpu-support.md
index 0c5c0c5b5176..bb76058fc9f0 100644
--- a/docs/user-guide/gpu-support.md
+++ b/docs/user-guide/gpu-support.md
@@ -30,7 +30,13 @@ You can install the GPU backend for Polars with a feature flag as part of a norm
 
 Having built a query using the lazy API [as normal](lazy/index.md), GPU-enabled execution is requested by running `.collect(engine="gpu")` instead of `.collect()`.
 
-{{code_block('user-guide/lazy/gpu', 'simple', ['collect'])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:setup"
+
+result = q.collect(engine="gpu")
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:setup"
@@ -39,7 +45,12 @@ Having built a query using the lazy API [as normal](lazy/index.md), GPU-enabled
 
 For more detailed control over the execution, for example to specify which GPU to use on a multi-GPU node, we can provide a `GPUEngine` object. By default, the GPU engine will use a configuration applicable to most use cases.
 
-{{code_block('user-guide/lazy/gpu', 'engine', ['GPUEngine'])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:engine-setup"
+result = q.collect(engine=pl.GPUEngine(device=1))
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:engine-setup"
@@ -87,16 +98,32 @@ The release of the GPU engine in Open Beta implies that we expect things to work
 
 When running in verbose mode, any queries that cannot execute on the GPU will issue a `PerformanceWarning`:
 
-{{code_block('user-guide/lazy/gpu', 'fallback-warning', ['Config'])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:fallback-setup"
+
+with pl.Config() as cfg:
+    cfg.set_verbose(True)
+    result = q.collect(engine="gpu")
+
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:fallback-setup"
---8<-- "python/user-guide/lazy/gpu.py:fallback-result"
+print(
+    "PerformanceWarning: Query execution with GPU not supported, reason: \n"
+    "<class 'NotImplementedError'>: Grouped rolling window not implemented"
+)
+print("# some details elided")
+print()
+print(q.collect())
 ```
 
 To disable fallback, and have the GPU engine raise an exception if a query is unsupported, we can pass an appropriately configured `GPUEngine` object:
 
-```python exec="on" result="text" session="user-guide/lazy"
+{{ code_header("python", [], []) }}
+```python
 q.collect(engine=pl.GPUEngine(raise_on_fail=True))
 ```
 
diff --git a/docs/user-guide/lazy/gpu.md b/docs/user-guide/lazy/gpu.md
index 026a3f8c70e1..60069412bf79 100644
--- a/docs/user-guide/lazy/gpu.md
+++ b/docs/user-guide/lazy/gpu.md
@@ -4,7 +4,13 @@ Polars provides an in-memory, GPU-accelerated execution engine for the Lazy API
 
 If you install Polars with the [GPU feature flag](../installation.md), you can trigger GPU-based execution by running `.collect(engine="gpu")` instead of `.collect()`.
 
-{{code_block('user-guide/lazy/gpu', 'simple', [])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:setup"
+
+result = q.collect(engine="gpu")
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:setup"