# Yet another JAX tutorial

Kevin Murphy (murphyk@gmail.com).
Last update: September 2021.

[JAX](https://github.com/google/jax) is a  version of NumPy that runs fast on CPU, GPU and TPU, by compiling the computational graph to XLA (Accelerated Linear Algebra). It also has an excellent automatic differentiation library, extending the earlier [autograd](https://github.com/hips/autograd) package. This library makes it easy to compute higher order derivatives, gradients of complex functions (e.g., optimize an iterative solver), etc.
The JAX interface is almost identical to NumPy (by design), but with some small differences, and lots of additional features.
We give a brief introduction below.
For more details, see [this list of JAX tutorials](https://github.com/probml/probml-notebooks/blob/main/markdown/jax_tutorials.md)

# Setup

```
# Standard Python libraries
from __future__ import absolute_import, division, print_function, unicode_literals

from functools import partial
import os
import time
import numpy as np
np.set_printoptions(precision=3)
import glob
import matplotlib.pyplot as plt
import PIL
import imageio

from typing import Tuple, NamedTuple

from IPython import display
%matplotlib inline

import sklearn

```

```
import jax
import jax.numpy as jnp

from jax import random, vmap, jit, grad, value_and_grad, hessian, jacfwd, jacrev
print("jax version {}".format(jax.__version__))
# Check the jax backend
print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
key = random.PRNGKey(0)
```

    jax version 0.2.19
    jax backend gpu

# Hardware accelerators

Colab makes it easy to use GPUs and TPUs for speeding up some workflows, especially related to deep learning.

## GPUs

Colab offers graphics processing units (GPUs) which can be much faster than CPUs (central processing units), as we illustrate below.

```
# Check if GPU is available and its model, memory ...etc.
!nvidia-smi

```

    Sat Sep 11 03:46:08 2021
    +-----------------------------------------------------------------------------+
    | NVIDIA-SMI 470.63.01    Driver Version: 460.32.03    CUDA Version: 11.2     |
    |-------------------------------+----------------------+----------------------+
    | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
    | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
    |                               |                      |               MIG M. |
    |===============================+======================+======================|
    |   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
    | N/A   66C    P0    29W /  70W |  13610MiB / 15109MiB |      0%      Default |
    |                               |                      |                  N/A |
    +-------------------------------+----------------------+----------------------+

    +-----------------------------------------------------------------------------+
    | Processes:                                                                  |
    |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
    |        ID   ID                                                   Usage      |
    |=============================================================================|
    |  No running processes found                                                 |
    +-----------------------------------------------------------------------------+

```

# Check if JAX is using GPU
print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
# Check the devices avaiable for JAX
jax.devices()
```

    jax backend gpu

    [GpuDevice(id=0, process_index=0)]

Let's see how JAX can speed up things like matrix-matrix multiplication.

First the numpy/CPU version.

```
# Parameters for the experiment
size = int(1e3)
number_of_loops=int(1e2)

```

```
# Standard numpy CPU

def f(x=None):
  if not isinstance(x, np.ndarray):
    x=np.ones((size, size), dtype=np.float32)
  return np.dot(x, x.T)

```

```
%timeit -o -n $number_of_loops f()
```

    100 loops, best of 5: 16.4 ms per loop

    <TimeitResult : 100 loops, best of 5: 16.4 ms per loop>

```
res = _ # get result of last cell
time_cpu = res.best
print(time_cpu)
```

    0.016404767970000192

Now we look at the JAX version. JAX supports execution on [XLA](https://www.tensorflow.org/xla) devices, which can be CPU, GPU or even TPU. We added that block_until_ready because JAX uses [asynchronous execution](https://jax.readthedocs.io/en/latest/async_dispatch.html) by default.

```
# JAX device execution
# https://github.com/google/jax/issues/1598

def jf(x=None):
  if not isinstance(x, jnp.ndarray):
    x=jnp.ones((size, size), dtype=jnp.float32)
  return jnp.dot(x, x.T)

f_gpu = jit(jf, backend='gpu')
f_cpu = jit(jf, backend='cpu')
```

```
# Time the CPU version

%timeit -o -n $number_of_loops f_cpu()
```

    100 loops, best of 5: 25.7 ms per loop

    <TimeitResult : 100 loops, best of 5: 25.7 ms per loop>

```
res = _
time_jcpu = res.best
print(time_jcpu)
```

    0.025695679049999854

```
# Time the GPU version

%timeit -o -n $number_of_loops f_gpu().block_until_ready()
```

    The slowest run took 40.61 times longer than the fastest. This could mean that an intermediate result is being cached.
    100 loops, best of 5: 485 µs per loop

    <TimeitResult : 100 loops, best of 5: 485 µs per loop>

```
res = _
time_jgpu = res.best
print(time_jgpu)
```

    0.00048540602000002764

```
print('JAX CPU time {:0.6f}, Numpy CPU time {:0.6f}, speedup {:0.6f}'.format(
    time_jcpu, time_cpu, time_cpu/time_jcpu))
print('JAX GPU time {:0.6f}, Numpy CPU time {:0.6f}, speedup {:0.6f}'.format(
    time_jgpu, time_cpu, time_cpu/time_jgpu))
```

    JAX CPU time 0.025696, Numpy CPU time 0.016405, speedup 0.638425
    JAX GPU time 0.000485, Numpy CPU time 0.016405, speedup 33.795971

In the above example we see that JAX GPU is much faster than Numpy CPU.
However we also see that JAX CPU is slower than Numpy CPU - this can happen with simple functions, but usually JAX provides a speedup, even on CPU, if you JIT compile a complex function (see below).

We can move numpy arrays to the GPU for speed. The result will be transferred back to CPU for printing, saving, etc.

```
from jax import device_put

x = np.ones((size, size)).astype(np.float32)
print(type(x))
%timeit -o -n $number_of_loops f(x)

x = device_put(x)
print(type(x))
%timeit -o -n $number_of_loops jf(x)
```

    <class 'numpy.ndarray'>
    100 loops, best of 5: 14.8 ms per loop
    <class 'jaxlib.xla_extension.DeviceArray'>
    100 loops, best of 5: 452 µs per loop

    <TimeitResult : 100 loops, best of 5: 452 µs per loop>

## TPUs

We can turn on the tensor processing unit by selecting from the Colab runtime.
Everything else "just works" as before.

```
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
```

If everything is set up correctly, the following command should return a list of 8 TPU devices.

```
jax.local_devices()
```

    [GpuDevice(id=0, process_index=0)]

# Vmap

We often write a function to process a single vector or matrix, and then want to apply it to a batch of data. Using for loops is slow, and manually batchifying code is complex. Fortunately we can use the `vmap` function, which will map our function across a set of inputs, automatically batchifying it.

## Example: 1d convolution

(This example is from the Deepmind tutorial.)

Consider standard 1d convolution of two vectors.

```
x = jnp.arange(5)
w = jnp.array([2., 3., 4.])

def convolve(x, w):
  output = []
  for i in range(1, len(x)-1):
    output.append(jnp.dot(x[i-1:i+2], w))
  return jnp.array(output)

convolve(x, w)
```

    DeviceArray([11., 20., 29.], dtype=float32)

Now suppose we want to convolve multiple vectors with multiple kernels. The simplest way is to use a for loop, but this is slow.

```
xs = jnp.stack([x, x])
ws = jnp.stack([w, w])

def manually_batched_convolve(xs, ws):
  output = []
  for i in range(xs.shape[0]):
    output.append(convolve(xs[i], ws[i]))
  return jnp.stack(output)

manually_batched_convolve(xs, ws)
```

    DeviceArray([[11., 20., 29.],
                 [11., 20., 29.]], dtype=float32)

We can manually vectorize the code, but it is complex.

```
def manually_vectorised_convolve(xs, ws):
  output = []
  for i in range(1, xs.shape[-1] -1):
    output.append(jnp.sum(xs[:, i-1:i+2] * ws, axis=1))
  return jnp.stack(output, axis=1)

manually_vectorised_convolve(xs, ws)
```

    DeviceArray([[11., 20., 29.],
                 [11., 20., 29.]], dtype=float32)

Fortunately vmap can do this for us!

```
auto_batch_convolve = jax.vmap(convolve)

auto_batch_convolve(xs, ws)
```

    DeviceArray([[11., 20., 29.],
                 [11., 20., 29.]], dtype=float32)

## Axes

By default, vmap vectorizes over the first axis of each of its inputs. If the first argument has a batch and the second does not, ,specify `in_axes=[0,None]`, so the second argument is not vectorized over.

```
jax.vmap(convolve, in_axes=[0, None])(xs, w)

```

    DeviceArray([[11., 20., 29.],
                 [11., 20., 29.]], dtype=float32)

We can also vectorize over other dimensions.

```

print(xs.shape)
xst = jnp.transpose(xs)
print(xst.shape)

wst = jnp.transpose(ws)

auto_batch_convolve_v2 = jax.vmap(convolve, in_axes=1, out_axes=1)
auto_batch_convolve_v2(xst, wst)
```

    (2, 5)
    (5, 2)

    DeviceArray([[11., 11.],
                 [20., 20.],
                 [29., 29.]], dtype=float32)

## Example: logistic regression

We now give another example, using binary logistic regression.
Let us start with a predictor for a single example
.

```

D = 2
N = 3

w = np.random.normal(size=(D,))
X = np.random.normal(size=(N,D))

def sigmoid(x): return 0.5 * (jnp.tanh(x / 2.) + 1)

def predict_single(x):
    return sigmoid(jnp.dot(w, x)) # <(D) , (D)> = (1) # inner product

```

```
print(predict_single(X[0,:])) # works

```

    0.23207036

```

print(predict_single(X)) # fails
```

We can manually vectorize the code by remembering the shapes, so
$X w$ multiplies each row of $X$ with $w$.

```
def predict_batch(X):
    return sigmoid(jnp.dot(X, w)) # (N,D) * (D,1) = (N,1) # matrix-vector multiply

print(predict_batch(X))
```

    [0.232 0.121 0.878]

But it easier to use vmap.

```
print(vmap(predict_single)(X))
```

    [0.232 0.121 0.878]

## Failure cases

Vmap requires that the shapes of all the variables that are created by the function that is being mapped are the same for all values of the input arguments, as explained [here](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html). So vmap cannot be used to do any kind of embarassingly parallel task. Below we give a simple example of where this fails, since internally we create a vector whose length depends on the input 'length'.

```
def example_fun(length, val=4):
  return jnp.sum(jnp.ones((length,)) * val)

xs = jnp.arange(1,10)

# Python map works fine
v = list(map(example_fun, xs))
print(v)
```

    [DeviceArray(4., dtype=float32), DeviceArray(8., dtype=float32), DeviceArray(12., dtype=float32), DeviceArray(16., dtype=float32), DeviceArray(20., dtype=float32), DeviceArray(24., dtype=float32), DeviceArray(28., dtype=float32), DeviceArray(32., dtype=float32), DeviceArray(36., dtype=float32)]

The following fails.

```
v = vmap(example_fun)(xs)
print(v)
```

# Stochastics

JAX is designed to be deterministic, but in some cases, we want to introduce randomness in a controlled way, and to reason about it. We discuss this below

## Random number generation

One of the biggest differences from NumPy is the way Jax treates pseudo random number generation (PRNG).
This is because Jax does not maintain any global state, i.e., it is purely functional.
This design "provides reproducible results invariant to compilation boundaries and backends,
while also maximizing performance by enabling vectorized generation and parallelization across random calls"
(to quote [the official page](https://github.com/google/jax#a-brief-tour)).

For example, consider this Numpy snippet. Each call to np.random.uniform updates the global state. The value of foo() is therefore only guaranteed to give the same result every time if we evaluate bar() and baz() in the same order (eg left to right). This is why foo1 and foo2 give different answers even though mathematically they shouldn't (we cannot just substitute in the value of a variable and derive the result, so we are violating "referential transparency").

```
import numpy as np

def bar(): return np.random.uniform(size=(3))
def baz(): return np.random.uniform(size=(3))

def foo(seed):
  np.random.seed(seed)
  return bar() + 2*baz()

def foo1(seed):
  np.random.seed(seed)
  a = bar()
  b =  2*baz()
  return a+b

def foo2(seed):
  np.random.seed(seed)
  a = 2*baz()
  b = bar()
  return a+b

seed = 0

print(foo(seed))
print(foo1(seed))
print(foo2(seed))
```

    [1.639 1.562 1.895]
    [1.639 1.562 1.895]
    [1.643 1.854 1.851]

Jax may evaluate parts of expressions such as `bar() + baz()` in parallel, which would violate reproducibility. To prevent this, the user must pass in an explicit PRNG key to every function that requires a source of randomness. Using the same key will give the same results. See the example below.

```

key = random.PRNGKey(0)
print(random.normal(key, shape=(3,)))  # [ 1.81608593 -0.48262325  0.33988902]
print(random.normal(key, shape=(3,)))  # [ 1.81608593 -0.48262325  0.33988902]  ## identical results

```

    [ 1.816 -0.483  0.34 ]
    [ 1.816 -0.483  0.34 ]

When generating independent samples, it is important to use different keys, to ensure results are not correlated. We can do this by *splitting* the key into the the 'master' key (which will be used in later parts of the code via splitting), and the 'subkey', which is used temporarily to generate randomness and then thrown away, as we illustrate below.

```
# To make a new key, we split the current key into two pieces.
key, subkey = random.split(key)
print(random.normal(subkey, shape=(3,)))  # [ 1.1378783  -1.22095478 -0.59153646]

# We can continue to split off new pieces from the global key.
key, subkey = random.split(key)
print(random.normal(subkey, shape=(3,)))  # [-0.06607265  0.16676566  1.17800343]

```

    [ 1.138 -1.221 -0.592]
    [-0.066  0.167  1.178]

We now reimplement the numpy example in Jax and show that we get the result no matter the order of evaluation of bar and baz.

```

def bar(key):
  return jax.random.uniform(key,shape=(3,))

def baz(key):
  return jax.random.uniform(key,shape=(3,))

def foo(key):
  subkey1, subkey2 = random.split(key, num=2)
  return bar(subkey1) + 2 * baz(subkey2)

def foo1(key):
  subkey1, subkey2 = random.split(key, num=2)
  a = bar(subkey1)
  b =  2 * baz(subkey2)
  return a+b

def foo2(key):
  subkey1, subkey2 = random.split(key, num=2)
  a = 2 * baz(subkey2)
  b = bar(subkey1)
  return a+b

key = random.PRNGKey(0)
key, subkey = random.split(key)
print(foo(subkey))
print(foo1(subkey))
print(foo2(subkey))
```

    [2.079 2.002 1.089]
    [2.079 2.002 1.089]
    [2.079 2.002 1.089]

In Jax (but not in python),  a random draw of N samples in parallel will not give the same results as N draws of individual samples, as we show below.

```
key = random.PRNGKey(42)
subkeys = random.split(key, 3)
sequence = np.stack([jax.random.normal(subkey) for subkey in subkeys])
print("individually:", sequence)

key = random.PRNGKey(42)
print("all at once: ", jax.random.normal(key, shape=(3,)))
```

    individually: [-0.048  0.108 -1.223]
    all at once:  [ 0.187 -1.281 -1.559]

```
np.random.seed(0)
sequence = np.stack([np.random.normal() for i in range(3)])
print("individually:", sequence)

np.random.seed(0)
print("all at once: ", np.random.normal(size=(3,)))

```

    individually: [1.764 0.4   0.979]
    all at once:  [1.764 0.4   0.979]

## Probability distributions

The [distrax library](https://github.com/deepmind/distrax) is a JAX-native implementation of some parts of  the distrbitions library from [Tensorflow Probabilty (TFP)](https://www.tensorflow.org/probability). The main advantage is that the distrax source code is much easier to read and understand. For distributions not in distrax, it is possible to use TFP instead.

Here is a brief example.

```
%%capture
!pip install git+git://github.com/deepmind/distrax.git
```

```
import distrax
import jax
import jax.numpy as jnp

from tensorflow_probability.substrates import jax as tfp
tfd = tfp.distributions

key = jax.random.PRNGKey(1234)
mu = jnp.array([-1., 0., 1.])
sigma = jnp.array([0.1, 0.2, 0.3])

dist_distrax = distrax.MultivariateNormalDiag(mu, sigma)
dist_tfp = tfd.MultivariateNormalDiag(mu, sigma)

samples = dist_distrax.sample(seed=key)

# Both print 1.775
print(dist_distrax.log_prob(samples))
print(dist_tfp.log_prob(samples))
```

    1.7750063
    1.7750063

# Autograd

In this section, we illustrate automatic differentation using JAX.
For details, see see  [this video](https://www.youtube.com/watch?v=wG_nF1awSSY&t=697s)  or [The Autodiff Cookbook](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html).

## Derivatives

We can compute $(\nabla f)(x)$ using `grad(f)(x)`. For example, consider

$f(x) = x^3 + 2x^2 - 3x + 1$

$f'(x) = 3x^2 + 4x -3$

$f''(x) = 6x + 4$

$f'''(x) = 6$

$f^{iv}(x) = 0$

```
f = lambda x: x**3 + 2*x**2 - 3*x + 1

dfdx = jax.grad(f)
d2fdx = jax.grad(dfdx)
d3fdx = jax.grad(d2fdx)
d4fdx = jax.grad(d3fdx)

print(dfdx(1.))
print(d2fdx(1.))
print(d3fdx(1.))
print(d4fdx(1.))
```

    4.0
    10.0
    6.0
    0.0

## Partial derivatives

$$
\begin{align}
f(x,y) &= x^2 + y \\
\frac{\partial f}{\partial x} &= 2x \\
\frac{\partial f}{\partial y} &= 1
\end{align}
$$

```
def f(x,y):
  return x**2 + y

# Partial derviatives
x = 2.0; y= 3.0;
v, gx = value_and_grad(f, argnums=0)(x,y)
print(v)
print(gx)

gy = grad(f, argnums=1)(x,y)
print(gy)

```

    7.0
    4.0
    1.0

## Gradients

Linear function: multi-input, scalar output.

$$
\begin{align}
f(x; a) &= a^T x\\
\nabla_x f(x;a) &= a
\end{align}
$$

```

def fun1d(x):
    return jnp.dot(a, x)[0]

Din = 3; Dout = 1;
a = np.random.normal(size=(Dout, Din))
x = np.random.normal(size=(Din,))

g = grad(fun1d)(x)
assert np.allclose(g, a)

# It is often useful to get the function value and gradient at the same time
val_grad_fn = jax.value_and_grad(fun1d)
v, g = val_grad_fn(x)
print(v)
print(g)
assert np.allclose(v, fun1d(x))
assert np.allclose(a, g)

```

    1.9472518
    [ 2.241  1.868 -0.977]

Linear function: multi-input, multi-output.

$$
\begin{align}
f(x;A) &= A x \\
\frac{\partial f(x;A)}{\partial x} &= A
\end{align}
$$

```
# We construct a multi-output linear function.
# We check forward and reverse mode give same Jacobians.

def fun(x):
    return jnp.dot(A, x)

Din = 3; Dout = 4;
A = np.random.normal(size=(Dout, Din))
x = np.random.normal(size=(Din,))
Jf = jacfwd(fun)(x)
Jr = jacrev(fun)(x)
assert np.allclose(Jf, Jr)
assert np.allclose(Jf, A)
```

Quadratic form.

$$
\begin{align}
f(x;A) &= x^T A x \\
\nabla_x f(x;A) &= (A+A^T) x
\end{align}
$$

```

D = 4
A = np.random.normal(size=(D,D))
x = np.random.normal(size=(D,))
quadfun = lambda x: jnp.dot(x, jnp.dot(A, x))

g = grad(quadfun)(x)
assert np.allclose(g, jnp.dot(A+A.T, x))

```

Chain rule applied to sigmoid function.

$$
\begin{align}
\mu(x;w) &=\sigma(w^T x) \\
\nabla_w \mu(x;w) &= \sigma'(w^T x) x \\
\sigma'(a) &= \sigma(a) * (1-\sigma(a))
\end{align}
$$

```

D = 4
w = np.random.normal(size=(D,))
x = np.random.normal(size=(D,))
y = 0

def sigmoid(x): return 0.5 * (jnp.tanh(x / 2.) + 1)
def mu(w): return sigmoid(jnp.dot(w,x))
def deriv_mu(w): return mu(w) * (1-mu(w)) * x
deriv_mu_jax =  grad(mu)

print(deriv_mu(w))
print(deriv_mu_jax(w))

assert np.allclose(deriv_mu(w), deriv_mu_jax(w), atol=1e-3)

```

    [-0.13  -0.017 -0.072  0.031]
    [-0.13  -0.017 -0.072  0.031]

## Auxiliary return values

A function can return its value and other auxiliary results; the latter are not differentiated.

```
def f(x,y):
  return x**2+y, 42

x = 2.0
y = 3.0
(v,aux), g = value_and_grad(f, has_aux=True)(x,y)
print(v)
print(aux)
print(g)
```

    7.0
    42
    4.0

## Jacobians

Example: Linear function: multi-input, multi-output.

$$
\begin{align}
f(x;A) &= A x \\
\frac{\partial f(x;A)}{\partial x} &= A
\end{align}
$$

```
# We construct a multi-output linear function.
# We check forward and reverse mode give same Jacobians.

def fun(x):
    return jnp.dot(A, x)

Din = 3; Dout = 4;
A = np.random.normal(size=(Dout, Din))
x = np.random.normal(size=(Din,))
Jf = jacfwd(fun)(x)
Jr = jacrev(fun)(x)
assert np.allclose(Jf, Jr)
```

## Hessians

Quadratic form.

$$
\begin{align}
f(x;A) &= x^T A x \\
\nabla_x^2 f(x;A) &= A + A^T
\end{align}
$$

```

D = 4
A = np.random.normal(size=(D,D))
x = np.random.normal(size=(D,))

quadfun = lambda x: jnp.dot(x, jnp.dot(A, x))

H1 = hessian(quadfun)(x)
assert np.allclose(H1, A+A.T)

def my_hessian(fun):
  return jacfwd(jacrev(fun))

H2 = my_hessian(quadfun)(x)
assert np.allclose(H1, H2)
```

## Example: Binary logistic regression

```

def sigmoid(x): return 0.5 * (jnp.tanh(x / 2.) + 1)

def predict_single(w, x):
    return sigmoid(jnp.dot(w, x)) # <(D) , (D)> = (1) # inner product

def predict_batch(w, X):
    return sigmoid(jnp.dot(X, w)) # (N,D) * (D,1) = (N,1) # matrix-vector multiply

# negative log likelihood
def loss(weights, inputs, targets):
    preds = predict_batch(weights, inputs)
    logprobs = jnp.log(preds) * targets + jnp.log(1 - preds) * (1 - targets)
    return -jnp.sum(logprobs)

D = 2
N = 3
w = jax.random.normal(key, shape=(D,))
X = jax.random.normal(key, shape=(N,D))
y = jax.random.choice(key, 2, shape=(N,)) # uniform binary labels
#logits = jnp.dot(X, w)
#y = jax.random.categorical(key, logits)

print(loss(w, X, y))

# Gradient function
grad_fun = grad(loss)

# Gradient of each example in the batch - 2 different ways
grad_fun_w = partial(grad_fun, w)
grads = vmap(grad_fun_w)(X,y)
print(grads)
assert grads.shape == (N,D)

grads2 = vmap(grad_fun, in_axes=(None, 0, 0))(w, X, y)
assert np.allclose(grads, grads2)

# Gradient for entire batch
grad_sum = jnp.sum(grads, axis=0)
assert grad_sum.shape == (D,)
print(grad_sum)
```

    2.384371
    [[ 0.596 -0.621]
     [ 0.315  0.837]
     [-0.233 -0.448]]
    [ 0.677 -0.232]

```
# Textbook implementation of gradient
def NLL_grad(weights, batch):
    X, y = batch
    N = X.shape[0]
    mu = predict_batch(weights, X)
    g = jnp.sum(jnp.dot(jnp.diag(mu - y), X), axis=0)
    return g

grad_sum_batch = NLL_grad(w, (X,y))
print(grad_sum_batch)
assert np.allclose(grad_sum, grad_sum_batch)
```

    [ 0.677 -0.232]

```
# We can also compute Hessians, as we illustrate below.

hessian_fun = hessian(loss)

# Hessian on one example
H0 = hessian_fun(w, X[0,:], y[0])
print('Hessian(example 0)\n{}'.format(H0))

# Hessian for batch
Hbatch = vmap(hessian_fun, in_axes=(None, 0, 0))(w, X, y)
print('Hbatch shape {}'.format(Hbatch.shape))

Hbatch_sum = jnp.sum(Hbatch, axis=0)
print('Hbatch sum\n {}'.format(Hbatch_sum))
```

    Hessian(example 0)
    [[ 0.185 -0.193]
     [-0.193  0.201]]
    Hbatch shape (3, 2, 2)
    Hbatch sum
     [[0.357 0.225]
     [0.225 1.239]]

```
# Textbook implementation of Hessian

def NLL_hessian(weights, batch):
  X, y = batch
  mu = predict_batch(weights, X)
  S = jnp.diag(mu * (1-mu))
  H = jnp.dot(jnp.dot(X.T, S), X)
  return H

H2 = NLL_hessian(w, (X,y) )

assert np.allclose(Hbatch_sum, H2, atol=1e-2)
```

## Vector Jacobian Products (VJP) and Jacobian Vector Products (JVP)

Suppose $f1(x) = W x$ for fixed W, so $J(x) = W$, and $u^T J(x) = W^T u$.
Instead of computing $J$ explicitly and then multiplying by $u$,
wecan do this in one operation.

```
n = 3; m = 2;
W = jax.random.normal(key, shape=(m,n))
x = jax.random.normal(key, shape=(n,))
u = jax.random.normal(key, shape=(m,))

def f1(x): return jnp.dot(W,x)

J1 = jacfwd(f1)(x)
print(J1.shape)

assert np.allclose(J1, W)
tmp1 = jnp.dot(u.T, J1)
print(tmp1)

(val, jvp_fun) = jax.vjp(f1, x)

tmp2 = jvp_fun(u)

assert np.allclose(tmp1, tmp2)

tmp3 = np.dot(W.T, u)
assert np.allclose(tmp1, tmp3)

```

    (2, 3)
    [ 0.888 -0.538 -0.539]

Suppose
$f2(W) = W x$ for fixed $x$.
Now $J(W) = \text{something complex}$,
but $u^T J(W) = J(W)^T u = u x^T$.

```

def f2(W): return jnp.dot(W,x)

J2 = jacfwd(f2)(W)
print(J2.shape)

tmp1 = jnp.dot(u.T, J2)
print(tmp1)
print(tmp1.shape)

(val, jvp_fun) = jax.vjp(f2, W)
tmp2 = jvp_fun(u)
assert np.allclose(tmp1, tmp2)

tmp3 = np.outer(u, x)
assert np.allclose(tmp1, tmp3)

```

    (2, 2, 3)
    [[-0.009 -0.032 -0.474]
     [ 0.005  0.019  0.286]]
    (2, 3)

## Stop-gradient

Sometimes we want to take the gradient of a complex expression wrt some parameters $\theta$, but treating $\theta$ as a constant for some parts of the expression. For example, consider the TD(0) update in reinforcement learning, which has the following form:

$\Delta \theta = (r_t + v_{\theta}(s_t) - v_{\theta}(s_{t-1})) \nabla v_{\theta}(s_{t-1})$

where $s$ is the state, $r$ is the reward, and $v$ is the value function.
This update is not the gradient of any loss function.
However, if the dependency of the target $r_t + v_{\theta}(s_t)$ on the parameter $\theta$ is ignored, it can be written as the gradient of the pseudo loss function

$L(\theta) = [r_t + v_{\theta}(s_t) - v_{\theta}(s_{t-1})]^2$

since

$\nabla_{\theta} L(\theta) = 2 [r_t + v_{\theta}(s_t) - v_{\theta}(s_{t-1})] \nabla v_{\theta}(s_{t-1})$

. We can implement this in JAX using `stop_gradient`, as we show below.

```
def td_loss(theta, s_prev, r_t, s_t):
  v_prev = value_fn(theta, s_prev)
  target = r_t + value_fn(theta, s_t)
  return 0.5*(jax.lax.stop_gradient(target) - v_prev) ** 2

td_update = jax.grad(td_loss)

# An example transition.
s_prev = jnp.array([1., 2., -1.])
r_t = jnp.array(1.)
s_t = jnp.array([2., 1., 0.])

# Value function and initial parameters
value_fn = lambda theta, state: jnp.dot(theta, state)
theta = jnp.array([0.1, -0.1, 0.])

print(td_update(theta, s_prev, r_t, s_t))

```

    [-1.2 -2.4  1.2]

## Straight through estimator

The straight-through estimator is a trick for defining a 'gradient' of a function that is otherwise non-differentiable. Given a non-differentiable function $f : \mathbb{R}^n \to \mathbb{R}^n$ that is used as part of a larger function that we wish to find a gradient of, we simply pretend during the backward pass that $f$ is the identity function, so gradients pass through $f$ ignoring the $f'$ term. This can be implemented neatly using `jax.lax.stop_gradient`.

Here is an example of a non-differentiable function that converts a soft probability distribution to a one-hot vector (discretization).

```
def onehot(labels, num_classes):
  y = (labels[..., None] == jnp.arange(num_classes)[None])
  return y.astype(jnp.float32)

def quantize(y_soft):
  y_hard = onehot(jnp.argmax(y_soft), 3)[0]
  return y_hard

y_soft = np.array([0.1, 0.2, 0.7])
print(quantize(y_soft))

```

    [0. 0. 1.]

Now suppose we define some linear function of the quantized variable of the form $f(y) = w^T q(y)$. If $w=[1,2,3]$ and $q(y)=[0,0,1]$, we get $f(y) = 3$. But the gradient is 0 because $q$ is not differentiable.

```
def f(y):
  w = jnp.array([1,2,3])
  yq = quantize(y)
  return jnp.dot(w, yq)

print(f(y_soft))
print(grad(f)(y_soft))

```

    3.0
    [0. 0. 0.]

To use the straight-through estimator, we replace $q(y)$ with
$$y + SG(q(y)-y)$$, where SG is stop gradient. In the forwards pass, we have $y+q(y)-y=q(y)$. In the backwards pass, the gradient of SG is 0, so we effectively replace $q(y)$ with $y$. So in the backwarsd pass we have
$$
\begin{align}
f(y) &= w^T q(y) \approx w^T  y \\
\nabla_y f(y) &\approx w
\end{align}
$$

```

def f_ste(y):
  w = jnp.array([1,2,3])
  yq = quantize(y)
  yy = y + jax.lax.stop_gradient(yq - y) # gives yq on fwd, and y on backward
  return jnp.dot(w, yy)

print(f_ste(y_soft))
print(grad(f_ste)(y_soft))
```

    3.0
    [1. 2. 3.]

## Per-example gradients

In some applications, we want to compute the gradient for every example in a batch, not just the sum of gradients over the batch. This is hard in other frameworks like TF and PyTorch but easy in JAX, as we show below.

```
def loss(w, x):
  return jnp.dot(w,x)

w = jnp.ones((3,))
x0 = jnp.array([1.0, 2.0, 3.0])
x1 = 2*x0
X = jnp.stack([x0, x1])
print(X.shape)

perex_grads = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0)))
print(perex_grads(w, X))

```

    (2, 3)
    [[1. 2. 3.]
     [2. 4. 6.]]

To explain the above code in more depth, note that the vmap converts the function loss to take  a batch of inputs for each of its arguments, and returns a batch of outputs. To make it work with a single weight vector, we specify in_axes=(None,0), meaning the first argument (w) is not replicated, and the second argument (x) is replicated along dimension 0.

```
gradfn = jax.grad(loss)

W = jnp.stack([w, w])
print(jax.vmap(gradfn)(W, X))

print(jax.vmap(gradfn, in_axes=(None,0))(w, X))

```

    [[1. 2. 3.]
     [2. 4. 6.]]
    [[1. 2. 3.]
     [2. 4. 6.]]

# Optimization

The [Optax library](https://github.com/deepmind/optax) implements many common optimizers. Below is a simple example.

```
%%capture

!pip install git+git://github.com/deepmind/optax.git
import optax
```

```
num_weights = 2
params = {'w': jnp.ones((num_weights,))}
num_ex = 3
xs = 2*jnp.ones((num_ex, num_weights))
ys = jnp.ones(num_ex)
compute_loss_single = lambda params, x, y: optax.l2_loss(params['w'].dot(x), y)
compute_loss = lambda params, xs, ys: jnp.sum(jax.vmap(compute_loss_single, in_axes=[None, 0, 0])(params, xs, ys))

print('original params ', params)
print('loss ', compute_loss(params, xs, ys))

# create a stateful optimizer
learning_rate = 0.1
optimizer = optax.adam(learning_rate)
opt_state = optimizer.init(params)
print('original state ', opt_state)

# compute gradients
grads = jax.grad(compute_loss)(params, xs, ys)
print('grads ', grads)

# updage params (and optstate) given gradients
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

print('updated params ', params)
print('updated state ', opt_state)
```

    original params  {'w': DeviceArray([1., 1.], dtype=float32)}
    loss  13.5
    original state  [ScaleByAdamState(count=DeviceArray(0, dtype=int32), mu={'w': DeviceArray([0., 0.], dtype=float32)}, nu={'w': DeviceArray([0., 0.], dtype=float32)}), EmptyState()]
    grads  {'w': DeviceArray([18., 18.], dtype=float32)}
    updated params  {'w': DeviceArray([0.9, 0.9], dtype=float32)}
    updated state  [ScaleByAdamState(count=DeviceArray(1, dtype=int32), mu={'w': DeviceArray([1.8, 1.8], dtype=float32)}, nu={'w': DeviceArray([0.324, 0.324], dtype=float32)}), EmptyState()]

# JIT (just in time compilation)

In this section, we illustrate how to use the Jax JIT compiler to make code go much faster (even on a CPU). It does this by compiling the computational graph into low-level XLA primitives, potentially fusing multiple sequential operations into a single op. However, it does not work on arbitrary Python code, as we explain below.

```

def slow_f(x):
  # Element-wise ops see a large benefit from fusion
  return x * x + x * 2.0

x = jnp.ones((5000, 5000))
%timeit  slow_f(x)

fast_f = jit(slow_f)
%timeit fast_f(x)

assert np.allclose(slow_f(x), fast_f(x))
```

    100 loops, best of 5: 3.11 ms per loop
    The slowest run took 15.74 times longer than the fastest. This could mean that an intermediate result is being cached.
    1000 loops, best of 5: 918 µs per loop

We can also add the `@jit` decorator in front of a function.

```
@jit
def faster_f(x):
  return x * x + x * 2.0

%timeit faster_f(x)
assert np.allclose(faster_f(x), fast_f(x))
```

    The slowest run took 20.30 times longer than the fastest. This could mean that an intermediate result is being cached.
    1000 loops, best of 5: 918 µs per loop

## How it works: Jaxprs and tracing

In this section, we briefly explain the mechanics behind JIT, which will help you understand when it does not work.

First, consider this function.

```

def f(x):
  y = jnp.ones((1,5)) * x
  return y

```

When a function is first executed (applied to an argument), it is converted to an intermediate representatio called a JAX expression or jaxpr, by a process called tracing, as we show below.

```
print(f(3.0))
print(jax.make_jaxpr(f)(3.0))
```

    [[3. 3. 3. 3. 3.]]
    { lambda  ; a.
      let b = broadcast_in_dim[ broadcast_dimensions=(  )
                                shape=(1, 5) ] 1.0
          c = convert_element_type[ new_dtype=float32
                                    weak_type=False ] a
          d = mul b c
      in (d,) }

The XLA JIT compiler can then convert the jaxpr to code that runs fast on a CPU, GPU or TPU; the original python code is no longer needed.

```
f_jit = jit(f)
print(f_jit(3.0))
```

    [[3. 3. 3. 3. 3.]]

However, the jaxpr is created by tracing the function for a specific value. If different code is executed depending on the value of the input arguments, the resulting jaxpr will be different, so the function cannot be JITed, as we illustrate below.

```

def f(x):
  if x > 0:
    return x
  else:
    return 2 * x

print(f(3.0))

f_jit = jit(f)
print(f_jit(3.0))
```

Jit will create a new compiled version for each different ShapedArray, but will reuse the code for different values of the same shape. If the code path depends on the concrete value,  we can either just jit a subfunction (whose code path is constant), or we can create a different jaxpr for each concrete value of the input arguments as we explain below.

## Static argnum

Note that JIT compilation requires that the control flow through the function  can be determined by the shape (but not concrete value) of its inputs. The function below violates this, since when x<0, it takes one branch, whereas when x>0, it takes the other.

```
@jit
def f(x):
  if x > 0:
    return x
  else:
    return 2 * x

# This will fail!
try:
  print(f(3))
except Exception as e:
  print("ERROR:", e)

```

    ERROR: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
    The problem arose with the `bool` function.
    While tracing the function f at <ipython-input-92-94e6eda28128>:1 for jit, this concrete value was not available in Python because it depends on the value of the argument 'x'.

    See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

We can fix this by telling JAX to trace the control flow through the function using concrete values of some of its arguments. JAX will then compile different versions, depending on the input values. See below for an example.

```

def f(x):
  if x > 0:
    return x
  else:
    return 2 * x

f = jit(f, static_argnums=(0,))

print(f(3))

```

    3

```
@partial(jit, static_argnums=(0,))
def f(x):
  if x > 0:
    return x
  else:
    return 2 * x
print(f(3))
```

    3

## Jit and vmap

Unfortunately, the static argnum method fails when the function is passed to  vmap, because the latter can take arguments of different shape.

```
xs = jnp.arange(5)

@partial(jit, static_argnums=(0,))
def f(x):
  if x > 0:
    return x
  else:
    return 2 * x

ys = vmap(f)(xs)
```

## Side effects

Since the jaxpr is created only once, if your function has global side-effects, such as using print, they will only happen once, even if the function is called multiple times. See example below.

```
def f(x):
  print('x', x)
  y = 2 * x
  print('y', y)
  return y

y1 = f(2)
print('f', y1)
print('\ncall function a second time')
y1 = f(2)
print('f', y1)

print('\njit version follows')
g = jax.jit(f)
y2 = g(2)
print('f', y2)

print('\ncall jitted function a second time')
y2 = g(2)
print('f', y2)
```

    x 2
    y 4
    f 4

    call function a second time
    x 2
    y 4
    f 4

    jit version follows
    x Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
    y Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
    f 4

    call jitted function a second time
    f 4

## Caching

If you write `g=jax.jit(f)`, then f will get compiled and the XLA code will be cahced. Subsequent calls to g reuse the cached code for speed. But if the jit is called inside a loop, it is effectively making a new f each time, which is slow. So typically jit occurs in the outermost scope (modulo being constant shape).

Also, if you specify `static_argnums`, then the cached code will be used only for the same values of arguments labelled as static. If any of them change, recompilation occurs.

## Strings

Jit does not work with functions that consume or return strings.

```
def f(x: int, y: str):
  if y=='add':
    return x+1
  else:
    return x-1

print(f(42, 'add'))
print(f(42, 'sub'))

fj = jax.jit(f)
print(fj(42, 'add'))

```

# Pytrees

A Pytree is a container of leaf elements and/or more pytrees. Containers include lists, tuples, and dicts. A leaf element is anything that’s not a pytree, e.g. an array.
Pytrees are useful for representing hierarchical sets of parameters for DNNs (and other structured dsta).

## Simple example

```
from jax import tree_util

# a simple pytree
t1 = [1, {"k1": 2, "k2": (3, 4)}, 5]
print('tree', t1)
leaves = jax.tree_leaves(t1)
print('num leaves', len(leaves))
print(leaves)

t4 = [jnp.array([1, 2, 3]), "foo"]
print('tree', t4)
leaves = jax.tree_leaves(t4)
print('num leaves', len(leaves))
print(leaves)

```

    tree [1, {'k1': 2, 'k2': (3, 4)}, 5]
    num leaves 5
    [1, 2, 3, 4, 5]
    tree [DeviceArray([1, 2, 3], dtype=int32), 'foo']
    num leaves 2
    [DeviceArray([1, 2, 3], dtype=int32), 'foo']

## Treemap

We can map functions down a pytree in the same way that we can map a function down a list. We can also combine elements in two pytrees that have the same shape to make a third pytree.

```
t1 = [1, {"k1": 2, "k2": (3, 4)}, 5]
print(t1)

t2 = tree_util.tree_map(lambda x: x*x, t1)
print('square each element', t2)

t3 = tree_util.tree_map(lambda x,y: x+y, t1, t2)
print('t1+t2', t3)
```

    [1, {'k1': 2, 'k2': (3, 4)}, 5]
    square each element [1, {'k1': 4, 'k2': (9, 16)}, 25]
    t1+t2 [2, {'k1': 6, 'k2': (12, 20)}, 30]

If we have a list of dicts, we can convert to a dict of lists, as shown below.

```

data = [dict(t=1, obs='a', val=-1), dict(t=2, obs='b', val=-2), dict(t=3, obs='c', val=-3)]

data2 = jax.tree_map(lambda d0, d1, d2: list((d0, d1, d2)),
                          data[0], data[1], data[2])
print(data2)

def join_trees(list_of_trees):
  d = jax.tree_map(lambda *xs: list(xs), *list_of_trees)
  return d

print(join_trees(data))
```

    {'obs': ['a', 'b', 'c'], 't': [1, 2, 3], 'val': [-1, -2, -3]}
    {'obs': ['a', 'b', 'c'], 't': [1, 2, 3], 'val': [-1, -2, -3]}

## Flattening / Unflattening

```

t1 = [1, {"k1": 2, "k2": (3, 4)}, 5]
print(t1)
leaves, treedef = jax.tree_util.tree_flatten(t1)
print(leaves)
print(treedef)

t2 = jax.tree_util.tree_unflatten(treedef, leaves)
print(t2)
```

    [1, {'k1': 2, 'k2': (3, 4)}, 5]
    [1, 2, 3, 4, 5]
    PyTreeDef([*, {'k1': *, 'k2': (*, *)}, *])
    [1, {'k1': 2, 'k2': (3, 4)}, 5]

## Example: Linear regression

In this section we show how to use pytrees as a container for parameters of a linear reregression model.
The code is based on the [flax JAX tutorial](https://flax.readthedocs.io/en/latest/notebooks/jax_for_the_impatient.html). When we compute the gradient, it will also be a pytree, and will have the same shape as the parameters, so we can add the params to the gradient without having to flatten and unflatten the parameters.

```

# Create the predict function from a set of parameters
def make_predict_pytree(params):
  def predict(x):
    return jnp.dot(params['W'],x)+params['b']
  return predict

# Create the loss from the data points set
def make_mse_pytree(x_batched,y_batched): # returns fn(params)->real
  def mse(params):
    # Define the squared loss for a single pair (x,y)
    def squared_error(x,y):
      y_pred = make_predict_pytree(params)(x)
      return jnp.inner(y-y_pred,y-y_pred)/2.0
    # We vectorize the previous to compute the average of the loss on all samples.
    return jnp.mean(jax.vmap(squared_error)(x_batched,y_batched), axis=0)
  return jax.jit(mse) # And finally we jit the result.
```

```
# Set problem dimensions
N = 20
xdim = 10
ydim = 5

# Generate random ground truth W and b
key = random.PRNGKey(0)
Wtrue = random.normal(key, (ydim, xdim))
btrue = random.normal(key, (ydim,))
params_true = {'W': Wtrue, 'b': btrue}
true_predict_fun = make_predict_pytree(params_true)

# Generate data with additional observation noise
X = random.normal(key, (N, xdim))
Ytrue = jax.vmap(true_predict_fun)(X)
Y = Ytrue + 0.1*random.normal(key, (N, ydim))

# Generate MSE for our samples
mse_fun = make_mse_pytree(X, Y)
```

```
# Initialize estimated W and b with zeros.
params = {'W': jnp.zeros_like(Wtrue), 'b': jnp.zeros_like(btrue)}

mse_pytree = make_mse_pytree(X, Y)
print(mse_pytree(params_true))
print(mse_pytree(params))
print(jax.grad(mse_pytree)(params))
```

    0.022292053
    24.97824
    {'W': DeviceArray([[-0.039,  0.755,  0.542,  0.36 ,  0.224,  1.651,  1.534,
                  -1.342, -0.15 , -1.638],
                 [-0.324,  0.141, -0.402,  0.498,  1.829,  4.308,  2.138,
                  -2.43 , -0.381, -2.178],
                 [ 1.7  , -0.707, -0.656, -0.568,  1.824, -2.194, -0.477,
                   0.96 ,  1.622,  1.408],
                 [-0.862,  0.321, -0.388, -0.74 , -0.82 ,  0.441,  0.772,
                  -1.713, -1.592, -0.557],
                 [ 1.338, -0.632, -0.968, -1.127,  1.775,  0.323,  1.405,
                  -0.638,  1.077, -0.739]], dtype=float32), 'b': DeviceArray([ 0.036,  1.092, -0.413, -1.389, -0.862], dtype=float32)}

```
alpha = 0.3 # Gradient step size
print('Loss for "true" W,b: ', mse_pytree(params_true))
for i in range(101):
  gradients = jax.grad(mse_pytree)(params)
  params = jax.tree_map(lambda old,grad: old-alpha*grad, params, gradients)
  if (i%10==0):
    print("Loss step {}: ".format(i), mse_pytree(params))

```

    Loss for "true" W,b:  0.022292053
    Loss step 0:  6.5597453
    Loss step 10:  0.17232798
    Loss step 20:  0.04339735
    Loss step 30:  0.024473602
    Loss step 40:  0.017078908
    Loss step 50:  0.013489487
    Loss step 60:  0.011695375
    Loss step 70:  0.010795272
    Loss step 80:  0.010343453
    Loss step 90:  0.01011666
    Loss step 100:  0.010002809

```
print(jax.tree_map(lambda x,y: np.allclose(x,y, atol=1e-1), params, params_true))
```

    {'W': True, 'b': True}

Compare the above to what the training code would look like
if W and b were passed in as separate arguments:
```
for i in range(101):
  grad_W = jax.grad(mse_fun,0)(What,bhat)
  grad_b = jax.grad(mse_fun,1)(What,bhat)
  What = What - alpha*grad_W
  bhat = bhat - alpha*grad_b
  if (i%10==0):
    print("Loss step {}: ".format(i), mse_fun(What,bhat)
```

## Example: MLPs

We now show a more interesting example, from the Deepmind tutorial, where we fit an MLP using SGD. The basic structure is similar to the linear regression case.

```
# define the model
def init_mlp_params(layer_widths):
  params = []
  for n_in, n_out in zip(layer_widths[:-1], layer_widths[1:]):
    params.append(
        dict(weights=np.random.normal(size=(n_in, n_out)) * np.sqrt(2/n_in),
             biases=np.ones(shape=(n_out,))
            )
    )
  return params

def forward(params, x):
  *hidden, last = params
  for layer in hidden:
    x = jax.nn.relu(x @ layer['weights'] + layer['biases'])
  return x @ last['weights'] + last['biases']

def loss_fn(params, x, y):
  return jnp.mean((forward(params, x) - y) ** 2)
```

```
# MLP with 2 hidden layers and linear output
np.random.seed(0)
params = init_mlp_params([1, 128, 128, 1])
jax.tree_map(lambda x: x.shape, params)

```

    [{'biases': (128,), 'weights': (1, 128)},
     {'biases': (128,), 'weights': (128, 128)},
     {'biases': (1,), 'weights': (128, 1)}]

```
LEARNING_RATE = 0.0001

@jax.jit
def update(params, x, y):
  grads = jax.grad(loss_fn)(params, x, y)
  return jax.tree_map(
      lambda p, g: p - LEARNING_RATE * g, params, grads)

np.random.seed(0)
xs = np.random.normal(size=(200, 1))
ys = xs ** 2

for _ in range(1000):
  params = update(params, xs, ys)

plt.scatter(xs, ys, label='truth')
plt.scatter(xs, forward(params, xs), label='Prediction')
plt.legend()
```

# Looping constructs

For loops in Python are slow, even when JIT-compiled. However, there are built-in primitives for loops that are fast, as we illustrate below.

## For loops.

The semantics of the for loop function in JAX is as follows:
```
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
```
We see that ```val``` is used to accumulate the results across iterations.

Below is an example.

```
# sum from 1 to N = N*(N+1)/2

def sum_exact(N):
  return int(N*(N+1)/2)

def sum_slow(N):
  s = 0
  for i in range(1,N+1):
    s += i
  return s

N = 10

assert sum_slow(N) == sum_exact(N)

def sum_fast(N):
  s = jax.lax.fori_loop(1, N+1, lambda i,partial_sum: i+partial_sum, 0)
  return s

assert sum_fast(N) == sum_exact(N)
```

```
N = 1000
%timeit sum_slow(N)
%timeit sum_fast(N)
```

    10000 loops, best of 5: 54.8 µs per loop
    10 loops, best of 5: 21.9 ms per loop

```
# Let's do more compute per step of the for loop

D = 10
X = jax.random.normal(key, shape=(D,D))

def sum_slow(N):
  s = jnp.zeros_like(X)
  for i in range(1,N+1):
    s += jnp.dot(X, X)
  return s

def sum_fast(N):
  s = jnp.zeros_like(X)
  s = jax.lax.fori_loop(1, N+1, lambda i,s: s+jnp.dot(X,X), s)
  return s

N = 10
assert np.allclose(sum_fast(N), sum_slow(N))
```

```
N = 1000
%timeit sum_slow(N)
%timeit sum_fast(N)
```

    1 loop, best of 5: 298 ms per loop
    10 loops, best of 5: 27.8 ms per loop

## While loops

Here is the semantics of the JAX while loop

```
def while_loop(cond_fun, body_fun, init_val):
  val = init_val
  while cond_fun(val):
    val = body_fun(val)
  return val
```

Below is an example.

```

def sum_slow_while(N):
  s = 0
  i = 0
  while (i <= N):
    s += i
    i += 1
  return s

def sum_fast_while(N):
  init_val = (0,0)
  def cond_fun(val):
    s,i = val
    return i<=N
  def body_fun(val):
    s,i = val
    s += i
    i += 1
    return (s,i)
  val = jax.lax.while_loop(cond_fun, body_fun, init_val)
  s2 = val[0]
  return s2

N = 10
assert sum_slow_while(N) == sum_exact(N)
assert sum_slow_while(N) == sum_fast_while(N)
```

```
N = 1000
%timeit sum_slow(N)
%timeit sum_fast(N)
```

    1 loop, best of 5: 312 ms per loop
    10 loops, best of 5: 28.3 ms per loop

## Scan

Here is the semantics of scan:
```
def scan(f, init, xs, length=None):
  if xs is None:
    xs = [None] * length
  carry = init
  ys = []
  for x in xs:
    carry, y = f(carry, x)
    ys.append(y)
  return carry, np.stack(ys)
  ```

Here is an example where we use scan to sample from a discrete-time, discrete-state Markov chain.

```
init_dist = jnp.array([0.8, 0.2])
trans_mat = jnp.array([[0.9, 0.1], [0.5, 0.5]])
rng_key = jax.random.PRNGKey(0)
from jax.scipy.special import logit
seq_len = 15

initial_state = jax.random.categorical(rng_key, logits=logit(init_dist), shape=(1,))

def draw_state(prev_state, key):
    logits = logit(trans_mat[:, prev_state])
    state = jax.random.categorical(key, logits=logits.flatten(), shape=(1,))
    return state, state

rng_key, rng_state, rng_obs = jax.random.split(rng_key, 3)
keys = jax.random.split(rng_state, seq_len - 1)

final_state, states = jax.lax.scan(draw_state, initial_state, keys)

print(states)
```

    [[0]
     [0]
     [0]
     [0]
     [0]
     [0]
     [0]
     [1]
     [1]
     [1]
     [1]
     [1]
     [1]
     [1]]

# Common gotchas

## Handling state

In this section, we discuss how to transform code that uses object-oriented programming (which can be stateful) to pure functional programming, which is stateless, as required by JAX. Our presentation is based on the Deepmind tutorial.

To start, consider a simple class that maintains an internal counter, and when called, increments the counter and returns the next number from some sequence.

```
#import string
#DICTIONARY = list(string.ascii_lowercase)
SEQUENCE = jnp.arange(0,100,2)

class Counter:

  def __init__(self):
    self.n = 0

  def count(self) -> int:
    #res = DICTIONARY[self.n]
    res = SEQUENCE[self.n]
    self.n += 1
    return res

  def reset(self):
    self.n = 0

counter = Counter()

for _ in range(3):
  print(counter.count())
```

    0
    2
    4

The trouble with the above code is that the call to `count` depends on the internal state of the object (the value `n`), even though this is not an argument to the function. (The code is therefoe said to violate 'referential transparency'.) When we Jit compile it, Jax will only call the code once (to convert to a jaxpr), so the side effect of updating `n` will not happen, resulting in incorrect behavior, as we show below,

```
counter.reset()
fast_count = jax.jit(counter.count)

for _ in range(3):
  print(fast_count())
```

    0
    0
    0

We can solve this problem by passing the state as an argument into the function.

```
CounterState = int
Result = int

class CounterV2:

  def count(self, n: CounterState) -> Tuple[Result, CounterState]:
    return SEQUENCE[n], n+1

  def reset(self) -> CounterState:
    return 0

counter = CounterV2()
state = counter.reset()

for _ in range(3):
  value, state = counter.count(state)
  print(value)
```

    0
    2
    4

This version is functionally pure, so jit-compiles nicely.

```
state = counter.reset()
fast_count = jax.jit(counter.count)

for _ in range(3):
  value, state = fast_count(state)
  print(value)
```

    0
    2
    4

We can apply the same process to any stateful method to convert it into a stateless one. We took a class of the form

```
class StatefulClass

  state: State

  def stateful_method(*args, **kwargs) -> Output:
```

and turned it into a class of the form

```
class StatelessClass

  def stateless_method(state: State, *args, **kwargs) -> (Output, State):
```

This is a common [functional programming](https://en.wikipedia.org/wiki/Functional_programming) pattern, and, essentially, is the way that state is handled in all JAX programs (as we saw with the way Jax handles random number state, or parameters of a model that get updated).
Note that the stateless version of the code no longer needs to use a class, but can instead group the functions into a common namespace using modules.

In some cases (eg when working with DNNs), it is more convenient to write code in an OO way. There are several libraries (notably [Flax](https://github.com/google/flax) and [Haiku](https://github.com/deepmind/dm-haiku))  that let you define a model in an OO way, and then generate functionally pure code.

## Mutation of arrays

Since JAX is functional, you cannot mutate arrays in place,
since this makes program analysis and transformation very difficult. JAX requires a pure functional expression of a numerical program.
Instead, JAX offers the functional update functions: `index_update`, `index_add`, `index_min`, `index_max`, and the `index` helper. These are illustrated below. However it is best to avoid these if possible, since they are slow.

Note: If the input values of `index_update` aren't reused, jit-compiled code will perform these operations in-place, rather than making a copy.

```
# You cannot assign directly to elements of an array.

A = jnp.zeros((3,3), dtype=np.float32)

# In place update of JAX's array will yield an error!
try:
  A[1, :] = 1.0
except:
  print('must use index_update')
```

    must use index_update

```
from jax.ops import index, index_add, index_update

D = 3
A = 2*jnp.ones((D,D))
print("original array:")
print(A)

A2 = index_update(A, index[1, :], 42.0) # A[1,:] = 42
print("original array:")
print(A) # unchanged
print("new array:")
print(A2)

A3 = A.at[1,:].set(42.0) # A3=np.copy(A),  A3[1,:] = 42
print("original array:")
print(A) # unchanged
print("new array:")
print(A3)

A4 = A.at[1,:].mul(42.0) # A4=np.copy(A),  A4[1,:] *= 42
print("original array:")
print(A) # unchanged
print("new array:")
print(A4)

```

    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    new array:
    [[ 2.  2.  2.]
     [42. 42. 42.]
     [ 2.  2.  2.]]
    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    new array:
    [[ 2.  2.  2.]
     [42. 42. 42.]
     [ 2.  2.  2.]]
    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    new array:
    [[ 2.  2.  2.]
     [84. 84. 84.]
     [ 2.  2.  2.]]

## Implicitly casting lists to vectors

You cannot treat a list of numbers as a vector. Instead you must explicitly create the vector using the np.array() constructor.

```
# You cannot treat a list of numbers as a vector.
try:
  S = jnp.diag([1.0, 2.0, 3.0])
except:
  print('must convert indices to np.array')
```

    must convert indices to np.array

```
# Instead you should explicitly construct the vector.

S = jnp.diag(jnp.array([1.0, 2.0, 3.0]))
```


AGILE TEST LAB
         --- coded by Xiaoyuan Cheng

```python
import cvxpy as cp
import functools
import numpy as np
import pandas as pd
import datetime as DT
import matplotlib.pyplot as plt
import math
import pickle
import re
import sympy as sy
from scipy.optimize import minimize
import time
from itertools import islice
```

study framework

day-ahead price mechinism

```python
#function wrapper
def Xiaoyuan(func):
    @functools.wraps(func)
    def wrapper(*arg,**kw):
        time_1 = time.time()
        r = func(*arg, **kw)
        print('{} consumes {}ms'.format(func.__name__, 1000*(time.time()-time_1)))
        return r
    return wrapper
```

```python
#define the real-time price of main grid electricity price
#time interval is 15 mins
time = [i/4 for i in range(24*4)]
grid_price = np.random.uniform(0.8,1,24*4)
fig = plt.plot(time,grid_price)
plt.xlabel('Time, hour')
plt.ylabel('The grid electricity price in one day, kw/$')
plt.show()
```

EV state mechanism

```python
#define a case study of EVs behaviors in a time period
#here we use a discrete markov chain to simulate it, there is three states; charge, off-line and discharge
#use 'charging', 'offline', 'discharge' indicated the
transition_matrix = np.array([[0.4,0.4,0.2],
                 [0.25,0.4,0.35],
                 [0.3,0.35,0.35]])

#you can check the state of each EV, I labeled each EV NO. and the start state
EV_No = [i for i in range(1,301)]
EV_start_state  =np.random.choice(['charging','offline','discharging'],300,[0.33,0.34,0.33])

#use an example to show you the case,just choose first 10 EVs
#x-bar is the No. of EV
x = EV_No[0:10]
y=[i for i in range(10)]
plt.scatter(x, y)
for i in range(len(x)):
    plt.annotate(EV_start_state[i], xy = (x[i], y[i]), xytext = (x[i]+0.1, y[i]+0.1))

plt.xlabel('EV No. in start state')
plt.show()
```

```python
#define a function to describe the EV state in next time interval
@Xiaoyuan
def EV_state(EV_current_state, transition_matrix):
    next_EV_state = []
    for state in EV_current_state:
        if state == 'charging':
            next_state = np.random.choice(['charging','offline','discharging'], p = transition_matrix[0].ravel())
        elif state == 'offline':
            next_state = np.random.choice(['charging','offline','discharging'], p = transition_matrix[1].ravel())
        else:
            next_state = np.random.choice(['charging','offline','discharging'], p = transition_matrix[2].ravel())
        pass
        next_EV_state.append(next_state)
    return next_EV_state

#based on this function, simulate the dynamics of EVs' state in one day
if __name__ == '__main__':
    EV_state_matrix = []
    EV_current_state = EV_start_state
    # EV_state_matrix.append(EV_current_state)
    for i in range(4*24-1):
        EV_next_state = EV_state(EV_current_state, transition_matrix)
        EV_current_state = EV_next_state
        EV_state_matrix.append(EV_current_state)
        pass
```

```python
# EV states show
# np.array(EV_state_matrix).shape
# i = 0
# while i<=94:
#     print('in time {}  charging EVs: {}, offline EVs: {}, discharging EVs: {}'.format(i*0.25+0.25,
#                                                                                EV_state_matrix[i].count('charging'),
#                                                                                EV_state_matrix[i].count('offline'),
#                                                                                EV_state_matrix[i].count('discharging')))
charging = [EV_state_matrix[i].count('charging') for i in range(95)]
discharging = [EV_state_matrix[i].count('discharging') for i in range(95)]
plt.plot([i*0.25+0.25 for i in range(95)], [EV_state_matrix[i].count('charging') for i in range(95)])
plt.plot([i*0.25+0.25 for i in range(95)], [EV_state_matrix[i].count('offline') for i in range(95)])
plt.plot([i*0.25+0.25 for i in range(95)], [EV_state_matrix[i].count('discharging') for i in range(95)])
plt.legend(['charging', 'offline', 'diischarging'], loc='upper right')
plt.xlabel('Time, hour')
plt.ylabel('EVs number')
plt.show()
```

```python
#define the aggregator characteristics
#in this case study, we assume there are two aggregators
#if it is feasible, then we extend to a larger scale

#in a micro-grid, we define class type, each of aggregator will follows the defined class
class aggregator(object):
    start_storage = 0

    def __init__(self,input_power,output_power,storage_limitation,start_storage,t):
        self.input = input_power #charging power of EVs to aggregator
        self.output = output_power # required power of EVs
        self.storage_limitation = storage_limitation # limitation of aggregator SOC
        # aggregator.start_storage = start_storage # aggregator SOC
        self.t = t #time stamp
        if (self.input-self.output)>0 and self.start_storage < self.storage_limitation and (self.start_storage+0.25*(self.input-self.output))< self.storage_limitation:
            aggregator.start_storage += 0.25*(self.input-self.output)
            self.mode = 1 #represent the current behaviors of aggregator
            print('aggregator is charging')
        elif (self.input-self.output)>0 and self.start_storage < self.storage_limitation and (self.start_storage+0.25*(self.input-self.output))> self.storage_limitation:
            aggregator.start_storage = storage_limitation
            self.mode = 2
            print('aggregator is full after charging')
        elif (self.input-self.output)<0 and (self.start_storage - 0.25*(self.input-self.output))>0:
            aggregator.start_storage += -0.25*(self.input-self.output)
            self.mode = 3
            print('aggregator is discharging')
        elif (self.input-self.output)<0 and (self.start_storage - 0.25*(self.input-self.output))<0:
            aggregator.start_storage = 0
            self.mode = 4
            print('aggregator is purchasing electricity from main grid')
        else:
            print('errors, please check')
            pass

    def next_time_state(self):
        return self.input,self.output,self.storage_limitation,aggregator.start_storage,self.t,self.mode
    # def profit(self, offer_price, real_makret):
    #     if self.mode == 1:
    #         profit_state =

```

here we have done a lot of jobs for preparing the boundaries, this part will enter into game theory to simulate the competitiveness among two aggregators

```python
# we assume that the EVs are sensitive to historical trading price, here we use a relatively simple function to desribe EVs when
#they choosing aggregators>> In the future, more complex functions can be employed, such as RNN, LSTM etc.

#let's define two prices
#day_ahead_price, means the price EVs selling electricity to aggregators, our target is to optimize this value
#real_market_price, means price aggregators selling to EVs, which is same to grid price

#then, define how many people will choose to different aggregators

@Xiaoyuan
def EVs_num_aggregator_n(aggregator_1_current_price, aggregator_2_current_price,current_grid_price,t):
    num_1 = (current_grid_price - aggregator_1_current_price)/(2*current_grid_price - aggregator_2_current_price -aggregator_1_current_price)
    num_2 = (current_grid_price - aggregator_2_current_price)/(2*current_grid_price - aggregator_2_current_price -aggregator_1_current_price)
    return num_1, num_2
```

```python
#liz mentioned the contract price should be set in a reasonable range
#in each time series, the day-ahead price should not beyond real-time price x in [0.6, current_grid_price)

#the detailed game theory process has been defined in my report, let's finished the last technical part

#for every step, each aggregator will try to find optimal price to maximize their profits, the function is:

# @Xiaoyuan

def fun(args):
    b,g,dis1,cha1,t = args
    v = lambda a: (g-a[0])/(2*g-a[0]-b)*(cha1*g-dis1*a[0])
    return v

def con(args1):
    x1min, x1max= args1
    cons = ({'type': 'ineq', 'fun': lambda x: x[0] - x1min},\
            {'type': 'ineq', 'fun': lambda x: -x[0] + x1max})
    return cons
```

```python
# let's start the game,
aggregator = [0.7]

args = (0.7,grid_price[1],discharging[0],charging[0],i)
args1 = (0.6,grid_price[1]*1.2)
x0 = np.asarray([0.82])
res = minimize(fun(args), x0, method='SLSQP',constraints=con(args1))
aggregator.append(res.x[0])

for i in range(1,95):
    args = (res.x[0],grid_price[i+1],discharging[i],charging[i],i%2)
    args1 = (0.6,grid_price[i]*1.2)
    x0 = np.asarray([0.82])
    res = minimize(fun(args), x0, method='SLSQP',constraints=con(args1))
    # print(res.fun)
    # print(res.success)
    # print(res.x)
    aggregator.append(res.x[0])
```

```python
#the day-ahead price of aggregator is:
plt.plot([i*0.5 for i in range(48)], [aggregator[i*2] for i in range(48)])
plt.plot([i*0.5 for i in range(48)], [aggregator[i*2+1] for i in range(48)])
plt.plot([i*0.5 for i in range(48)], [grid_price[i*2] for i in range(48)])
plt.legend(['aggregator 1 - day ahead', 'aggregator 2 - day ahead', 'real-time price'], loc='upper right')
plt.xlabel('Time, hour')
plt.ylabel('electricity price in one day, kw/$')
plt.show()
```

please note, the orange and blue line is the price we wanted


# Lecture 3.6: Logistic Regression

[**Lecture Slides**](https://docs.google.com/presentation/d/1gk8b6M08T_D8BNbOts2CYkD4qxOeQVWefx7pWTt-YjI/edit?usp=sharing)

This lecture, we are going to use logistic regression to predict the quality of wines.

**Learning goals:**

- solve a classification problem
- fit a logistic regression model
- predict using a logistic regression model
- visualize the decision boundary of a logistic regression model

## 1. Introduction

Humans have been making wine for more than [8000 years](https://www.nationalgeographic.com/travel/destinations/asia/georgia/sponsor-content-secret-birthplace-of-wine/),  and that's a lot of bottles. 🍷 There's too many options! Let's use data science to help choose the _tastiest_ red wines. Or at least, to classify high rated wines vs low rated wines based on some chemical features.

## 2. Data Munging

Let's load the dataset into a `DataFrame`:

```python
import pandas as pd

df = pd.read_csv('wine_quality.csv')
df.head()
```

      fixed acidity
      volatile acidity
      citric acid
      residual sugar
      chlorides
      free sulfur dioxide
      total sulfur dioxide
      density
      pH
      sulphates
      alcohol
      tasty

      0
      -0.528360
      0.961877
      -1.391472
      -0.453218
      -0.243707
      -0.466193
      -0.379133
      0.558274
      1.288643
      -0.579207
      -0.960246
      0

      1
      -0.298547
      1.967442
      -1.391472
      0.043416
      0.223875
      0.872638
      0.624363
      0.028261
      -0.719933
      0.128950
      -0.584777
      0

      2
      -0.298547
      1.297065
      -1.186070
      -0.169427
      0.096353
      -0.083669
      0.229047
      0.134264
      -0.331177
      -0.048089
      -0.584777
      0

      3
      1.654856
      -1.384443
      1.484154
      -0.453218
      -0.264960
      0.107592
      0.411500
      0.664277
      -0.979104
      -0.461180
      -0.584777
      1

      4
      -0.528360
      0.961877
      -1.391472
      -0.453218
      -0.243707
      -0.466193
      -0.379133
      0.558274
      1.288643
      -0.579207
      -0.960246
      0

```python
df.describe()
```

      fixed acidity
      volatile acidity
      citric acid
      residual sugar
      chlorides
      free sulfur dioxide
      total sulfur dioxide
      density
      pH
      sulphates
      alcohol
      tasty

      count
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1.599000e+03
      1599.000000

      mean
      3.554936e-16
      1.866341e-16
      -7.109871e-17
      -1.244227e-16
      3.910429e-16
      -4.443669e-17
      3.554936e-17
      -3.469617e-14
      2.843948e-15
      6.754377e-16
      5.332403e-17
      0.534709

      std
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      1.000313e+00
      0.498950

      min
      -2.137045e+00
      -2.278280e+00
      -1.391472e+00
      -1.162696e+00
      -1.603945e+00
      -1.422500e+00
      -1.230584e+00
      -3.538731e+00
      -3.700401e+00
      -1.936507e+00
      -1.898919e+00
      0.000000

      25%
      -7.007187e-01
      -7.699311e-01
      -9.293181e-01
      -4.532184e-01
      -3.712290e-01
      -8.487156e-01
      -7.440403e-01
      -6.077557e-01
      -6.551405e-01
      -6.382196e-01
      -8.663789e-01
      0.000000

      50%
      -2.410944e-01
      -4.368911e-02
      -5.636026e-02
      -2.403750e-01
      -1.799455e-01
      -1.793002e-01
      -2.574968e-01
      1.760083e-03
      -7.212705e-03
      -2.251281e-01
      -2.093081e-01
      1.000000

      75%
      5.057952e-01
      6.266881e-01
      7.652471e-01
      4.341614e-02
      5.384542e-02
      4.901152e-01
      4.723184e-01
      5.768249e-01
      5.759223e-01
      4.240158e-01
      6.354971e-01
      1.000000

      max
      4.355149e+00
      5.877976e+00
      3.743574e+00
      9.195681e+00
      1.112703e+01
      5.367284e+00
      7.375154e+00
      3.680055e+00
      4.528282e+00
      7.918677e+00
      4.202453e+00
      1.000000

Each row of our dataset is a wine.

The `tasty` column indicates if a wine is "tasty" or not. i.e if it has a high rating. Even if it is represented by an integer, it is a _categorical_ column, as values are either 0 or 1. This will be the _label_ we are aiming to _predict_ in this _classification_ task.

The other columns are numerical features. Their means are close to 0, which suggests that this dataset was normalized, like for lecture 3.5.

To keep things simple and easy to visualize, let's start by using only two features to predict the "tastiness" of these wines: `sulphates` and `alcohol`.

```python
X = df[['sulphates', 'alcohol']]
y = df['tasty']
X.shape
```

    (1599, 2)

Since our features are already a _matrix_, we do not need to convert our `X` and `y` to numpy arrays and reshape them. Which means we're ready for training!

## 3. Training

We'll use sklearn's [`LogisticRegression`](https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression) model which, as always, will make things simple with the `.fit()` and `.predict()` methods.

```python
from sklearn.linear_model import LogisticRegression

clf = LogisticRegression(random_state=0).fit(X, y)
```

🧠 Can you recall all the optimization steps that sklearn takes care of in the `.fit()` method? Refer to lectures 3.4 & 3.5 if this isn't clear.

Just like for linear regression, we can investigate the _parameters_ of our model, $\theta$.
They won't be _interpreted_ in the same way however, since our hypothesis uses the logistic function. Nevertheless, sklearn still calls these parameters `intercept_` and `coef_`.

```python
import numpy as np

theta = np.append(clf.intercept_, clf.coef_)
theta
```

    array([0.24435503, 0.4510355 , 1.09938009])

This means that the optimization procedure found that the logistic hypothesis which best predits the `tasty` labels is:

$$\begin{align}
h_{\theta}(x_{1}, x_{2}) &= \frac{1}{1 + e^{-(\theta_{0} + \theta_{1} x_{1} + \theta_{2} x_{2})}} \\
h_{\theta}(x_{1}, x_{2}) &= \frac{1}{1 + e^{-(0.244 + 0.451 x_{1} + 1.099 x_{2})}}
\end{align}$$

where $x_{1}$ is our first feature, the `sulphates` column, and $x_{2}$ is our second feature, the `alcohol` column.

This is a function with two variables that we can plot in 3D:

```python
%matplotlib notebook
import matplotlib.pyplot as plt
import seaborn as sns
from mpl_toolkits import mplot3d
import math

sns.set()

def hypothesis(x1, x2):
    linear_term = theta[0] + theta[1] * x1 + theta[2] * x2
    return 1 / (1 + math.exp(-linear_term))

v_hypothesis = np.vectorize(hypothesis)

def create_mesh(theta):
    x_values = np.linspace(-8, 8, 40)
    y_values = np.linspace(-8, 8, 40)
    x_mesh, y_mesh = np.meshgrid(x_values, y_values)
    z_mesh = v_hypothesis(x_mesh, y_mesh)
    return x_mesh, y_mesh, z_mesh

x_mesh, y_mesh, z_mesh = create_mesh(theta)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

ax.plot_wireframe(x_mesh, y_mesh, z_mesh, rstride=5, cstride=5, alpha=0.4, color='g')

ax.set_title('Logistic Regression Model of Wine Tastiness')
ax.set_xlabel('Sulphates (norm)')
ax.set_ylabel('Alcohol (norm)')
ax.set_zlabel('Probability of Being Tasty');

```

    <IPython.core.display.Javascript object>

🧠 Is this the shape you expected? Can you relate this to the 1D plots of the logistic function we saw in the slides?

We could add our data points to this 3D graph. However, things would get crowded, and _encoding_ a binary label on the z-axis of a 3D plot isn't very clear. We end up with some points near the "roof" of the volume, and others near the "floor", without a good way of comparing them.

Instead, we can use _color_ to encode the labels of the data points. We'll visualize the wines with the `tasty=1` label in red, and the non `tasty` wines in blue. In short, we are using a _2D graph_, where:
- on the x-axis, the `sulphates` feature
- on the y-axis, the `alcohol` feature
- as color, the `tasty` binary label

```python
%matplotlib inline

colors = sns.color_palette('husl').as_hex()

fig = plt.figure(dpi=120)
ax = fig.add_subplot(111)

# plot points
positives = X[y == 1]
negatives = X[y == 0]
ax.scatter(positives.iloc[:, 0], positives.iloc[:, 1], alpha=0.4, s=8,c=colors[0], label='tasty')
ax.scatter(negatives.iloc[:, 0], negatives.iloc[:, 1], alpha=0.4, s=8, c=colors[4], label='not tasty')

# formatting
ax.set_ylim((-2, 4))
ax.set_xlim((-2, 4))
ax.set_xlabel('Sulphates (norm)')
ax.set_ylabel('Alcohol (norm)')
ax.legend()
ax.set_title('Wine Tastiness vs Sulphates & Alcohol Contents');

```

😨 This data is messy! This is a good example of what _real world data_ looks like. No clean clusters, no clean cut classes. How easy it is to split the dataset into its categorical labels is often called the _separability_ of the data. A dataset that is _separable_ has clear optimal decision boundaries that would lead to few prediction errors. In this case, it's not even clear to a human where to draw the line between `tasty` and not `tasty`!

Let's look at our logistic regression model's attempt at separating the labels... maybe it has tasted more bottles than us 😏. Here, the prediction probabilities we visualized in the 3D graph isn't really what we wish to see. We would like to compare the _binary predictions_ against the binary labels.

For this, we have to show the _decision boundary_ of our model. But all we have are the model parameters, $\theta$ ! How do we go from the shape of the _hypothesis_ to the shape of the _decision boundary_ ? 🤨

## 4. Decision Boundary

Remember we defined how to go from prediction probability to binary prediction in the slides. We choose a threshold such that:

if $h_{\theta}(\textbf{x}) >= 0.5$ :
    _predict 1_
else:
    _predict 0_

The decision boundary is the set of points where predictions "switch" from 0 to 1. This means we can represent the it as all the points where:

$$
h_{\theta}(x_{1}, x_{2}) = 0.5 \label{eq:threshold} \tag{1}
$$

However, this doesn't help us much with plotting it in matplotlib... Ideally, we expect a line equation of the form:

$$
y = b + ax
$$

And since we're using a 2D plot where our feature `sulphates` ($x_{1}$) is on the x-axis, and our feature `alcohol` ($x_{2}$) is on the y-axis, we can rewrite the above as:

$$
x_{2} = b + ax_{1}
$$

So this is our goal. Now, let's use a bit of mathematical trickery to get there. 🧙♀
We can substitute the logistic hypothesis in equation $\eqref{eq:threshold}$:

$$
\begin{align}
\frac{1}{1 + e^{-(\theta_{0} + \theta_{1} x_{1} + \theta_{2} x_{2})}} &= 0.5 \label{eq:logistic} \tag{2} \\
\end{align}
$$

You might recall from the [plot](https://upload.wikimedia.org/wikipedia/commons/8/88/Logistic-curve.svg) of the logistic function that it is equal to $0.5$ when its input is $0$.
So we can rewrite equation $\eqref{eq:logistic}$ as:

$$
\begin{align}
\theta_{0} + \theta_{1} x_{1} + \theta_{2} x_{2} &= 0 \\
\end{align}
$$

Finally, we can re-arrange the terms to get the line equation:

$$
\begin{align}
\theta_{0} + \theta_{1} x_{1} + \theta_{2} x_{2} &= 0 \\
x_{2} &= \frac{-(\theta_{0} + \theta_{1} x_{1})}{\theta_{2}} \\
x_{2} &= \frac{-\theta_{0}}{\theta_{2}} + \frac{ \theta_{1}}{\theta_{2}}x_{1} \\
\end{align}
$$

This is the equation of a _line_ with intercept $\frac{-\theta_{0}}{\theta_{2}}$ and slope $\frac{ \theta_{1}}{\theta_{2}}$. We can plot this in our 2D graph! 🎊

We'll add our logistic regression model's decision boundary as a green line:

```python
%matplotlib inline

colors = sns.color_palette('husl').as_hex()

fig = plt.figure(dpi=120)
ax = fig.add_subplot(111)

# plot points
positives = X[y == 1]
negatives = X[y == 0]
ax.scatter(positives.iloc[:, 0], positives.iloc[:, 1], alpha=0.4, s=8,c=colors[0], label='tasty')
ax.scatter(negatives.iloc[:, 0], negatives.iloc[:, 1], alpha=0.4, s=8, c=colors[4], label='not tasty')

def decision_boundary(x1, theta):
    return -(x1 * theta[1] + theta[0])/theta[2]

# plot decision boundary
x1 = np.array(ax.get_xlim())
x2 = decision_boundary(x1, theta)
ax.plot(x1, x2, '--', c=colors[2]);

# formatting
ax.set_ylim((-2, 4))
ax.set_xlim((-2, 4))
ax.set_xlabel('Sulphates (norm)')
ax.set_ylabel('Alcohol (norm)')
ax.legend()
ax.set_title('Wine Tastiness vs Sulphates & Alcohol Contents');

```

The green line is all the points where $h_{\theta}(x_{1}, x_{2}) = 0.5$. The points _above_ the line are the model's _positive_ predictions, i.e it thinks those wines are `tasty`, whilst the points _below_ the line are the model's _negative_ predictions, i.e it thinks those wines are not so `tasty`. Take you time to understand this graph, as it is typical of classification tasks.

ℹ Notice that the predictions don't always match the true _label_ of the data. For example, take the blue point at $(-1.75, 2.4)$. The blue color shows that the label is 0, i.e a _not_ `tasty` wine. However, it is located _above_ the decision boundary, so the model think it is `tasty`. It's a classification error! Based on the _separability_ of the dataset, we can understand that the model makes a few mistakes... the goal of supervised learning algorithms isn't to perfectly predict a variable, but to find the model parameters that lead to the best _possible_ predictions.

🧠 Can you match this 2D graph with the 3D representation of the prediction probabilities of section 3? Where would the decision boundary be on the 3D graph?

## 5. Prediction

We are happy with our model, and want to try it out as soon as possible. During a visit to the supermarket, we find a strange looking wine bottle. It piques our curiosity, so we read the label: `sulphates=1.942` and `alcohol=-0.343` (in normalized units). Let's figure out if our model thinks this is a tasty bottle!

🧠 Can you graphically read the model prediction for this strange wine using the graph above?

We know by now that sklearn's `.predict()` method expects a matrix. Since we have only one example with two features, $x_{1}$ and $x_{2}$, it will be automatically instantiated as a 2D vector. Instead, let's reshape it to a ($1 \times 2$) `ndarray` matrix:

```python
sulphates = 1.942
alcohol = -0.343
x_strange_wine = np.asarray([sulphates, alcohol]).reshape(1, 2)
# Or np.asarray([[sulphates, alcohol])
```

We can then predict the wine's tastiness using the logistic regression model that we trained earlier:

```python
y_strange_wine = clf.predict(x_strange_wine)
y_strange_wine
```

    array([1])

The prediction is positive, $1$, meaning our model thinks that this wine is `tasty`. Hopefully this means it doesn't taste as strange as it looks!

We can compare this bottle to the rest of our dataset, by plotting it on the 2D graph from section 4:

```python
%matplotlib inline

colors = sns.color_palette('husl').as_hex()

fig = plt.figure(dpi=120)
ax = fig.add_subplot(111)

# plot points
positives = X[y == 1]
negatives = X[y == 0]
ax.scatter(positives.iloc[:, 0], positives.iloc[:, 1], alpha=0.4, s=8,c=colors[0], label='tasty')
ax.scatter(negatives.iloc[:, 0], negatives.iloc[:, 1], alpha=0.4, s=8, c=colors[4], label='not tasty')

def decision_boundary(x1, theta):
    return -(x1 * theta[1] + theta[0])/theta[2]

# plot decision boundary
x1 = np.array(ax.get_xlim())
x2 = decision_boundary(x1, theta)
ax.plot(x1, x2, '--', c=colors[2]);

# plot strange wine
ax.plot(x_strange_wine[:, 0], x_strange_wine[:, 1], marker='x', markersize=20, markeredgewidth=4, color='r')

# formatting
ax.set_ylim((-2, 4))
ax.set_xlim((-2, 4))
ax.set_xlabel('Sulphates (norm)')
ax.set_ylabel('Alcohol (norm)')
ax.legend()
ax.set_title('Wine Tastiness vs Sulphates & Alcohol Contents');

```

We can indeed see that this strange looking wine is placed _above_ the decision boundary, which is why the `.predict()` method returned a value of $1$. 🙌

You might recall that we had many more features available in our dataset `df`. Going beyond two features means that we won't be able to plot the data or the decision boundary in a meaningful manner. However, the predictions are likely to be more accurate!

💪💪 Train a new logistic regression model named `full_clf` using _all_ the features available in our dataset `df`. Then predict the tastiness of the `mystery_wine` 🕵♀, and store it in a variable called `y_mystery`.
Pro-tip: don't use the `tasty` column in your feature matrix X! That should be reserved as label vector, y.
Pro-tip 2: please use the `random_state=0` argument when constructing your sklearn class, to make sure your results are reproducible. Else, the unit test won't pass!
Pro-tip 3: If you also don't like typing out a big list of column names, you might want to consider [`.drop()`](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.drop.html)

```python
mystery_wine = np.array([-0.12, 0.82, 1.43, -0.29, 0.03, -0.90, -0.55, 2.19, -0.62, 0.43, 1.85]).reshape(1, -1)

# INSERT CODE

df_new = df.drop(columns=['tasty'])

X = df_new

y = df['tasty']

full_clf = LogisticRegression(random_state=0).fit(X, y)

y_mystery = full_clf.predict(mystery_wine)

def test_mystery_wine_prediction():
    assert y_mystery == 1, 'The mystery wine is actually tasty! But your prediction says otherwise 😢'
    prob_mystery = full_clf.predict_proba(mystery_wine)[:, 1]
    assert math.isclose(prob_mystery, 0.7615, rel_tol=1e-3), f'Your prediction probability for the mystery wine is wrong'
    print('Success! 🎉')

test_mystery_wine_prediction()

```

    Success! 🎉

```python
# Just to check :
mystery_wine
```

    array([[-0.12,  0.82,  1.43, -0.29,  0.03, -0.9 , -0.55,  2.19, -0.62,
             0.43,  1.85]])

## 6. Summary

Today, we learned about our first classification method, **logistic regression**. We restated that classification tasks aim to predict **categorical** labels, and visualized this for one feature in a two dimensional plot. Linear models were a bad fit binary labels, so we discovered the **logistic function** that could predict **probabilities** between 0 & 1. We described how **thresholds** can turn these probabilities into a binary prediction. The logistic function can be **parameterized** by feeding a **linear model** as its argument, like a set of russian dolls. This means that we can **optimize** our logistic model using the **cross-entropy** cost function and **gradient descent**. We explained what **decision boundaries** are, and how important they are for **classifiers**. We also extended our logistic regression to the **multi-class** setting, where it could predict **multiple labels** with a **softmax** hypothesis. Finally, we put this knowledge to practice by predicting if wines were **tasty**, and visualizing the **decision boundary** of our trained **logistic regression** model.

# Resources
### Core Resources

- [**Slides**](https://docs.google.com/presentation/d/1gk8b6M08T_D8BNbOts2CYkD4qxOeQVWefx7pWTt-YjI/edit?usp=sharing)
- [Coursera Machine Learning - Logistic Regression](https://www.coursera.org/lecture/machine-learning/classification-wlPeP)
Andrew Ng's always excellent course goes into detail for this section on Logistic Regression.
- [Kaggle dataset - red wine quality](https://www.kaggle.com/uciml/red-wine-quality-cortez-et-al-2009)
The dataset used in this notebook.

### Additional Resources

- [ML cheatsheet cost functions](https://ml-cheatsheet.readthedocs.io/en/latest/loss_functions.html)
A handy and mathematical summary of common cost functions, including MSE and cross-entropy.
- [Cross-entropy loss a visual explanation](https://towardsdatascience.com/understanding-binary-cross-entropy-log-loss-a-visual-explanation-a3ac6025181a)
A clear visual yet formal explanation of the cross-entropy cost function.
- [Binary vs multi-class logistic regression](https://chrisyeh96.github.io/2018/06/11/logistic-regression.html#multinomial-logistic-regression-via-cross-entropy)
Comparison of logistic vs softmax functions for logistic regression models.
- [Understanding the softmax function in minutes](https://medium.com/data-science-bootcamp/understand-the-softmax-function-in-minutes-f3a59641e86d)
Comprehensive blogpost defining the softmax function and its uses in Machine Learning models.


```python
from math import exp

from sympy.stats import E, variance, Die, Normal, Binomial
from sympy.stats import P as Prob

import matplotlib.pyplot as plt
```

### terminology

* 'negative' ~ $y=0$
* 'positive' ~ $y=1$
* 'natural' ~ a sample from $X=X_1 + Lap\left(0, \frac{\Delta f}{\epsilon_\text{true}}\right)$
* 'alternative' ~ a special value defined as $\bot=X_1-1000(\frac{\Delta f}{\epsilon_\text{true}}, \frac{\Delta f}{\epsilon_\text{true}}, \frac{\Delta f}{\epsilon_\text{true}})$

```python
from empirical_privacy.laplace_mechanism import EvaluateKNNLaplaceStatDist
from empirical_privacy.laplace_mechanism import FitKNNModelLaplace

from math import sqrt

import numpy as np
from matplotlib.colors import ListedColormap
from sklearn import neighbors, datasets

cmap_light = ListedColormap(['#FFAAAA', '#AAFFAA', '#AAAAFF'])
cmap_bold = ListedColormap(['#FF0000', '#00FF00', '#0000FF'])

def plot_knn(dataset_settings, training_set_size, test_set_size,
             plot_alternative_samples=True, random_seed='hats'):
    scale = eps/ds['epsilon']

    knn = FitKNNModelLaplace(dataset_settings=dataset_settings,
                            samples_per_class=int(round(training_set_size / 2)),
                            random_seed='hats',
                            in_memory=True)
    knn.requires()
    knn.run()

    X = np.vstack((knn.reqs_['samples_positive']['X'],
                   knn.reqs_['samples_negative']['X']))
    y = np.concatenate((knn.reqs_['samples_positive']['y'].ravel(),
                   knn.reqs_['samples_negative']['y'].ravel()))
    if not plot_alternative_samples:
        I = X[:,1] <= -1000 * scale
        X = X[~I,:]
        y = y[~I]
        print('not plotting {} alternative samples'.format(np.argwhere(I).size))

    plot_decision(knn.output_['KNN'], X, y)

def plot_decision(clf, X, y, num=5000):
    plt.figure(figsize=(10,10))
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xs = np.linspace(x_min, x_max, num=int(sqrt(num)))
    ys = np.linspace(y_min, y_max, num=int(sqrt(num)))
    xx, yy= np.meshgrid(xs, ys, )
    Z = clf.predict(np.c_[xx.ravel(), yy.ravel(), np.zeros_like(yy.ravel())])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    plt.pcolormesh(xx, yy, Z, cmap=cmap_light)
    plt.scatter(X[:, 0], X[:, 1], c=y, cmap=cmap_bold ,s=20)

```

### Experiment Settings

```python
eps = 0.01
claimed_eps = eps * 0.9

ds = {
        'database_0' : (0, 0, 0),
        'database_1' : (1, 0, 0),
        'sensitivity': 1.,
        'epsilon'    : eps,
        'claimed_epsilon': claimed_eps,
        'delta'      : 0.
}

CLAIMED_EPS = 0.01
probability_of_natural_sample = 1 / (exp(CLAIMED_EPS))
probability_of_alternative_sample = 1 - probability_of_natural_sample

training_set_size = int(2**14)
test_set_size = int(2**10)
prior_positive = 0.5
prior_negative = 1 - prior_positive  #

plot_knn(ds, training_set_size, test_set_size, plot_alternative_samples=True)
```

## sympy

```python

```

```python
vars_ = [n_natural_negative, n_natural_positive, n_alternative]
P = 1.0
for v in vars_:
    P*=Prob(v(test_set_size) >=1)
    P*=Prob(v(training_set_size) >= 1)

print(P)
```

    0.987930624264561

```python
def n_natural_negative(n):
    return Binomial('n_natural_negative',
                              n=n,
                              p=prior_negative
                         )
def n_natural_positive(n):
    return Binomial('n_natural_positive',
                             n=n,
                             p=prior_positive*probability_of_natural_sample)
def n_alternative(n):
    return Binomial('n_alternative',
                        n=n,
                        p=prior_positive*probability_of_alternative_sample)

```

Notebooks/development/Calculating error as a function of training or test size.ipynb