Repository: google/jax

Issue:
non-Python type signatures in docs are confusing for a Python package
The docs include confusing type signatures, using what I assume is Haskell syntax (please correct me if I'm wrong---I never learned Haskell). Two examples I found are:
https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan
https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html#jax.lax.fori_loop

PS: sorry, I didn't see an option to report an issue with the docs, so just used the "bugs" issue template.


Discussion:
I think the choice is made consiously, it's more 'math like notation' than just 'haskell syntax'.
That may be so, but I'd be surprised if I'm the only reader that can't parse this.
I would argue that it depends on the reader. I personally find the syntax of the documentation much more readable than the convention Python uses for `Callable`s, so I'd err on the side of keeping it.
Perhaps a good compromise would be to link to some resource each time we use haskell signatures?
Yes, that sounds reasonable. At the moment the docs I linked to don't even mention that these are haskell signatures, leaving the reader completely in the dark if they don't happen to be familiar with them already.
Does anyone know a good reference for this haskell signature syntax?
I think the problem people have is not "type signature syntax" per se, but automatic currying. Which is also why Haskell like syntax is so much better, as it provides a lot more affordance. So a reference might be [this](https://wiki.haskell.org/Currying). 

Resolving patch:
diff --git a/docs/autodidax.py b/docs/autodidax.py
--- a/docs/autodidax.py
+++ b/docs/autodidax.py
@@ -1884,7 +1884,9 @@ def pprint_xla_call(names: DefaultDict[Var, str], eqn: JaxprEqn) -> PPrint:
 # ### `linearize`
 #
 # In the case of `linearize`, we want to stage out the linear part of a `jvp`
-# computation. That is, if we have `jvp : (a -> b) -> (a, T a) -> (b, T b)`,
+# computation. That is, in terms of
+# [Haskell-like type signatures](https://wiki.haskell.org/Type_signature),
+# if we have `jvp : (a -> b) -> (a, T a) -> (b, T b)`,
 # then we write `linearize : (a -> b) -> a -> (b, T a -o T b)`, using `T a` to
 # mean "the tangent type of `a`" and using the "lollipop" `-o` rather than the
 # arrow `->` to indicate a _linear_ function. We define the semantics of
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -781,8 +781,8 @@ def custom_gradient(fun):
   and the VJP (gradient) function. See
   https://www.tensorflow.org/api_docs/python/tf/custom_gradient.
 
-  If the mathematical function to be differentiated has type signature ``a ->
-  b``, then the Python callable ``fun`` should have signature
+  If the mathematical function to be differentiated has Haskell-like signature
+  ``a -> b``, then the Python callable ``fun`` should have the signature
   ``a -> (b, CT b --o CT a)`` where we use ``CT x`` to denote a cotangent type
   for ``x`` and the ``--o`` arrow to denote a linear function. See the example
   below. That is, ``fun`` should return a pair where the first element
@@ -1001,7 +1001,7 @@ def linear_call(fun: Callable, fun_transpose: Callable, residual_args,
                 linear_args):
   """Call a linear function, with a custom implementation for its transpose.
 
-  The type signatures of ``fun`` and ``fun_transpose`` are:
+  The `Haskell-like type signatures`_ of ``fun`` and ``fun_transpose`` are:
 
   .. code-block:: haskell
 
@@ -1081,6 +1081,7 @@ def linear_call(fun: Callable, fun_transpose: Callable, residual_args,
   Returns:
     The call result, i.e. ``fun(residual_args, linear_args)``.
 
+  .. _Haskell-like type signatures: https://wiki.haskell.org/Type_signature
   """
   operands_res, res_tree = tree_flatten(residual_args)
   operands_lin, lin_tree = tree_flatten(linear_args)
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -153,7 +153,7 @@ def scanned_fun(loop_carry, _):
 def fori_loop(lower, upper, body_fun, init_val):
   """Loop from ``lower`` to ``upper`` by reduction to :func:`jax.lax.while_loop`.
 
-  The type signature in brief is
+  The `Haskell-like type signature`_ in brief is
 
   .. code-block:: haskell
 
@@ -191,6 +191,8 @@ def fori_loop(lower, upper, body_fun, init_val):
 
   Returns:
     Loop value from the final iteration, of type ``a``.
+
+  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
   """
   if not callable(body_fun):
     raise TypeError("lax.fori_loop: body_fun argument should be callable.")
@@ -235,7 +237,7 @@ def while_loop(cond_fun: Callable[[T], BooleanNumeric],
                init_val: T) -> T:
   """Call ``body_fun`` repeatedly in a loop while ``cond_fun`` is True.
 
-  The type signature in brief is
+  The `Haskell-like type signature`_ in brief is
 
   .. code-block:: haskell
 
@@ -275,6 +277,8 @@ def while_loop(cond_fun, body_fun, init_val):
 
   Returns:
     The output from the final iteration of body_fun, of type ``a``.
+
+  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
   """
   if not (callable(body_fun) and callable(cond_fun)):
     raise TypeError("lax.while_loop: body_fun and cond_fun arguments should be callable.")
@@ -1354,7 +1358,7 @@ def scan(f: Callable[[Carry, X], Tuple[Carry, Y]],
          unroll: int = 1) -> Tuple[Carry, Y]:
   """Scan a function over leading array axes while carrying along state.
 
-  The type signature in brief is
+  The `Haskell-like type signature`_ in brief is
 
   .. code-block:: haskell
 
@@ -1422,6 +1426,8 @@ def scan(f, init, xs, length=None):
     A pair of type ``(c, [b])`` where the first element represents the final
     loop carry value and the second element represents the stacked outputs of
     the second output of ``f`` when scanned over the leading axis of the inputs.
+
+  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
   """
   if not callable(f):
     raise TypeError("lax.scan: f argument should be a callable.")


Test patch: