We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
codediff section not respecting html container 
In the new `sphinx_book_theme` our custom `codediff` directive is rendering its content outside the container, overflowing until the end of the page. As shown in this screenshot, it appear to be rendered underneath the `contents` section because its not respecting its section boundaries:

![image](https://user-images.githubusercontent.com/5862228/173973531-5e548ee8-53f6-4978-9744-f9a3a54b5470.png)



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `docs/conf.py`
Content:
```
1 # Copyright 2022 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Configuration file for the Sphinx documentation builder."""
16 
17 
18 # This file only contains a selection of the most common options. For a full
19 # list see the documentation:
20 # https://www.sphinx-doc.org/en/master/usage/configuration.html
21 
22 # -- Path setup --------------------------------------------------------------
23 
24 # If extensions (or modules to document with autodoc) are in another directory,
25 # add these directories to sys.path here. If the directory is relative to the
26 # documentation root, use os.path.abspath to make it absolute, like shown here.
27 #
28 # import os
29 # import sys
30 # sys.path.insert(0, os.path.abspath('.'))
31 
32 import os
33 import sys
34 sys.path.insert(0, os.path.abspath('..'))
35 # Include local extension.
36 sys.path.append(os.path.abspath('./_ext'))
37 
38 # patch sphinx
39 import docs.conf_sphinx_patch
40 # -- Project information -----------------------------------------------------
41 
42 project = 'Flax'
43 copyright = '2020, The Flax authors'  # pylint: disable=redefined-builtin
44 author = 'The Flax authors'
45 
46 
47 # -- General configuration ---------------------------------------------------
48 
49 # Add any Sphinx extension module names here, as strings. They can be
50 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
51 # ones.
52 extensions = [
53     'sphinx.ext.autodoc',
54     'sphinx.ext.autosummary',
55     'sphinx.ext.autosectionlabel',
56     'sphinx.ext.doctest',
57     'sphinx.ext.intersphinx',
58     'sphinx.ext.mathjax',
59     'sphinx.ext.napoleon',
60     'sphinx.ext.viewcode',
61     'myst_nb',
62     'codediff',
63     'sphinx_markdown_tables'
64 ]
65 
66 # Add any paths that contain templates here, relative to this directory.
67 templates_path = ['_templates']
68 
69 # List of patterns, relative to source directory, that match files and
70 # directories to ignore when looking for source files.
71 # This pattern also affects html_static_path and html_extra_path.
72 exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
73 
74 # The suffix(es) of source filenames.
75 # Note: important to list ipynb before md here: we have both md and ipynb
76 # copies of each notebook, and myst will choose which to convert based on
77 # the order in the source_suffix list. Notebooks which are not executed have
78 # outputs stored in ipynb but not in md, so we must convert the ipynb.
79 source_suffix = ['.rst', '.ipynb', '.md']
80 
81 autosummary_generate = True
82 
83 master_doc = 'index'
84 
85 autodoc_typehints = 'none'
86 
87 
88 # -- Options for HTML output -------------------------------------------------
89 
90 # The theme to use for HTML and HTML Help pages.  See the documentation for
91 # a list of builtin themes.
92 #
93 html_theme = 'sphinx_book_theme'
94 # html_style = 'css/flax_theme.css'
95 
96 # The name of an image file (relative to this directory) to place at the top
97 # of the sidebar.
98 html_logo = './flax.png'
99 
100 # Add any paths that contain custom static files (such as style sheets) here,
101 # relative to this directory. They are copied after the builtin static files,
102 # so a file named "default.css" will overwrite the builtin "default.css".
103 html_static_path = ['_static']
104 
105 html_theme_options = {
106     "repository_url": "https://github.com/google/flax",
107     "use_repository_button": True,     # add a "link to repository" button
108     "use_issues_button": False,        # add an "Open an Issue" button
109     "path_to_docs": "docs",            # used to compute the path to launch notebooks in colab
110     "launch_buttons": {
111         "colab_url": "https://colab.research.google.com/",
112     },
113 }
114 
115 # -- Options for myst ----------------------------------------------
116 # Notebook cell execution timeout; defaults to 30.
117 nb_execution_timeout = 100
118 # List of patterns, relative to source directory, that match notebook
119 # files that will not be executed.
120 nb_execution_excludepatterns = [
121   'notebooks/annotated_mnist.ipynb', # <-- times out 
122 ]
123 # raise exceptions on execution so CI can catch errors
124 nb_execution_allow_errors = False
125 nb_execution_raise_on_error = True
126 
127 # -- Extension configuration -------------------------------------------------
128 
129 # Tell sphinx-autodoc-typehints to generate stub parameter annotations including
130 # types, even if the parameters aren't explicitly documented.
131 always_document_param_types = True
132
```

Path: `docs/_ext/codediff.py`
Content:
```
1 # Copyright 2022 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Sphinx directive for creating code diff tables.
16 
17 Use directive as follows:
18 
19 .. codediff::
20   :title_left: <LEFT_CODE_BLOCK_TITLE>
21   :title_right: <RIGHT_CODE_BLOCK_TITLE>
22 
23   <CODE_BLOCK_LEFT>
24   ---
25   <CODE_BLOCK_RIGHT>
26 
27 In order to highlight a line of code, prepend it with "#!".
28 """
29 import itertools
30 
31 from docutils import nodes
32 from docutils.parsers.rst import directives
33 from docutils.statemachine import ViewList
34 
35 import sphinx
36 from sphinx.util.docutils import SphinxDirective
37 
38 
39 class CodeDiffParser:
40 
41   def parse(self, lines, title_left='Base', title_right='Diff', code_sep='---'):
42     if code_sep not in lines:
43       raise ValueError('Code separator not found! Code snippets should be '
44                        f'separated by {code_sep}.')
45     idx = lines.index(code_sep)
46     code_left = self._code_block(lines[0: idx])
47     test_code = lines[idx+1:]
48     code_right = self._code_block(test_code)
49 
50     self.max_left = max(len(x) for x in code_left + [title_left])
51     self.max_right = max(len(x) for x in code_right + [title_right])
52 
53     output = [
54         self._hline(),
55         self._table_row(title_left, title_right),
56         self._hline(),
57     ]
58 
59     for l, r in itertools.zip_longest(code_left, code_right, fillvalue=''):
60       output += [self._table_row(l, r)]
61 
62     return output + [self._hline()], test_code
63 
64   def _code_block(self, lines):
65     """Creates a codeblock."""
66     # Remove right trailing whitespace so we can detect the comments.
67     lines = [x.rstrip() for x in lines]
68     highlight = lambda x: x.endswith('#!')
69     code = map(lambda x: x[:-2].rstrip() if highlight(x) else x, lines)
70     highlights = [i + 1 for i in range(len(lines)) if highlight(lines[i])]
71     highlights = ','.join(str(i) for i in highlights)
72 
73     directive = ['.. code-block:: python']
74     if highlights:
75       directive += [f'  :emphasize-lines: {highlights}']
76 
77     # Indent code and add empty line so the code is picked up by the directive.
78     return directive + [''] + list(map(lambda x: '  ' + x, code))
79 
80   def _hline(self):
81     return '+' + '-'*(self.max_left+2) + '+' + '-'*(self.max_right+2) + '+'
82 
83   def _rfill(self, text, max_len):
84     return text + ' ' * (max_len-len(text))
85 
86   def _table_row(self, left, right):
87     text_left = self._rfill(left, self.max_left)
88     text_right = self._rfill(right, self.max_right)
89     return '| ' + text_left + ' | ' + text_right + ' |'
90 
91 
92 class CodeDiffDirective(SphinxDirective):
93   has_content = True
94   option_spec = {
95       'title_left': directives.unchanged,
96       'title_right': directives.unchanged,
97       'code_sep': directives.unchanged,
98   }
99 
100   def run(self):
101     table_code, test_code = CodeDiffParser().parse(
102         list(self.content), **self.options)
103 
104     # Create a test node as a comment node so it won't show up in the docs.
105     # We add attribute "testnodetype" so it is be picked up by the doctest
106     # builder. This functionality is not officially documented but can be found
107     # in the source code:
108     # https://github.com/sphinx-doc/sphinx/blob/3.x/sphinx/ext/doctest.py
109     # (search for 'testnodetype').
110     test_code = '\n'.join(test_code)
111     test_node = nodes.comment(test_code, test_code, testnodetype='testcode')
112     # Set the source info so the error message is correct when testing.
113     self.set_source_info(test_node)
114     test_node['options'] = {}
115     test_node['language'] = 'python3'
116 
117     # The table node is the side-by-side diff view that will be shown on RTD.
118     table_node = nodes.paragraph()
119     self.content = ViewList(table_code, self.content.parent)
120     self.state.nested_parse(self.content, self.content_offset, table_node)
121 
122     return [table_node, test_node]
123 
124 
125 def setup(app):
126   app.add_directive('codediff', CodeDiffDirective)
127 
128   return {
129       'version': sphinx.__display_version__,
130       'parallel_read_safe': True,
131       'parallel_write_safe': True,
132   }
133
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/docs/_ext/codediff.py b/docs/_ext/codediff.py
--- a/docs/_ext/codediff.py
+++ b/docs/_ext/codediff.py
@@ -24,9 +24,10 @@
   ---
   <CODE_BLOCK_RIGHT>
 
-In order to highlight a line of code, prepend it with "#!".
+In order to highlight a line of code, append "#!" to it.
 """
 import itertools
+from typing import List, Tuple
 
 from docutils import nodes
 from docutils.parsers.rst import directives
@@ -35,10 +36,14 @@
 import sphinx
 from sphinx.util.docutils import SphinxDirective
 
+MISSING = object()
 
 class CodeDiffParser:
 
-  def parse(self, lines, title_left='Base', title_right='Diff', code_sep='---'):
+  def parse(
+    self, lines, title_left='Base', title_right='Diff', code_sep='---', sync=MISSING):
+    sync = sync is not MISSING
+
     if code_sep not in lines:
       raise ValueError('Code separator not found! Code snippets should be '
                        f'separated by {code_sep}.')
@@ -47,19 +52,10 @@
     test_code = lines[idx+1:]
     code_right = self._code_block(test_code)
 
-    self.max_left = max(len(x) for x in code_left + [title_left])
-    self.max_right = max(len(x) for x in code_right + [title_right])
-
-    output = [
-        self._hline(),
-        self._table_row(title_left, title_right),
-        self._hline(),
-    ]
+    output = self._tabs(
+      (title_left, code_left), (title_right, code_right), sync=sync)
 
-    for l, r in itertools.zip_longest(code_left, code_right, fillvalue=''):
-      output += [self._table_row(l, r)]
-
-    return output + [self._hline()], test_code
+    return output, test_code
 
   def _code_block(self, lines):
     """Creates a codeblock."""
@@ -77,17 +73,20 @@
     # Indent code and add empty line so the code is picked up by the directive.
     return directive + [''] + list(map(lambda x: '  ' + x, code))
 
-  def _hline(self):
-    return '+' + '-'*(self.max_left+2) + '+' + '-'*(self.max_right+2) + '+'
-
-  def _rfill(self, text, max_len):
-    return text + ' ' * (max_len-len(text))
+  def _tabs(self, *contents: Tuple[str, List[str]], sync):
+    output = ['.. tab-set::'] + ['  ']
+    
+    for title, content in contents:
+      output += [f'  .. tab-item:: {title}']
+      
+      if sync:
+        key = title.strip()
+        output += [f'    :sync: {key}']
 
-  def _table_row(self, left, right):
-    text_left = self._rfill(left, self.max_left)
-    text_right = self._rfill(right, self.max_right)
-    return '| ' + text_left + ' | ' + text_right + ' |'
+      output += ['    ']
+      output += ['    ' + line for line in content]
 
+    return output
 
 class CodeDiffDirective(SphinxDirective):
   has_content = True
@@ -95,6 +94,7 @@
       'title_left': directives.unchanged,
       'title_right': directives.unchanged,
       'code_sep': directives.unchanged,
+      'sync': directives.flag,
   }
 
   def run(self):
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -60,7 +60,7 @@
     'sphinx.ext.viewcode',
     'myst_nb',
     'codediff',
-    'sphinx_markdown_tables'
+    'sphinx_design',
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -91,7 +91,7 @@
 # a list of builtin themes.
 #
 html_theme = 'sphinx_book_theme'
-# html_style = 'css/flax_theme.css'
+html_css_files = ["css/flax_theme.css"]
 
 # The name of an image file (relative to this directory) to place at the top
 # of the sidebar.

{"golden_diff": "diff --git a/docs/_ext/codediff.py b/docs/_ext/codediff.py\n--- a/docs/_ext/codediff.py\n+++ b/docs/_ext/codediff.py\n@@ -24,9 +24,10 @@\n   ---\n   <CODE_BLOCK_RIGHT>\n \n-In order to highlight a line of code, prepend it with \"#!\".\n+In order to highlight a line of code, append \"#!\" to it.\n \"\"\"\n import itertools\n+from typing import List, Tuple\n \n from docutils import nodes\n from docutils.parsers.rst import directives\n@@ -35,10 +36,14 @@\n import sphinx\n from sphinx.util.docutils import SphinxDirective\n \n+MISSING = object()\n \n class CodeDiffParser:\n \n-  def parse(self, lines, title_left='Base', title_right='Diff', code_sep='---'):\n+  def parse(\n+    self, lines, title_left='Base', title_right='Diff', code_sep='---', sync=MISSING):\n+    sync = sync is not MISSING\n+\n     if code_sep not in lines:\n       raise ValueError('Code separator not found! Code snippets should be '\n                        f'separated by {code_sep}.')\n@@ -47,19 +52,10 @@\n     test_code = lines[idx+1:]\n     code_right = self._code_block(test_code)\n \n-    self.max_left = max(len(x) for x in code_left + [title_left])\n-    self.max_right = max(len(x) for x in code_right + [title_right])\n-\n-    output = [\n-        self._hline(),\n-        self._table_row(title_left, title_right),\n-        self._hline(),\n-    ]\n+    output = self._tabs(\n+      (title_left, code_left), (title_right, code_right), sync=sync)\n \n-    for l, r in itertools.zip_longest(code_left, code_right, fillvalue=''):\n-      output += [self._table_row(l, r)]\n-\n-    return output + [self._hline()], test_code\n+    return output, test_code\n \n   def _code_block(self, lines):\n     \"\"\"Creates a codeblock.\"\"\"\n@@ -77,17 +73,20 @@\n     # Indent code and add empty line so the code is picked up by the directive.\n     return directive + [''] + list(map(lambda x: '  ' + x, code))\n \n-  def _hline(self):\n-    return '+' + '-'*(self.max_left+2) + '+' + '-'*(self.max_right+2) + '+'\n-\n-  def _rfill(self, text, max_len):\n-    return text + ' ' * (max_len-len(text))\n+  def _tabs(self, *contents: Tuple[str, List[str]], sync):\n+    output = ['.. tab-set::'] + ['  ']\n+    \n+    for title, content in contents:\n+      output += [f'  .. tab-item:: {title}']\n+      \n+      if sync:\n+        key = title.strip()\n+        output += [f'    :sync: {key}']\n \n-  def _table_row(self, left, right):\n-    text_left = self._rfill(left, self.max_left)\n-    text_right = self._rfill(right, self.max_right)\n-    return '| ' + text_left + ' | ' + text_right + ' |'\n+      output += ['    ']\n+      output += ['    ' + line for line in content]\n \n+    return output\n \n class CodeDiffDirective(SphinxDirective):\n   has_content = True\n@@ -95,6 +94,7 @@\n       'title_left': directives.unchanged,\n       'title_right': directives.unchanged,\n       'code_sep': directives.unchanged,\n+      'sync': directives.flag,\n   }\n \n   def run(self):\ndiff --git a/docs/conf.py b/docs/conf.py\n--- a/docs/conf.py\n+++ b/docs/conf.py\n@@ -60,7 +60,7 @@\n     'sphinx.ext.viewcode',\n     'myst_nb',\n     'codediff',\n-    'sphinx_markdown_tables'\n+    'sphinx_design',\n ]\n \n # Add any paths that contain templates here, relative to this directory.\n@@ -91,7 +91,7 @@\n # a list of builtin themes.\n #\n html_theme = 'sphinx_book_theme'\n-# html_style = 'css/flax_theme.css'\n+html_css_files = [\"css/flax_theme.css\"]\n \n # The name of an image file (relative to this directory) to place at the top\n # of the sidebar.\n", "issue": "codediff section not respecting html container \nIn the new `sphinx_book_theme` our custom `codediff` directive is rendering its content outside the container, overflowing until the end of the page. As shown in this screenshot, it appear to be rendered underneath the `contents` section because its not respecting its section boundaries:\r\n\r\n![image](https://user-images.githubusercontent.com/5862228/173973531-5e548ee8-53f6-4978-9744-f9a3a54b5470.png)\r\n\r\n\n", "before_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Configuration file for the Sphinx documentation builder.\"\"\"\n\n\n# This file only contains a selection of the most common options. For a full\n# list see the documentation:\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\n# import os\n# import sys\n# sys.path.insert(0, os.path.abspath('.'))\n\nimport os\nimport sys\nsys.path.insert(0, os.path.abspath('..'))\n# Include local extension.\nsys.path.append(os.path.abspath('./_ext'))\n\n# patch sphinx\nimport docs.conf_sphinx_patch\n# -- Project information -----------------------------------------------------\n\nproject = 'Flax'\ncopyright = '2020, The Flax authors'  # pylint: disable=redefined-builtin\nauthor = 'The Flax authors'\n\n\n# -- General configuration ---------------------------------------------------\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n    'sphinx.ext.autodoc',\n    'sphinx.ext.autosummary',\n    'sphinx.ext.autosectionlabel',\n    'sphinx.ext.doctest',\n    'sphinx.ext.intersphinx',\n    'sphinx.ext.mathjax',\n    'sphinx.ext.napoleon',\n    'sphinx.ext.viewcode',\n    'myst_nb',\n    'codediff',\n    'sphinx_markdown_tables'\n]\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = ['_templates']\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path.\nexclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']\n\n# The suffix(es) of source filenames.\n# Note: important to list ipynb before md here: we have both md and ipynb\n# copies of each notebook, and myst will choose which to convert based on\n# the order in the source_suffix list. Notebooks which are not executed have\n# outputs stored in ipynb but not in md, so we must convert the ipynb.\nsource_suffix = ['.rst', '.ipynb', '.md']\n\nautosummary_generate = True\n\nmaster_doc = 'index'\n\nautodoc_typehints = 'none'\n\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages.  See the documentation for\n# a list of builtin themes.\n#\nhtml_theme = 'sphinx_book_theme'\n# html_style = 'css/flax_theme.css'\n\n# The name of an image file (relative to this directory) to place at the top\n# of the sidebar.\nhtml_logo = './flax.png'\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named \"default.css\" will overwrite the builtin \"default.css\".\nhtml_static_path = ['_static']\n\nhtml_theme_options = {\n    \"repository_url\": \"https://github.com/google/flax\",\n    \"use_repository_button\": True,     # add a \"link to repository\" button\n    \"use_issues_button\": False,        # add an \"Open an Issue\" button\n    \"path_to_docs\": \"docs\",            # used to compute the path to launch notebooks in colab\n    \"launch_buttons\": {\n        \"colab_url\": \"https://colab.research.google.com/\",\n    },\n}\n\n# -- Options for myst ----------------------------------------------\n# Notebook cell execution timeout; defaults to 30.\nnb_execution_timeout = 100\n# List of patterns, relative to source directory, that match notebook\n# files that will not be executed.\nnb_execution_excludepatterns = [\n  'notebooks/annotated_mnist.ipynb', # <-- times out \n]\n# raise exceptions on execution so CI can catch errors\nnb_execution_allow_errors = False\nnb_execution_raise_on_error = True\n\n# -- Extension configuration -------------------------------------------------\n\n# Tell sphinx-autodoc-typehints to generate stub parameter annotations including\n# types, even if the parameters aren't explicitly documented.\nalways_document_param_types = True\n", "path": "docs/conf.py"}, {"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sphinx directive for creating code diff tables.\n\nUse directive as follows:\n\n.. codediff::\n  :title_left: <LEFT_CODE_BLOCK_TITLE>\n  :title_right: <RIGHT_CODE_BLOCK_TITLE>\n\n  <CODE_BLOCK_LEFT>\n  ---\n  <CODE_BLOCK_RIGHT>\n\nIn order to highlight a line of code, prepend it with \"#!\".\n\"\"\"\nimport itertools\n\nfrom docutils import nodes\nfrom docutils.parsers.rst import directives\nfrom docutils.statemachine import ViewList\n\nimport sphinx\nfrom sphinx.util.docutils import SphinxDirective\n\n\nclass CodeDiffParser:\n\n  def parse(self, lines, title_left='Base', title_right='Diff', code_sep='---'):\n    if code_sep not in lines:\n      raise ValueError('Code separator not found! Code snippets should be '\n                       f'separated by {code_sep}.')\n    idx = lines.index(code_sep)\n    code_left = self._code_block(lines[0: idx])\n    test_code = lines[idx+1:]\n    code_right = self._code_block(test_code)\n\n    self.max_left = max(len(x) for x in code_left + [title_left])\n    self.max_right = max(len(x) for x in code_right + [title_right])\n\n    output = [\n        self._hline(),\n        self._table_row(title_left, title_right),\n        self._hline(),\n    ]\n\n    for l, r in itertools.zip_longest(code_left, code_right, fillvalue=''):\n      output += [self._table_row(l, r)]\n\n    return output + [self._hline()], test_code\n\n  def _code_block(self, lines):\n    \"\"\"Creates a codeblock.\"\"\"\n    # Remove right trailing whitespace so we can detect the comments.\n    lines = [x.rstrip() for x in lines]\n    highlight = lambda x: x.endswith('#!')\n    code = map(lambda x: x[:-2].rstrip() if highlight(x) else x, lines)\n    highlights = [i + 1 for i in range(len(lines)) if highlight(lines[i])]\n    highlights = ','.join(str(i) for i in highlights)\n\n    directive = ['.. code-block:: python']\n    if highlights:\n      directive += [f'  :emphasize-lines: {highlights}']\n\n    # Indent code and add empty line so the code is picked up by the directive.\n    return directive + [''] + list(map(lambda x: '  ' + x, code))\n\n  def _hline(self):\n    return '+' + '-'*(self.max_left+2) + '+' + '-'*(self.max_right+2) + '+'\n\n  def _rfill(self, text, max_len):\n    return text + ' ' * (max_len-len(text))\n\n  def _table_row(self, left, right):\n    text_left = self._rfill(left, self.max_left)\n    text_right = self._rfill(right, self.max_right)\n    return '| ' + text_left + ' | ' + text_right + ' |'\n\n\nclass CodeDiffDirective(SphinxDirective):\n  has_content = True\n  option_spec = {\n      'title_left': directives.unchanged,\n      'title_right': directives.unchanged,\n      'code_sep': directives.unchanged,\n  }\n\n  def run(self):\n    table_code, test_code = CodeDiffParser().parse(\n        list(self.content), **self.options)\n\n    # Create a test node as a comment node so it won't show up in the docs.\n    # We add attribute \"testnodetype\" so it is be picked up by the doctest\n    # builder. This functionality is not officially documented but can be found\n    # in the source code:\n    # https://github.com/sphinx-doc/sphinx/blob/3.x/sphinx/ext/doctest.py\n    # (search for 'testnodetype').\n    test_code = '\\n'.join(test_code)\n    test_node = nodes.comment(test_code, test_code, testnodetype='testcode')\n    # Set the source info so the error message is correct when testing.\n    self.set_source_info(test_node)\n    test_node['options'] = {}\n    test_node['language'] = 'python3'\n\n    # The table node is the side-by-side diff view that will be shown on RTD.\n    table_node = nodes.paragraph()\n    self.content = ViewList(table_code, self.content.parent)\n    self.state.nested_parse(self.content, self.content_offset, table_node)\n\n    return [table_node, test_node]\n\n\ndef setup(app):\n  app.add_directive('codediff', CodeDiffDirective)\n\n  return {\n      'version': sphinx.__display_version__,\n      'parallel_read_safe': True,\n      'parallel_write_safe': True,\n  }\n", "path": "docs/_ext/codediff.py"}], "after_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Configuration file for the Sphinx documentation builder.\"\"\"\n\n\n# This file only contains a selection of the most common options. For a full\n# list see the documentation:\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\n# import os\n# import sys\n# sys.path.insert(0, os.path.abspath('.'))\n\nimport os\nimport sys\nsys.path.insert(0, os.path.abspath('..'))\n# Include local extension.\nsys.path.append(os.path.abspath('./_ext'))\n\n# patch sphinx\nimport docs.conf_sphinx_patch\n# -- Project information -----------------------------------------------------\n\nproject = 'Flax'\ncopyright = '2020, The Flax authors'  # pylint: disable=redefined-builtin\nauthor = 'The Flax authors'\n\n\n# -- General configuration ---------------------------------------------------\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n    'sphinx.ext.autodoc',\n    'sphinx.ext.autosummary',\n    'sphinx.ext.autosectionlabel',\n    'sphinx.ext.doctest',\n    'sphinx.ext.intersphinx',\n    'sphinx.ext.mathjax',\n    'sphinx.ext.napoleon',\n    'sphinx.ext.viewcode',\n    'myst_nb',\n    'codediff',\n    'sphinx_design',\n]\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = ['_templates']\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path.\nexclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']\n\n# The suffix(es) of source filenames.\n# Note: important to list ipynb before md here: we have both md and ipynb\n# copies of each notebook, and myst will choose which to convert based on\n# the order in the source_suffix list. Notebooks which are not executed have\n# outputs stored in ipynb but not in md, so we must convert the ipynb.\nsource_suffix = ['.rst', '.ipynb', '.md']\n\nautosummary_generate = True\n\nmaster_doc = 'index'\n\nautodoc_typehints = 'none'\n\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages.  See the documentation for\n# a list of builtin themes.\n#\nhtml_theme = 'sphinx_book_theme'\nhtml_css_files = [\"css/flax_theme.css\"]\n\n# The name of an image file (relative to this directory) to place at the top\n# of the sidebar.\nhtml_logo = './flax.png'\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named \"default.css\" will overwrite the builtin \"default.css\".\nhtml_static_path = ['_static']\n\nhtml_theme_options = {\n    \"repository_url\": \"https://github.com/google/flax\",\n    \"use_repository_button\": True,     # add a \"link to repository\" button\n    \"use_issues_button\": False,        # add an \"Open an Issue\" button\n    \"path_to_docs\": \"docs\",            # used to compute the path to launch notebooks in colab\n    \"launch_buttons\": {\n        \"colab_url\": \"https://colab.research.google.com/\",\n    },\n}\n\n# -- Options for myst ----------------------------------------------\n# Notebook cell execution timeout; defaults to 30.\nnb_execution_timeout = 100\n# List of patterns, relative to source directory, that match notebook\n# files that will not be executed.\nnb_execution_excludepatterns = [\n  'notebooks/annotated_mnist.ipynb', # <-- times out \n]\n# raise exceptions on execution so CI can catch errors\nnb_execution_allow_errors = False\nnb_execution_raise_on_error = True\n\n# -- Extension configuration -------------------------------------------------\n\n# Tell sphinx-autodoc-typehints to generate stub parameter annotations including\n# types, even if the parameters aren't explicitly documented.\nalways_document_param_types = True\n", "path": "docs/conf.py"}, {"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sphinx directive for creating code diff tables.\n\nUse directive as follows:\n\n.. codediff::\n  :title_left: <LEFT_CODE_BLOCK_TITLE>\n  :title_right: <RIGHT_CODE_BLOCK_TITLE>\n\n  <CODE_BLOCK_LEFT>\n  ---\n  <CODE_BLOCK_RIGHT>\n\nIn order to highlight a line of code, append \"#!\" to it.\n\"\"\"\nimport itertools\nfrom typing import List, Tuple\n\nfrom docutils import nodes\nfrom docutils.parsers.rst import directives\nfrom docutils.statemachine import ViewList\n\nimport sphinx\nfrom sphinx.util.docutils import SphinxDirective\n\nMISSING = object()\n\nclass CodeDiffParser:\n\n  def parse(\n    self, lines, title_left='Base', title_right='Diff', code_sep='---', sync=MISSING):\n    sync = sync is not MISSING\n\n    if code_sep not in lines:\n      raise ValueError('Code separator not found! Code snippets should be '\n                       f'separated by {code_sep}.')\n    idx = lines.index(code_sep)\n    code_left = self._code_block(lines[0: idx])\n    test_code = lines[idx+1:]\n    code_right = self._code_block(test_code)\n\n    output = self._tabs(\n      (title_left, code_left), (title_right, code_right), sync=sync)\n\n    return output, test_code\n\n  def _code_block(self, lines):\n    \"\"\"Creates a codeblock.\"\"\"\n    # Remove right trailing whitespace so we can detect the comments.\n    lines = [x.rstrip() for x in lines]\n    highlight = lambda x: x.endswith('#!')\n    code = map(lambda x: x[:-2].rstrip() if highlight(x) else x, lines)\n    highlights = [i + 1 for i in range(len(lines)) if highlight(lines[i])]\n    highlights = ','.join(str(i) for i in highlights)\n\n    directive = ['.. code-block:: python']\n    if highlights:\n      directive += [f'  :emphasize-lines: {highlights}']\n\n    # Indent code and add empty line so the code is picked up by the directive.\n    return directive + [''] + list(map(lambda x: '  ' + x, code))\n\n  def _tabs(self, *contents: Tuple[str, List[str]], sync):\n    output = ['.. tab-set::'] + ['  ']\n    \n    for title, content in contents:\n      output += [f'  .. tab-item:: {title}']\n      \n      if sync:\n        key = title.strip()\n        output += [f'    :sync: {key}']\n\n      output += ['    ']\n      output += ['    ' + line for line in content]\n\n    return output\n\nclass CodeDiffDirective(SphinxDirective):\n  has_content = True\n  option_spec = {\n      'title_left': directives.unchanged,\n      'title_right': directives.unchanged,\n      'code_sep': directives.unchanged,\n      'sync': directives.flag,\n  }\n\n  def run(self):\n    table_code, test_code = CodeDiffParser().parse(\n        list(self.content), **self.options)\n\n    # Create a test node as a comment node so it won't show up in the docs.\n    # We add attribute \"testnodetype\" so it is be picked up by the doctest\n    # builder. This functionality is not officially documented but can be found\n    # in the source code:\n    # https://github.com/sphinx-doc/sphinx/blob/3.x/sphinx/ext/doctest.py\n    # (search for 'testnodetype').\n    test_code = '\\n'.join(test_code)\n    test_node = nodes.comment(test_code, test_code, testnodetype='testcode')\n    # Set the source info so the error message is correct when testing.\n    self.set_source_info(test_node)\n    test_node['options'] = {}\n    test_node['language'] = 'python3'\n\n    # The table node is the side-by-side diff view that will be shown on RTD.\n    table_node = nodes.paragraph()\n    self.content = ViewList(table_code, self.content.parent)\n    self.state.nested_parse(self.content, self.content_offset, table_node)\n\n    return [table_node, test_node]\n\n\ndef setup(app):\n  app.add_directive('codediff', CodeDiffDirective)\n\n  return {\n      'version': sphinx.__display_version__,\n      'parallel_read_safe': True,\n      'parallel_write_safe': True,\n  }\n", "path": "docs/_ext/codediff.py"}]}