Haskell Type System

Haskell is a pure functional programming language with a strong static type system. This means that all functions are pure and all types are known at compile time. This is in contrast to Python, which is an impure functional programming language with a weak dynamic type system. This means that functions can be impure and types are not known until runtime.

Haskell started out as a research project in the 1980s. It has becoming increasingly popular, but not because it is a practical programming language. It is popular because it is a great language for learning about functional programming and type systems. It has a very simple syntax and a beautiful type system. This is why many of Jax authors use Haskell types in their documentation.

Primitive Types

There are some basic Haskell types that you need to know about. The most basic type is the type Int, which is the type of integers. The type Int is a primitive type, which means that it is built into the language. The type Int is a concrete type, which means that it is not parameterized with any other type. Another type is the type Bool, which is the type of booleans. It is also primitive and concrete. Other primitive and concrete types are Char, Double, and Float.

Composite Types

Lists are a generic or parametric type, which means that they are parameterized with another type. Their type is written as [T], where T is a type variable that can be replaced with any type. For example, the type [Int] is the type of lists of integers. [Float] is the type of lists of floats. [Bool] is the type of lists of booleans. [Char] is the type of lists of characters. And so on.

Lists of lists are written as [[T]], where T is a type variable that can be replaced with any type. And lists of lists of lists are written as [[[T]]], and so on.

Function Types

The type of functions is written as T -> U, where T and U are type variables that can be replaced with any type. For example a function that takes an integer and returns an integer has the type signature:

f:: Int -> Int

A type that takes a list of floats and returns a float has the type signature:

g: [Float] -> Float

You can also have a function that works with a list of any type and return a result of the same type:

h:: [T] -> T

This is called a polymorphic function, which means that it can work with multiple types.

Curried Functions

You can also have a function that takes two arguments of the same type and returns a result of the same type. Think of it as f(x, y) in Python. It has the type signature:

f:: T -> T -> T

Why this notation? You can think of the function f as a function that takes one argument of type T and returns a function that takes one argument of type T and returns a result of type T. Like this:

• Take and x of type T.

• Plug x into f(x, .) and get a function g:: T -> T defined by g(y) = f(x, y).

This is called a curried function, which means that it takes multiple arguments one at a time. The name comes from Haskell Curry, who was a logician and mathematician.

Another similar example is a function that takes two arguments of different types and returns a result of the same type as the second argument.

f:: T -> U -> U

Same deal. Think of the function f as a function that takes one argument of type T and returns a function that takes one argument of type U and returns a result of type U. So f(x, .) is a function that takes one argument of type U and returns a result of type U.

Now take a function with three arguments of different types, returning a result that is the same type as the third argument.

f:: T -> U -> V -> V

Do you get it? The function f takes one argument of type T and returns a function that takes one argument of type U and returns a function that takes one argument of type V and returns a result of type V.

Types of higher-order functions

Type signature of map

For example, the function Map which takes a function of one argument, and a list, and returns a list containing the result of applying the function to each element of the list. Here is its type signature:

map:: (T -> U) -> [T] -> [U]

Let’s break it down:

• It takes one argument of type T -> U. This is a function that takes an argument of type T and returns a result of type U.

• It takes another argument of type [T]. This is a list of values of type T.

• It returns a result of type [U]. This is a list of values of type U. It has to be a U because the function we provide as an argument returns a U.

Type signature of filter

Another example is the Filter function, which takes a function of one argument, and a list, and returns a list containing the elements of the list for which the function returns True. Here is its type signature:

filter:: (T -> Bool) -> [T] -> [T]

Type of function composition

Another example is the Function composition operator, which takes two functions and returns a function that is the composition of the two functions. Its type signature is:

compose:: (U -> V) -> (T -> U) -> (T -> V)

Meditate on this for a while:

• It takes one argument of type U -> V.

• It takes another argument of type T -> U.

• It returns a result of type T -> V.

Types of Jax functions

In Jax we work with arrays instead of lists. The type of the elements of an array is called the dtype of the array. It can be a primitive type, such as int or float, or a composite type, such as float32 or float64. In the Jax documentation, they write a, b, etc., for the primitive types, and [a], [b], etc., for the array types.

Type signature of vmap

Recall that the function takes a function of one argument, and returns a function that can be applied to an array. In its simplest form, the type signature of vmap is:

vmap:: (a -> b) -> ([a] -> [b])

Let’s break it down:

• It takes one argument of type a -> b. This is a function that takes an argument of type a and returns a result of type b.

• It takes another argument of type [a]. This is a list of values of type a.

Signature of fori_loop

Many of the functions in jax.lax are documented with Haskell type signatures. The , called as fori_loop(lower, upper, body_fun, init_val), is essentially doing something equivalent this:

def fori_loop(lower, upper, body_fun, init_val):
  val = init_val
  for i in range(lower, upper):
    val = body_fun(i, val)
  return val

The signature of fori_loop is:

fori_loop:: Int -> Int -> ((Int, a) -> a) -> a -> a

So:

• It takes one argument of type Int. This is the lower bound of the loop.

• It takes another argument of type Int. This is the upper bound of the loop.

• It takes another argument of type (Int, a) -> a. This is a function that takes a tuple of type (Int, a) and returns a result of type a. The integer is the loop index and the a is the accumulator.

• It takes another argument of type a. This is the initial value of the accumulator.

• It returns a result of type a. This is the final value of the accumulator.

Signature of grad

The function takes a function of one argument, and returns a function that computes the gradient of the function. In its simplest form, the type signature of grad is:

grad:: ([a] -> a) -> ([a] -> [a])

Let’s break it down:

• It takes one argument of type [a] -> a. This is the function we want to differentiate. It is a scalar function.

• It returns a result of type [a] -> [a]. This is the function that computes the gradient of the function. It is a vector function. This is the gradient with respect to the input argument.

Signature of jacobian

The function takes a function of one argument, and returns a function that computes the Jacobian of the function. In its simplest form, the type signature of jacobian is:

jacobian:: ([a] -> [a]) -> ([a] -> [[a]])

So:

• It takes one argument of type [a] -> [a]. This is the function we want to differentiate. It is a vector function.

• It returns a result of type [a] -> [[a]]. This is the function that computes the Jacobian of the function. It is a matrix function. This is the Jacobian with respect to the input argument.

Signature of hessian

The function takes a function of one argument, and returns a function that computes the Hessian of the function.

In its simplest form, the type signature of hessian is:

hessian:: ([a] -> a) -> ([a] -> [[a]])

So:

• It takes one argument of type [a] -> a. This is the function we want to differentiate. It is a scalar function.

• It returns a result of type [a] -> [[a]]. This is the function that computes the Hessian of the function. It is a matrix function. This is the Hessian with respect to the input argument.

Signature of jvp

The function takes a function of one argument, and returns a function that computes the Jacobian-vector product of the function. In its simplest form, the type signature of jvp is:

jvp:: ([a] -> a) -> [a] -> [a] -> [a]