How to read "type signatures"?

[nalzok]

The documentation occasionally contains "type signatures", e.g.

jax.lax.scan(f, init, xs, length=None, reverse=False, unroll=1)

Scan a function over leading array axes while carrying along state.

The type signature in brief is

scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])

where we use [t] here to denote the type t with an additional leading axis. That is, if t is an array type then [t] represents the type with an additional leading axis, and if t is a pytree (container) type with array leaves then [t] represents the type with the same pytree structure and corresponding leaves each with an additional leading axis.

What does that mean? It resembles a Haskell-style type signature, but I'm asking anyway just to be sure.

[jakevdp]

It is a Haskell-style type signature. More info here: #8355

I don't think that issue led to anything being done, but we should probably improve the documentation around these.

[jakevdp]

Hopefully #10700 will make this more clear. Thanks for the question!

[nalzok]

Thanks! I think it would also be helpful to have a short introduction to the syntax. To be honest, I am still confused about the signature after consulting Wikipedia, and nobody wants to learn Haskell just to understand the documentation!

scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])

Is it equivalent to the following one?

scan :: (c -> a -> (c, b)) -> (c -> [a] -> (c, [b]))

[sharadmv]

In Haskell,

scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])

is equivalent with

scan :: (c -> a -> (c, b)) -> (c -> [a] -> (c, [b]))

because in Haskell, the -> in type signatures is right-associative (i.e. a -> b -> c == a -> (b -> c)). Functions in Haskell are curried (you can call a function with partial arguments and it will return a function that takes in the remaining arguments).

However, in Python, functions are not curried so we should be careful not to stretch this type-signature analogy too far.