!pip install datasets
from IPython.display import clear_output
clear_output()

# **Q1.** Autograd for MLPs from scratch (20 pts)

Automatic differentiation is the backbone of modern deep learning toolkits like Torch, JAX, TensorFlow, and MLX. To build an understanding of how these tools work, we will define functions and their gradients/Jacobians and compose them together to as the building blocks of a neural network

The gradient of a real-valued function $f:\mathbb{R}^n\to\mathbb{R}$ is a vector where
$$\nabla f(\boldsymbol{x})_{i}=\frac{\partial f(\boldsymbol{x})}{\partial x_i}$$

The Jacobian is a generalization of the gradient for a vector-valued function $f:\mathbb{R}^m\to\mathbb{R}^n$. It is a $m\times n$ matrix where
$$J_f(\boldsymbol{x})_{i,j}=\frac{\partial f(\boldsymbol{x})_i}{\partial x_j}$$

Neural nets are generally built from composed functions. To calculate the Jacobian of composed functions, we can use matrix multiplication to apply the chain rule.

$$
J_{f\circ g}(\boldsymbol{x})
= J_f(g(\boldsymbol{x}))J_g(\boldsymbol{x})
$$

([Proof](https://math.stackexchange.com/questions/3260609/how-to-show-jacobian-of-a-composite-function-is-the-product-of-jacobians).)

We will define a series of functions which take 2D inputs and return 2D outputs.
```
arg1: (M, N)
arg2: (O, P)
...
returns: (Q, R)
```
We will also define Jacobian functions, which return the function output as well as the Jacobian vector of each of the inputs.
```
arg1: (M, N)
arg2: (O, P)
...
returns: (Q, R), (Q, R, M, N), (Q, R, O, P), ...
```
Note that the first two dimensions of the Jacobians match the output shape (Q, R), and the last two match the dimensions of the arguments so that they can be composed together

```python
output1, jac1 = func1_and_jac(x) # Shape: (M, N, O, P)
output2, jac2 = func2_and_jac(output1) # Shape: (O, P, Q, R)
jac = np.einsum("mnop,opqr->mnqr", jac1, jac2) # Shape (M, N, Q, R)
```

import jax
import jax.numpy as np
from numpy import testing
import math

def test(func, func_and_jac, *params):
  testing.assert_allclose(func(*params), func_and_jac(*params)[0])
  for i in range(len(params)):
    print(f"Testing Jacobian of parameter {i}")
    testing.assert_allclose(jax.jacfwd(func, argnums=i)(*params), func_and_jac(*params)[i + 1])
  print("All passed!")

# Reference implementation of a func_and_jac function.
def add_and_jac(x, y):
  output = x + y
  x_jac = np.eye(math.prod(x.shape)).reshape(*x.shape, *x.shape)
  x_jac = np.broadcast_to(x_jac, (*output.shape, *x.shape))
  y_jac = np.eye(math.prod(y.shape)).reshape(*y.shape, *y.shape)
  y_jac = np.broadcast_to(y_jac, (*output.shape, *y.shape))
  return output, x_jac, y_jac

x, y = np.arange(24).reshape(2, 3, 4).astype(float)
test(np.add, add_and_jac, x, y)
# print(add_and_jac(x, y))

def matmul_and_jac(x, y):
    # Perform matrix multiplication
    output = np.matmul(x, y)
    
    # Create zero-initialized Jacobians
    x_jac = np.zeros((*output.shape, *x.shape))
    y_jac = np.zeros((*output.shape, *y.shape))
    
    # Fill the Jacobian for x and y using the .at[] method
    for i in range(output.shape[0]):
        for j in range(output.shape[1]):
            for k in range(x.shape[1]):
                # Use the .at[] method to set values in JAX arrays
                x_jac = x_jac.at[i, j, i, k].set(y[k, j])
                y_jac = y_jac.at[i, j, k, j].set(x[i, k])
    
    return output, x_jac, y_jac

x = np.arange(12).reshape(3, 4).astype(float)
y = np.arange(20).reshape(4, 5).astype(float)
test(np.matmul, matmul_and_jac, x, y)

def affine(x, w, b):
  return np.dot(x, w) + b

def affine_and_jac(x, w, b):
    # First apply the matrix multiplication and get its Jacobians
    matmul_output, x_jac_w, w_jac = matmul_and_jac(x, w)
    
    # Now apply the addition and get its Jacobians
    affine_output, _, b_jac = add_and_jac(matmul_output, b)
    
    return affine_output, x_jac_w, w_jac, b_jac

x = np.arange(12).reshape(3, 4).astype(float)
w = np.arange(20).reshape(4, 5).astype(float)
b = np.arange(5).reshape(1, 5).astype(float)
test(affine, affine_and_jac, x, w, b)

def relu(x):
  return np.maximum(x, 0)

def relu_and_jac(x):
  # Some libraries define $d relu(0) / dx$ to be 0.5.
  x_jac = (np.diagflat(x > 0) + np.diagflat(x == 0) / 2).reshape(*x.shape, *x.shape)
  return np.maximum(x, 0), x_jac

def logsoftmax(x):
  stabilized_x = x - np.max(x, axis=1, keepdims=True)
  return stabilized_x - np.log(np.sum(np.exp(stabilized_x), axis=1, keepdims=True))

def logsoftmax_and_jac(x):
    ...

x = np.arange(12).reshape(3, 4).astype(float)
# print(jax.jacfwd(logsoftmax)(x))
# print("other", logsoftmax_and_jac(x))
test(logsoftmax, logsoftmax_and_jac, x)

implement logsoftmax_and_jac(x) using the examples above. Use the logsoftmax(x) function where necessary. Avoid importing new packages