From 6423741db02139d17d9c46e58350eb662b2ebd13 Mon Sep 17 00:00:00 2001
From: tornikeo <tornikeonoprishvili@gmail.com>
Date: Tue, 22 Feb 2022 13:49:48 +0400
Subject: [PATCH] Add copy button to code snippets in documentation

---
 docs/conf.py          | 1 +
 docs/requirements.txt | 1 +
 2 files changed, 2 insertions(+)

diff --git a/docs/conf.py b/docs/conf.py
index c75d97ec9..42e5bc261 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -70,6 +70,7 @@ extensions = [
     'sphinx_autodoc_typehints',
     'myst_nb',
     "sphinx_remove_toctrees",
+    'sphinx_copybutton',
     'jax_extensions',
 ]
 
diff --git a/docs/requirements.txt b/docs/requirements.txt
index 09454bb50..b62db0316 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -5,6 +5,7 @@ sphinx-book-theme
 sphinx-remove-toctrees
 # Newer versions cause issues; see https://github.com/google/jax/pull/6449
 sphinx-autodoc-typehints==1.11.1
+sphinx-copybutton>=0.5.0
 jupyter-sphinx>=0.3.2
 myst-nb