From 6dab89b8ece7e022bd3df5774c9ddf309e2eb2d9 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 27 Aug 2025 02:47:08 +0100
Subject: [PATCH] [Docs] Fix math rendering in docs (#23676)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 docs/mkdocs/javascript/mathjax.js | 20 ++++++++++++++++++++
 mkdocs.yaml                       |  7 ++++---
 requirements/docs.txt             |  1 -
 3 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 docs/mkdocs/javascript/mathjax.js

diff --git a/docs/mkdocs/javascript/mathjax.js b/docs/mkdocs/javascript/mathjax.js
new file mode 100644
index 0000000000..5da0d44357
--- /dev/null
+++ b/docs/mkdocs/javascript/mathjax.js
@@ -0,0 +1,20 @@
+// Enables MathJax rendering
+window.MathJax = {
+  tex: {
+    inlineMath: [["\\(", "\\)"]],
+    displayMath: [["\\[", "\\]"]],
+    processEscapes: true,
+    processEnvironments: true
+  },
+  options: {
+    ignoreHtmlClass: ".*|",
+    processHtmlClass: "arithmatex"
+  }
+};
+
+document$.subscribe(() => { 
+  MathJax.startup.output.clearCache()
+  MathJax.typesetClear()
+  MathJax.texReset()
+  MathJax.typesetPromise()
+})
diff --git a/mkdocs.yaml b/mkdocs.yaml
index 47fe1ebce9..507a80c41e 100644
--- a/mkdocs.yaml
+++ b/mkdocs.yaml
@@ -129,15 +129,16 @@ markdown_extensions:
   - toc:
       permalink: true
   # For math rendering
-  - mdx_math:
-      enable_dollar_delimiter: true
+  - pymdownx.arithmatex:
+      generic: true
 
 extra_css:
   - mkdocs/stylesheets/extra.css
 
 extra_javascript:
   - mkdocs/javascript/run_llm_widget.js
-  - https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-AMS_HTML
+  - mkdocs/javascript/mathjax.js
+  - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
   - mkdocs/javascript/edit_and_feedback.js
   - mkdocs/javascript/slack_and_forum.js
 
diff --git a/requirements/docs.txt b/requirements/docs.txt
index 3b72a8a9e7..d1c5463987 100644
--- a/requirements/docs.txt
+++ b/requirements/docs.txt
@@ -7,7 +7,6 @@ mkdocs-awesome-nav
 mkdocs-glightbox
 mkdocs-git-revision-date-localized-plugin
 mkdocs-minify-plugin
-python-markdown-math
 regex
 ruff