test codediff

chiamp · May 3, 2024 · 92dab2d · 92dab2d
1 parent d0e080d
commit 92dab2d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/docs/_ext/codediff.py b/docs/_ext/codediff.py
@@ -62,6 +62,13 @@ def parse(
       (title_left, code_left), (title_right, code_right), sync=sync
     )
 
+    print('\nPRINT START')
+    print('\n'+'='*20+'OUTPUT'+'='*20)
+    print(output)
+    print('\n'+'='*20+'TEST CODE'+'='*20)
+    print(test_code)
+    print('\nEND PRINT')
+
     return output, test_code
 
   def _code_block(self, lines):

diff --git a/docs/experimental/nnx/transforms.rst b/docs/experimental/nnx/transforms.rst
@@ -57,6 +57,7 @@ the transformed function.
   train_step(model, x, y)
 
   ---
+
   @jax.jit #!
   def train_step(graphdef, state, x, y): #!
     def loss_fn(graphdef, state): #!

diff --git a/docs/guides/converting_and_upgrading/haiku_migration_guide.rst b/docs/guides/converting_and_upgrading/haiku_migration_guide.rst
@@ -358,6 +358,7 @@ return value.
     params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)
 
     return params, new_state
+
   ---
 
   def train_step(params, batch_stats, inputs, labels):
-Original file line number
+Diff line change
@@ Expand Up / @@ -358,6 +358,7 @@ return value. @@
         params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)
         return params, new_state
       ---
       def train_step(params, batch_stats, inputs, labels):
@@ Expand Down @@