```{code-cell} ipython3

def K(

a_in: jnp.array, # a_in[i, z] is an asset grid

c_in: jnp.array, # c_in[i, z] = consumption at a_in[i, z]

a_in: jnp.array, # a_in[i, z] is an asset grid

ifp: IFP

):

"""

c_out = c_out.at[0, :].set(0)

a_out = a_out.at[0, :].set(0)

return a_out, c_out

return c_out, a_out

```

The next function solves for an approximation of the optimal consumption policy

Let's generate an approximation solution with JAX:

```{code-cell} ipython3

a_star, σ_star = solve_model(ifp, a_init, σ_init)

σ_star, a_star = solve_model(ifp, σ_init, a_init)

```

Let's try it again with a timer.

```{code-cell} python3

with qe.Timer(precision=8):

a_star, σ_star = solve_model(ifp, a_init, σ_init)

a_star.block_until_ready()

σ_star, a_star = solve_model(ifp, σ_init, a_init)

σ_star.block_until_ready()

```

## Simulation

n_z = len(ifp.P)

a_init = s_grid[:, None] * jnp.ones(n_z)

c_init = a_init

a_vec, c_vec = solve_model(ifp, a_init, c_init)

c_vec, a_vec = solve_model(ifp, c_init, a_init)

assets = compute_asset_stationary(c_vec, a_vec, ifp, num_households=200_000)

# Compute Gini coefficient for the plot

n_z_temp = len(ifp_temp.P)

a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)

c_init_temp = a_init_temp

a_vec_temp, c_vec_temp = solve_model(

ifp_temp, a_init_temp, c_init_temp

c_vec_temp, a_vec_temp = solve_model(

ifp_temp, c_init_temp, a_init_temp

)

# Simulate households

n_z_temp = len(ifp_temp.P)

a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)

c_init_temp = a_init_temp

a_vec_temp, c_vec_temp = solve_model(

ifp_temp, a_init_temp, c_init_temp

c_vec_temp, a_vec_temp = solve_model(

ifp_temp, c_init_temp, a_init_temp

)

# Simulate households