fix doc test failure

2025-04-19 05:16:06 +00:00 · 2020-11-24 10:45:03 -08:00 · 2020-11-24 10:45:03 -08:00 · 50cb604f2e
commit 50cb604f2e
parent 8057cf919e
1 changed files with 1 additions and 1 deletions
--- a/docs/jaxpr.rst
+++ b/docs/jaxpr.rst
@ -448,7 +448,7 @@ captured using the ``xla_pmap`` primitive. Consider this example
                                                           shape=(1,) ] 1.0
                                     e = add c d
                                     f = psum[ axis_index_groups=None
-                                               axis_name=rows ] b
+                                               axis_name=('rows',) ] b
                                     g = div e f
                                 in (g,) }
                    devices=None