Repository: google/jax

Issue:
Documentation issue - associative_scan
https://github.com/google/jax/blob/e95d5701e33bda30299a8c87aa43b61a686c1b01/jax/lax/lax_control_flow.py#L2352

Hi,

It should clearly be used from lax, and not numpy + it seems that the Example format does not follow the docstring style, so it doesn't show properly in the readthedocs page.

Adrien 


Discussion:
Thanks for raising the issue! We'll make certain this gets corrected.

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -2329,33 +2329,34 @@ def associative_scan(fn, elems):
 
   Args:
     fn: Python callable implementing an associative binary operation with
-      signature `r = fn(a, b)`. This must satisfy associativity:
-      `fn(a, fn(b, c)) == fn(fn(a, b), c)`. The inputs and result are
-      (possibly nested structures of) `Tensor`(s), matching `elems`. Each
-      `Tensor` has a leading batch dimension in place of `num_elems`; the `fn`
-      is expected to map over this dimension. The result `r` has the same shape
-      (and structure) as the two inputs `a` and `b`.
-    elems: A (possibly nested structure of) `Tensor`(s), each with leading
-      dimension `num_elems`, which must be known statically.
-  Returns:
-    result: A (possibly nested structure of) `Tensor`(s) of the same shape
-      and structure as `elems`, in which the `k`th element is the result of
-      recursively applying `fn` to combine the first `k` elements of
-      `elems`. For example, given `elems = [a, b, c, ...]`, the result
-      would be `[a, fn(a, b), fn(fn(a, b), c), ...]`.
 
-  #### Examples
+      signature ``r = fn(a, b)``. This must satisfy associativity:
+      ``fn(a, fn(b, c)) == fn(fn(a, b), c)``. The inputs and result are
+      (possibly nested structures of) array(s) matching ``elems``. Each
+      array has a leading dimension in place of ``num_elems``; the `fn`
+      is expected to be scanned over this dimension. The result `r` has the same
+      shape (and structure) as the two inputs ``a`` and ``b``.
+    elems: A (possibly nested structure of) array(s), each with leading
+      dimension ``num_elems``.
+
+  Returns:
+    result: A (possibly nested structure of) array(s) of the same shape
+      and structure as ``elems``, in which the ``k``th element is the result of
+      recursively applying ``fn`` to combine the first ``k`` elements of
+      ``elems``. For example, given ``elems = [a, b, c, ...]``, the result
+      would be ``[a, fn(a, b), fn(fn(a, b), c), ...]``.
 
-  ```python
-  # Example 1: Partials sums of numbers.
+  Example 1: partial sums of an array of numbers:
 
-  np.associative_scan(operator.add, np.arange(0, 4))
-  # ==> [ 0, 1, 3, 6]
+  >>> lax.associative_scan(jnp.add, jnp.arange(0, 4))
+  [ 0, 1, 3, 6]
 
-  # Example 2: Partial products of random matrices.
+  Example 2: partial products of an array of matrices
 
-  np.associative_scan(np.matmul, matrices)
-  ```
+  >>> mats = random.uniform(random.PRNGKey(0), (4, 2, 2))
+  >>> partial_prods = lax.associative_scan(jnp.matmul, mats)
+  >>> partial_prods.shape
+  (4, 2, 2)
   """
   elems_flat, tree = tree_flatten(elems)
 


Test patch: