diff --git a/luamml-tex-annotate.lua b/luamml-tex-annotate.lua
index cb7a8ed..0f852ad 100644
--- a/luamml-tex-annotate.lua
+++ b/luamml-tex-annotate.lua
@@ -2,7 +2,10 @@ local nest = tex.nest
 
 local properties = node.get_properties_table()
 
-local mark_environment = {}
+local mark_environment = {
+  data = {
+  },
+}
 
 local function annotate()
   local annotation, err = load( 'return {'
@@ -98,3 +101,5 @@ end
 funcid = luatexbase.new_luafunction'__luamml_annotate_end:e'
 token.set_lua('__luamml_annotate_end:e', funcid, 'protected')
 lua.get_functions_table()[funcid] = annotate
+
+return mark_environment
diff --git a/luamml-tex.lua b/luamml-tex.lua
index 8832025..49f645a 100644
--- a/luamml-tex.lua
+++ b/luamml-tex.lua
@@ -159,7 +159,8 @@ lua.get_functions_table()[funcid] = function()
   mlist_result = nil
 end
 
-require'luamml-tex-annotate'
+local annotate_context = require'luamml-tex-annotate'
+annotate_context.data.mathml = labelled_mathml
 
 return {
   save_result = save_result,
diff --git a/luamml.dtx b/luamml.dtx
index 493e7a7..9606158 100644
--- a/luamml.dtx
+++ b/luamml.dtx
@@ -366,16 +366,14 @@
 %    \end{macrocode}
 %
 % \subsection{Trigger for specific formula}
-% This only applies for pdf\TeX\ since in Lua\TeX\ everything is controlled by the callback.
-%
-%    \begin{macrocode}
-%<*pdftex>
-%    \end{macrocode}
+% This only applies for pdf\TeX\ since in Lua\TeX\ everything is controlled by the callback,
+% but for compatibility the function is defined anyway.
 %
 % \begin{macro}{\luamml_pdf_write:}
 % We could accept parameters for the flag and tag here, but for compatibility
 % with Lua\TeX they are passed in macros instead.
 %    \begin{macrocode}
+%<*pdftex>
 \cs_new_protected:Npn \luamml_pdf_write: {
   \int_gincr:N \g__luamml_formula_id_int
   \iow_now:Nx \l__luamml_pdf_stream {
@@ -393,11 +391,12 @@
     LUAMML_FORMULA_END
   }
 }
+%</pdftex>
+%<luatex>\cs_new_eq:NN \luamml_pdf_write: \scan_stop:
 %    \end{macrocode}
 % \end{macro}
 %
 %    \begin{macrocode}
-%</pdftex>
 %    \end{macrocode}
 %
 % \subsection{Further helpers}