In these notebooks we present some prerequisites for using coax. We do this by showing how to implement linear regression in pure JAX as well as in Haiku. The latter is the library on which the coax function approximators are based.
To keep things simple, we'll just do some very basic linear regression to illustrate the basics of JAX and Haiku.
jax haiku