Merge pull request #26289 from reikdas:reikdas/docfix

PiperOrigin-RevId: 722878611
jax-ml · Feb 4, 2025 · 6281b86 · 6281b86
2 parents bc1a706 + 72b6704
commit 6281b86
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/docs/autodidax.ipynb b/docs/autodidax.ipynb
@@ -146,7 +146,7 @@
     "around calls to `bind`. These wrappers let us control how arguments are passed\n",
     "to `bind`, and in particular we follow a handy internal convention: when we\n",
     "call `bind`, we pass values representing array data as positional arguments,\n",
-    "and we pass metadata like the `axis` argument to `sum_p` via keyword. This\n",
+    "and we pass metadata like the `axis` argument to `reduce_sum_p` via keyword. This\n",
     "calling convention simplifies some core logic (since e.g. instances of the\n",
     "`Tracer` class to be defined below can only occur in positional arguments to\n",
     "`bind`). The wrappers can also provide docstrings!\n",

diff --git a/docs/autodidax.md b/docs/autodidax.md
@@ -133,7 +133,7 @@ The functions that user code calls, like `add` and `sin`, are just wrappers
 around calls to `bind`. These wrappers let us control how arguments are passed
 to `bind`, and in particular we follow a handy internal convention: when we
 call `bind`, we pass values representing array data as positional arguments,
-and we pass metadata like the `axis` argument to `sum_p` via keyword. This
+and we pass metadata like the `axis` argument to `reduce_sum_p` via keyword. This
 calling convention simplifies some core logic (since e.g. instances of the
 `Tracer` class to be defined below can only occur in positional arguments to
 `bind`). The wrappers can also provide docstrings!

diff --git a/docs/autodidax.py b/docs/autodidax.py
@@ -123,7 +123,7 @@ def bind1(prim, *args, **params):
 # around calls to `bind`. These wrappers let us control how arguments are passed
 # to `bind`, and in particular we follow a handy internal convention: when we
 # call `bind`, we pass values representing array data as positional arguments,
-# and we pass metadata like the `axis` argument to `sum_p` via keyword. This
+# and we pass metadata like the `axis` argument to `reduce_sum_p` via keyword. This
 # calling convention simplifies some core logic (since e.g. instances of the
 # `Tracer` class to be defined below can only occur in positional arguments to
 # `bind`). The wrappers can also provide docstrings!