diff --git a/crates/gpui3_macros/src/derive_element.rs b/crates/gpui3_macros/src/derive_element.rs
index 6dcc84751da6f41a578fab3c01fb12d85ac26c41..bfe7a301c2092e23da24862caead19da4902f3b0 100644
--- a/crates/gpui3_macros/src/derive_element.rs
+++ b/crates/gpui3_macros/src/derive_element.rs
@@ -57,27 +57,34 @@ pub fn derive_element(input: TokenStream) -> TokenStream {
                 None
             }
 
-            fn layout(
+            fn initialize(
                 &mut self,
                 view_state: &mut Self::ViewState,
-                element_state: Option<Self::ElementState>,
-                cx: &mut gpui3::ViewContext<Self::ViewState>,
-            ) -> (gpui3::LayoutId, Self::ElementState) {
+                _: Option<Self::ElementState>,
+                cx: &mut gpui3::ViewContext<Self::ViewState>
+            ) -> Self::ElementState {
                 use gpui3::IntoAnyElement;
 
-                let mut rendered_element = self.render(view_state, cx).into_any();
-                let layout_id = rendered_element.layout(view_state, cx);
-                (layout_id, rendered_element)
+                self.render(view_state, cx).into_any()
+            }
+
+            fn layout(
+                &mut self,
+                view_state: &mut Self::ViewState,
+                rendered_element: &mut Self::ElementState,
+                cx: &mut gpui3::ViewContext<Self::ViewState>,
+            ) -> gpui3::LayoutId {
+                rendered_element.layout(view_state, cx)
             }
 
             fn paint(
                 &mut self,
                 bounds: gpui3::Bounds<gpui3::Pixels>,
                 view_state: &mut Self::ViewState,
-                element_state: &mut Self::ElementState,
+                rendered_element: &mut Self::ElementState,
                 cx: &mut gpui3::ViewContext<Self::ViewState>,
             ) {
-                element_state.paint(view_state, None, cx)
+                rendered_element.paint(view_state, None, cx)
             }
         }
     };
diff --git a/crates/ui2/src/theme.rs b/crates/ui2/src/theme.rs
index 479b81d9468663968928f8af0dbdc1e6723ee4d5..aa84e7c11f1954646879010e264e12481a0376aa 100644
--- a/crates/ui2/src/theme.rs
+++ b/crates/ui2/src/theme.rs
@@ -164,31 +164,42 @@ impl<E: Element> Element for Themed<E> {
         None
     }
 
-    fn layout(
+    fn initialize(
         &mut self,
-        state: &mut E::ViewState,
+        view_state: &mut Self::ViewState,
         element_state: Option<Self::ElementState>,
+        cx: &mut ViewContext<Self::ViewState>,
+    ) -> Self::ElementState {
+        cx.with_global(self.theme.clone(), |cx| {
+            self.child.initialize(view_state, element_state, cx)
+        })
+    }
+
+    fn layout(
+        &mut self,
+        view_state: &mut E::ViewState,
+        element_state: &mut Self::ElementState,
         cx: &mut ViewContext<E::ViewState>,
-    ) -> (LayoutId, Self::ElementState)
+    ) -> LayoutId
     where
         Self: Sized,
     {
         cx.with_global(self.theme.clone(), |cx| {
-            self.child.layout(state, element_state, cx)
+            self.child.layout(view_state, element_state, cx)
         })
     }
 
     fn paint(
         &mut self,
         bounds: Bounds<Pixels>,
-        state: &mut Self::ViewState,
+        view_state: &mut Self::ViewState,
         frame_state: &mut Self::ElementState,
         cx: &mut ViewContext<Self::ViewState>,
     ) where
         Self: Sized,
     {
         cx.with_global(self.theme.clone(), |cx| {
-            self.child.paint(bounds, state, frame_state, cx);
+            self.child.paint(bounds, view_state, frame_state, cx);
         });
     }
 }