---
title: "NumPy Broadcasting & JAX"
subtitle: "Lecture 06"
author: "Dr. Colin Rundel"
footer: "Sta 663 - Spring 2025"
format:
  revealjs:
    theme: slides.scss
    transition: fade
    slide-number: true
    self-contained: true
execute:
  echo: true
  warning: true
---

```{r r_setup}
#| message: false
#| warning: false
#| include: false

options(
  width=80
)

local({
  hook_error <- knitr::knit_hooks$get("error")
  hook_output <- knitr::knit_hooks$get("output")
  knitr::knit_hooks$set(
    error = function(x, options) {
      x = sub("## \n## Detailed traceback:\n.*$", "", x)
      x = sub("Error in py_call_impl\\(.*?\\)\\: ", "", x)
      hook_error(x, options)
    },
    output = function(x, options) {
      x = sub("<string>:1: (.*?)\n", "\\1\n\n", x)
      hook_output(x, options)
    }
  )
})
```

```{python setup}
#| include: false
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

plt.rcParams['figure.dpi'] = 200

np.set_printoptions(
  edgeitems=30, linewidth=60,
  precision = 5, suppress=True
)
```


# Basic file IO

## Reading and writing ndarrays

We will not spend much time on this as most data you will encounter is more likely to be a tabular format (e.g. data frame) and tools like Pandas are more appropriate.

For basic saving and loading of NumPy arrays there are the `save()` and `load()` functions which use a custom binary format.

```{python}
x = np.arange(1e5)
np.save("data/x.npy", x)
```

```{python}
new_x = np.load("data/x.npy")
np.all(x == new_x)
```

Additional functions for saving (`savez()`, `savez_compressed()`, `savetxt()`) exist for saving multiple arrays or saving a text representation of an array.

## Reading delimited data

While not particularly recommended, if you need to read delimited (csv, tsv, etc.) data into a NumPy array you can use `genfromtxt()`,

::: {.small}
```{python}
with open("data/mtcars.csv") as file:
    mtcars = np.genfromtxt(file, delimiter=",", skip_header=True)
    
mtcars
```
:::



# Broadcasting

## Broadcasting

This is an approach for deciding how to generalize operations between arrays with differing shapes.

```{python}
x = np.array([1, 2, 3])
```

. . .

```{python}
x * 2
```


. . .

```{python}
x * np.array([2,2,2])
```

. . .

```{python}
x * np.array([2])
```

## Efficiency

Using broadcasts can be more efficient as it does not copy the broadcast data,

::: {.small}
```{python}
x = np.arange(1e5)
y = np.array([2]).repeat(1e5)
```
:::

<br/>

. . .

::: {.small}
```{python}
#| eval: false
%timeit x * 2
```
```
13.1 μs ± 297 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
```
:::

. . .

::: {.small}
```{python}
#| eval: false
%timeit x * np.array([2])
```
```
18.3 μs ± 96.4 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
```
:::

. . .

::: {.small}
```{python}
#| eval: false
%timeit x * y
```
```
61.4 μs ± 465 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
```
:::

. . .

::: {.small}
```{python}
#| eval: false
%timeit x * np.array([2]).repeat(1e5)
```
```
99.2 μs ± 1.12 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
```
:::

## Rules for Broadcasting

> When operating on two arrays, NumPy compares their shapes element-wise. It starts with the trailing (i.e. rightmost) dimensions and works its way left. Two dimensions are compatible when
>
>    1. they are equal, or
>
>    2. one of them is 1
> 
> If these conditions are not met, a `ValueError: operands could not be broadcast together` exception is thrown, indicating that the arrays have incompatible shapes. The size of the resulting array is the size that is not 1 along each axis of the inputs.


::: {.aside}
If you want numpy style broadcasting in R you can check out the [`rray`](https://rray.r-lib.org/index.html) package.
:::

## Example

Why does the code on the left work but not the code on the right?

:::: {.columns .small}
::: {.column width='50%'}
```{python}
#| error: true
x = np.arange(12).reshape((4,3)); x

x + np.array([1,2,3])
```
:::

::: {.column width='50%'}
```{python}
#| error: true
x = np.arange(12).reshape((3,4)); x

x + np.array([1,2,3])
```
:::
::::

. . .

:::: {.columns}
::: {.column width='50%' .callout-note appearance="minimal"}
```
    x    (2d array): 4 x 3
    y    (1d array):     3 
    ----------------------
    x+y  (2d array): 4 x 3
```
:::

::: {.column width='50%' .callout-note appearance="minimal"}
```
    x    (2d array): 3 x 4
    y    (1d array):     3 
    ----------------------
    x+y  (2d array): Error
```
:::
::::



## A fix

::: {.medium}
```{python error=TRUE}
x = np.arange(12).reshape((3,4)); x
```

```{python}
x + np.array([1,2,3]).reshape(3,1)
```
:::

. . .

:::: {.columns}
::: {.column width="25%"}
&nbsp;
:::
::: {.column width="50%" .callout-note appearance="minimal"}
```
    x    (2d array): 3 x 4
    y    (2d array): 3 x 1
    ----------------------
    x+y  (2d array): 3 x 4
```
:::
::::

## Examples (2)

::: {.xsmall}
```{python}
x = np.array([0,10,20,30]).reshape((4,1))
y = np.array([1,2,3])
```
:::

:::: {.columns .xsmall}
::: {.column width='50%'}
```{python}
x
```
```{python}
y
```
:::

::: {.column width='50%' .fragment}

```{python}
x+y
```
:::
::::

. . .

![](imgs/numpy_broadcasting.png){fig-align="center" width="50%"}

::: {.aside}
From NumPy user guide - [Broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html)
:::

## Exercise 1

For each of the following combinations determine what the resulting dimension will be using broadcasting

* A [128 x 128 x 3] + B [3]

* A [8 x 1 x 6 x 1] + B [7 x 1 x 5]

* A [2 x 1] + B [8 x 4 x 3]

* A [3 x 1] + B [15 x 3 x 5]

* A [3] + B [4]


## Demo 1 - Standardization

Below we generate a data set with 3 columns of random normal values. Each column has a different mean and standard deviation which we can check with `mean()` and `std()`.

:::: {.columns .small}
::: {.column}
```{python}
rng = np.random.default_rng(1234)
d = rng.normal(
  loc=[-1,0,1], 
  scale=[1,2,3],
  size=(1000,3)
)
```

```{python}
d.shape
```
:::
::: {.column}
```{python}
d.mean(axis=0)
```

```{python}
d.std(axis=0)
```
:::
::::

Lets use broadcasting to standardize all three columns to have mean 0 and standard deviation 1. 



## Broadcasting and assignment

In addition to arithmetic operators, broadcasting can be used with assignment via array indexing,

::: {.small}
```{python error=TRUE}
x = np.arange(12).reshape((3,4))
y = -np.arange(4)
z = -np.arange(3)
```
:::

:::: {.columns .small}
::: {.column width='50%'}
```{python}
x[:] = y
x
```
```{python}
x[...] = y
x
```
:::

::: {.column width='50%'}
```{python}
#| error: true
x[:] = z
```
```{python}
x[:] = z.reshape((3,1))
x
```
:::
::::





# JAX

## JAX

> JAX is a library for array-oriented numerical computation (à la NumPy), with automatic differentiation and JIT compilation to enable high-performance machine learning research.
> 
> * JAX provides a unified NumPy-like interface to computations that run on CPU, GPU, or TPU, in local or distributed settings.
>
> * JAX features built-in Just-In-Time (JIT) compilation via Open XLA, an open-source machine learning compiler ecosystem.
>
> * JAX functions support efficient evaluation of gradients via its automatic differentiation transformations.
>
> * JAX functions can be automatically vectorized to efficiently map them over arrays representing batches of inputs.

::: {.small}
```{python}
import jax
jax.__version__
```
:::

## JAX & NumPy

:::: {.columns .small}
::: {.column width='50%'}
```{python}
import numpy as np

x_np = np.linspace(0, 5, 51)
y_np = 2 * np.sin(x_np) * np.cos(x_np)
plt.plot(x_np, y_np)
```
```{python}
type(x_np)
```
:::

::: {.column width='50%'}
```{python}
import jax.numpy as jnp

x_jnp = jnp.linspace(0, 5, 51)
y_jnp = 2 * jnp.sin(x_jnp) * jnp.cos(x_jnp)
plt.plot(x_jnp, y_jnp)
```
```{python}
type(x_jnp)
```
:::
::::



##

:::: {.columns .small}
::: {.column width='50%'}
```{python}
x_np
```
```{python}
x_np.dtype
```
:::

::: {.column width='50%'}
```{python}
x_jnp
```
```{python}
x_jnp.dtype
```
:::
::::



## Compatibility

::: {.small}
```{python}
y_mix = 2 * np.sin(x_jnp) * jnp.cos(x_np); y_mix
```
:::

. . .

::: {.small}
```{python}
type(y_mix)
```
:::

## JAX Arrays

As we've just seen a JAX array is very similar to a numpy array but there are some important differences.

- JAX arrays are immutable

::: {.xsmall}
```{python}
#| error: true
x = jnp.array([3, 2, 1])
x[0] = 2
```
:::

- related to the above JAX does not support inplace operations - these functions now create and return a copy of the array

::: {.columns .xsmall}
::: {.column}
```{python}
y = x.sort()
y
x
```
:::
::: {.column}
```{python}
np.shares_memory(x,y)
```
:::
:::

##

- The default JAX array dtypes 32 bits not 64 bits (i.e. `float32` not `float64` and `int32` not `int64`)

  ::: {.columns .xsmall}
  ::: {.column}
  ```{python}
  #| warning: true
  jnp.array([1, 2, 3])
  ```
  :::
  ::: {.column}
  ```{python}
  #| warning: true
  jnp.array([1., 2., 3.])
  ```
  :::
  :::

  ::: {.xsmall}
  ```{python}
  #| warning: true
  jnp.array([1, 2, 3], dtype=jnp.float64)
  ```
  :::

  64-bit dtypes can be enabled by setting `jax_enable_x64=True` in the JAX configuration.

  ::: {.xsmall}
  ```{python}
  jax.config.update("jax_enable_x64", True)
  ```
  :::

  ::: {.columns .xsmall}
  ::: {.column}
  ```{python}
  jnp.array([1, 2, 3])
  ```
  :::
  ::: {.column}
  ```{python}
  jnp.array([1., 2., 3.])
  ```
  :::
  :::




##

- JAX arrays are allocated to one *or more* devices

::: {.columns .small}
::: {.column}
```{python}
jax.devices()
```
:::
::: {.column}
```{python}
x.devices()
x.sharding
```
:::
:::

- Using JAX interactively allows for the use of standard Python control flow (`if`, `while`, `for`, etc.) but this is not supported for some of JAX's more advanced operations (e.g. `jit` and `grad`)

  There are replacements for most of these constructs in JAX, but they are beyond the scope of today.



# Random number generation

## JAX vs NumPy

Pseudo random number generation in JAX is a bit different than with NumPy - the latter depends on a global state that is updated each time a random function is called.

NumPy's PRNG guarantees something called sequential equivalence which amounts to sampling N numbers sequentially is the same as sampling N numbers at once (e.g. a vector of length N).

::: {.small}
```{python}
np.random.seed(0)
f"individually: {np.stack([np.random.uniform() for i in range(5)])}"
```

```{python}
np.random.seed(0)
f"at once: {np.random.uniform(size=5)}"
```
:::


## Parallelization & sequential equivalence

Sequential equivalence can be problematic in when using parallelization across multiple devices, consider the following code:

::: {.small}
```{python}
np.random.seed(0)

def bar(): 
  return np.random.uniform()

def baz(): 
  return np.random.uniform()

def foo(): 
  return bar() + 2 * baz()
```
:::

How do we guarantee that we get consistent results if we don't know the order that `bar()` and `baz()` will run?


## PRNG keys

JAX makes use of *random keys* which are just a fancier version of random seeds - all of JAX's random functions require a key as their first argument.

::: {.small}
```{python}
key = jax.random.PRNGKey(1234); key
```
:::

. . .

::: {.small}
```{python}
jax.random.normal(key)
```
```{python}
jax.random.normal(key)
```
:::

. . .

::: {.small}
```{python}
jax.random.normal(key, shape=(3,))
```
:::

Note that JAX does not provide a sequential equivalence guarantee - this is so that it can support vectorization for the generation of PRN.


## Splitting keys

Since a key is essentially a seed we do not want to reuse them (unless we want an identical output). Therefore to generate multiple different PRN we can split a key to deterministically generate two (or more) new keys.

:::: {.columns .xsmall}
::: {.column width='50%'}
```{python}
key11, key12 = jax.random.split(key)
f"{key=}"
f"{key11=}"
f"{key12=}"
```
:::

::: {.column width='50%'}
```{python}
key21, key22 = jax.random.split(key)
f"{key=}"
f"{key21=}"
f"{key22=}"
```
:::
::::

. . .

::: {.xsmall}
```{python}
key3 = jax.random.split(key, num=3)
key3
```
:::

## 

::: {.small}
```{python}
jax.random.normal(key, shape=(3,))
```
:::

. . .
 
::: {.columns .small}
::: {.column}
```{python}
jax.random.normal(key11, shape=(3,))
jax.random.normal(key12, shape=(3,))
```
:::
::: {.column}

```{python}
jax.random.normal(key21, shape=(3,))
jax.random.normal(key22, shape=(3,))
```
:::
:::

. . .

::: {.small}
```{python}
jax.random.normal(key3[0], shape=(3,))
jax.random.normal(key3[1], shape=(3,))
jax.random.normal(key3[2], shape=(3,))
```
:::


# JAX & jit

## JAX performance

::: {.small}
```{python}
key = jax.random.PRNGKey(1234)
x_jnp = jax.random.normal(key, (1000,1000))
x_np = np.array(x_jnp)
```
:::

::: {.small}
```{python}
#| eval: false
%timeit y = x_np @ x_np
```
```
1.09 ms ± 92.4 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```

```{python}
#| eval: false
%timeit y = x_jnp @ x_jnp
```
```
3.42 ms ± 122 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```
:::

. . .

::: {.small}
```{python}
#| eval: false
%timeit y = 3*x_np + x_np
```
```
514 μs ± 41 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```
```{python}
#| eval: false
%timeit y = 3*x_jnp + x_jnp
```
```
413 μs ± 24.2 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```
:::


## jit

:::: {.columns .xsmall}
::: {.column width='50%'}
```{python}
def SELU_np(x, α=1.67, λ=1.05):
  "Scaled Exponential Linear Unit"
  return λ * np.where(x > 0, x, α * np.exp(x) - α)
```
:::

::: {.column width='50%'}
```{python}
def SELU_jnp(x, α=1.67, λ=1.05):
  "Scaled Exponential Linear Unit"
  return λ * jnp.where(x > 0, x, α * jnp.exp(x) - α)
```
:::
::::

. . .

:::: {.columns .small}
::: {.column width='50%'}
```{python}
#| eval: false
x = np.arange(1e6)
%timeit y = SELU_np(x)
```
```
4.08 ms ± 80 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```
:::

::: {.column width='50%' .fragment}
```{python}
#| eval: false
x = jnp.arange(1e6)
%timeit y = SELU_jnp(x)
```
```
1.58 ms ± 68.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```
:::
::::


. . .

:::: {.columns .small}
::: {.column width='50%'}
```{python}
#| error: true
SELU_np_jit = jax.jit(SELU_np)
```
:::

::: {.column width='50%'}
```{python}
SELU_jnp_jit = jax.jit(SELU_jnp)
```
:::
::::

. . .

:::: {.columns .small}
::: {.column width='50%'}

```{python}
#| error: true
#| eval: false
%timeit y = SELU_np_jit(x)
```
```
TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[1000000]
```
:::

::: {.column width='50%' .fragment}

```{python}
#| eval: false
%timeit y = SELU_jnp_jit(x)
```
```
418 μs ± 13 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```
:::
::::


## jit limitations

When it works the jit tool is fantastic, but it does have a number of limitations,

* Must use [pure functions](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions) (no side effects)

* Must primarily use JAX functions
  * e.g. use `jnp.minimum()` not `np.minimum()` or `min()`

* Must generally avoid conditionals / control flow

* Issues around concrete values when tracing (static values)

* Check performance - there are not always gains + there is the initial cost of compilation


# Automatic differentiation

## Basics

The `grad()` function takes a numerical function, returning a scalar, and returns a function for calculating the gradient of that function.

:::: {.columns .xsmall}
::: {.column width='33%' .fragment}
```{python}
def f(x):
  return x**2
```

```{python}
f(3.)
```
```{python}
jax.grad(f)(3.)
```
```{python}
jax.grad(
  jax.grad(f)
)(3.)
```
:::

::: {.column width='33%' .fragment}
```{python}
def g(x):
  return jnp.exp(-x)
```

```{python}
g(1.)
```
```{python}
jax.grad(g)(1.)
```
```{python}
jax.grad(
  jax.grad(g)
)(1.)
```
:::

::: {.column width='33%' .fragment}
```{python}
def h(x):
  return jnp.maximum(0,x)
```

```{python}
h(-2.)
```
```{python}
h(2.)
```
```{python}
jax.grad(h)(-2.)
```
```{python}
jax.grad(h)(2.)
```
:::
::::

## Aside - `vmap()`

I would like to plot `h()` and `jax.grad(h)()` - lets see what happens,

::: {.small}
```{python}
#| error: true
x = jnp.linspace(-3,3,101)
y = h(x)
y_grad = jax.grad(h)(x)
```
:::

. . .

As just mentiond, we can only calculate the gradient for scalar valued functions. However, we can transform our scalar function into a vectorized function using `vmap()`.

::: {.small}
```{python}
h_grad = jax.vmap(
  jax.grad(h)
)
y_grad = h_grad(x)
```
:::

::: {.aside}
`vmap()` is significantly more powerful than just this as it allows for mapping over multiple input and output axes
:::

##

::: {.xsmall}
```{python}
y_grad
```
:::

```{python}
#| echo: false

plt.figure(figsize=(10, 4))

plt.subplot(121)
plt.plot(x, y, "-b")
plt.title("h(x)")

plt.subplot(122)
plt.plot(x, y_grad, "-r")
plt.title("∇h(x)")

plt.show()
```

## Another quick example

::: {.xsmall}
```{python}
x = jnp.linspace(-6,6,101)
f = lambda x: 0.5 * (jnp.tanh(x / 2) + 1)
y = f(x)
y_grad = jax.vmap(jax.grad(f))(x)
```
:::

```{python}
#| echo: false
plt.figure(figsize=(10, 4))

plt.subplot(121)
plt.plot(x, y, "-b")
plt.title("f(x)")

plt.subplot(122)
plt.plot(x, y_grad, "-r")
plt.title("∇f(x)")

plt.show()
```
