```python
import jax
import jax.numpy as jnp
from IPython.display import display, Latex
```

## If else condition with `lax`

$$
f(\mathbf{x}) = \sum_{x \in \mathbf{x}} \begin{cases}
    x^2,& \text{if } x \gt 5\\
    x^3,             & \text{otherwise}
\end{cases}
$$


```python
x = [jnp.array(10.0), jnp.array(2.0)]


@jax.jit
@jax.value_and_grad
def f(x):
    bool_val = jax.tree_map(lambda val: val > 5.0, x)
    ans = jax.tree_map(
        lambda val, bool: jax.lax.cond(bool, lambda: val**2, lambda: val**3),
        x,
        bool_val,
    )
    return jax.tree_util.tree_reduce(lambda a, b: a + b, ans)


value, grad = f(x)

display(Latex(f"$f(\mathbf{{x}}) = {value}$"))
for idx in range(len(x)):
    display(Latex(f"$\\frac{{df}}{{dx_{idx}}} = {grad[idx]}$"))
```


$f(\mathbf{x}) = 108.0$



$\frac{df}{dx_0} = 20.0$



$\frac{df}{dx_1} = 12.0$


## Pair-wise distance with `vmap`


```python
# create vour pairwise function
def distance(a, b):
    return jnp.linalg.norm(a - b)


# map based combinator to operate on all pairs
def all_pairs(f):
    f = jax.vmap(f, in_axes=(None, 0))
    f = jax.vmap(f, in_axes=(0, None))
    return f


# transform to operate over sets
distances = all_pairs(distance)

# Example
x = jnp.array([1.0, 2.0, 3.0])
y = jnp.array([3.0, 4.0, 5.0])
distances(x, y)
```




    DeviceArray([[2., 3., 4.],
                 [1., 2., 3.],
                 [0., 1., 2.]], dtype=float32)



## Compute Hessian with `jax`

Let us consider Linear regression loss function

\begin{align}
\mathcal{L}(\boldsymbol{\theta}) &= (\boldsymbol{y} - X\boldsymbol{\theta})^T(\boldsymbol{y} - X\boldsymbol{\theta})\\
\frac{d\mathcal{L}}{d\boldsymbol{\theta}} &= -2X^T\boldsymbol{y} + 2X^TX\boldsymbol{\theta}\\
H_{\mathcal{L}}(\boldsymbol{\theta}) &= 2X^TX
\end{align}


```python
def loss_function_per_point(theta, x, y):
    y_pred = x.T @ theta
    return jnp.square(y_pred - y)


def loss_function(theta, x, y):
    loss_per_point = jax.vmap(loss_function_per_point, in_axes=(None, 0, 0))(
        theta, x, y
    )
    return jnp.sum(loss_per_point)


def gt_loss(theta, x, y):
    return jnp.sum(jnp.square(x @ theta - y))


def gt_grad(theta, x, y):
    return 2 * (x.T @ x @ theta - x.T @ y)


def gt_hess(theta, x, y):
    return 2 * x.T @ x
```

### Simulate dataset 


```python
key = jax.random.PRNGKey(0)
key, subkey1, subkey2 = jax.random.split(key, num=3)
N = 100
D = 11
x = jax.random.uniform(key, shape=(N, D))
y = jax.random.uniform(subkey1, shape=(N,))
theta = jax.random.uniform(subkey2, shape=(D,))
```

### Verify loss and gradient values


```python
loss_and_grad_function = jax.value_and_grad(loss_function)

loss_val, grad = loss_and_grad_function(theta, x, y)

assert jnp.allclose(loss_val, gt_loss(theta, x, y))
assert jnp.allclose(grad, gt_grad(theta, x, y))
```

### Verify hessian matrix

#### Way-1 


```python
hess = jax.hessian(loss_function)(theta, x, y)

assert jnp.allclose(hess, gt_hess(theta, x, y))
```

#### Way-2


```python
hess = jax.jacfwd(jax.jacrev(loss_function))(theta, x, y)

assert jnp.allclose(hess, gt_hess(theta, x, y))
```

## `tree_map` in JAX

The only requirement for `tree_map` to work is, output should have the same structure as the first argument (as explained [here](https://github.com/deepmind/distrax/issues/147)). For example:


```python
import jax
import jax.numpy as jnp
try:
    import distrax
except:
    %pip install -qqq distrax
    import distrax

dists = {"Normal": distrax.Normal(3.0, 4.0), "Gamma": distrax.Gamma(3.0, 4.0)}
samples = {"Normal": jnp.array(2.0), "Gamma": jnp.array(3.0)}
try:
    log_probs = jax.tree_map(lambda dist, sample: dist.log_prob(sample), dists, samples)
except Exception as e:
    print(e)
```

    Custom node type mismatch: expected type: <class 'distrax._src.distributions.normal.Normal'>, value: DeviceArray(2., dtype=float32, weak_type=True).


The problem here is that `dists` do not have same structure as `log_probs` (`log_probs` structure matches with `samples`). So, we should keep `samples` as the first argument: 


```python
log_probs = jax.tree_map(lambda sample, dist: dist.log_prob(sample), samples, dists)
log_probs
```




    {'Gamma': DeviceArray(-6.337041, dtype=float32, weak_type=True),
     'Normal': DeviceArray(-2.336483, dtype=float32, weak_type=True)}



## Use of `lax.scan` to accelerate a training loop

Here we create a dummy training loop and check the performance of `lax.scan`. The example also shows how to convert a training loop to `lax.scan` version of it.


```python
value_and_grad_fun = jax.jit(jax.value_and_grad(lambda x: jnp.sum(x**2)))


def training_loop(n_iters, params):
    values = []
    for i in range(n_iters):
        value, grad = value_and_grad_fun(params)
        params = params - learning_rate * grad
        values.append(value)
    return value


@jax.jit
def one_step(params, xs):
    value, grad = value_and_grad_fun(params)
    params = params - learning_rate * grad
    return params, value
```


```python
key = jax.random.PRNGKey(0)
N = 1000
n_iters = 10000
learning_rate = 0.01

params = jax.random.uniform(key, (N,))
```


```python
training_loop(1, params)  # warn up
%timeit -n 1 -r 1 training_loop(n_iters, params)
```

    1.26 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)



```python
%timeit -n 1 -r 1 jax.lax.scan(one_step, params, xs=None, length=n_iters)
```

    221 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)


Note that `xs` array can be passed in case we want to scan over it. An example of it can be found in this [blackjax documentation](https://blackjax-devs.github.io/blackjax/examples/Introduction.html).

ec4d738b0523bedb3fae06d1f03baedfbf9dcba0
ipynb
Jupyter Notebook
notebooks/Practical_JAX_Tips.ipynb
patel-zeel/probml-notebooks
1ff09bfddb2bd6b3932d81845546770e7e2fce3a
MIT
notebooks/Practical_JAX_Tips.ipynb
patel-zeel/probml-notebooks
1ff09bfddb2bd6b3932d81845546770e7e2fce3a
MIT
2022-03-30T20:00:48.000Z
2022-03-30T20:30:42.000Z
notebooks/Practical_JAX_Tips.ipynb
patel-zeel/probml-notebooks
1ff09bfddb2bd6b3932d81845546770e7e2fce3a
MIT
Qwen/Qwen-72B
1. YES
2. YES