# Kesten Processes and Firm Dynamics

```{admonition} GPU in use

```{admonition} GPU

:class: warning

This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming.

This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming.

Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU.

Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax) for installing JAX with GPU support.

Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax) for installing JAX with GPU support. If you would like to install jax running on the `cpu` only you can use `pip install jax[cpu]`

```

```{index} single: Linear State Space Models

---

!pip install quantecon

!pip install --upgrade yfinance

# If your machine has CUDA support, please follow the guide in GPU Warning.

# Otherwise, run the line below:

!pip install --upgrade "jax[CPU]"

```

## Overview

:class: dropdown

```

Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html).

Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html).

First let's import the necessary modules and check the backend for JAX

exp_a = jnp.exp(a_random[t, :])

exp_b = jnp.exp(b_random[t, :])

exp_e = jnp.exp(e_random[t, :])

s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar,

s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar,

exp_e,

exp_a * s[:, t] + exp_b))

return s[:, -1]

%time data = generate_draws().block_until_ready()

The plot produces a straight line, consistent with a Pareto tail.

It is possible to further speed up our code by replacing the `for` loop with [`lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)

It is possible to further speed up our code by replacing the `for` loop with [`lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)

to reduce the loop overhead in the compilation of the jitted function

```{code-cell} ipython3

M=1_000_000,

s_init=1.0,

seed=123):

key = random.PRNGKey(seed)

keys = random.split(key, 3)

# Generate random draws and initial values

a_random = μ_a + σ_a * random.normal(keys[0], (T, M))

b_random = μ_b + σ_b * random.normal(keys[1], (T, M))

# Define the function for each update

def update_s(s, a_b_e_draws):

a, b, e = a_b_e_draws

res = jnp.where(s < s_bar,

jnp.exp(e),

res = jnp.where(s < s_bar,

jnp.exp(e),

jnp.exp(a) * s + jnp.exp(b))

return res, res

# Use lax.scan to perform the calculations on all states

s_final, _ = lax.scan(update_s, s, (a_random, b_random, e_random))

return s_final

```

```{solution-end}

```

```