Repository: google/flax

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)




Discussion:

Resolving patch:
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 @@ def parse(self, lines, title_left='Base', title_right='Diff', code_sep='---'):
     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 @@ def _code_block(self, lines):
     # 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 @@ class CodeDiffDirective(SphinxDirective):
       '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.


Test patch: