This repository contains JAX implementations of Dynamic Programming Techniques commonly used in economics:
- Brute-force Value Function Iteration.
- Standard (F.O.C. based) Value Function Iteration.
The goal of this repository is to provide optimized versions of the algorithms that can run on CPU, GPU or TPU efficiently, using the simplity allowed by jax to leverage automatic vectorization (jax.vmap), automatic differntiation (jax.grad), parallelization (jax.pmap) and compiler-based optimizations (jax.jit). The intention is for them to be simple and readable implementations to build research on top of. If you use econjax in your work, please cite this repository in publications:
@misc{jaxrl,
author = {Covarrubias, Matias},
doi = {},
month = {},
title = {{ECONJAX: Implementations of Dynamic Programming Algorithms using JAX}},
url = {https://github.com/MatiasCovarrubias/econjax/edit/main/README.md},
year = {2021}
}
- Created the private repository and an implementation of the Brute-Force VFI algorithm.
Explain the intended use.
Explain Troubleshooting
Launch tensorboard to see training and evaluation logs
tensorboard --logdir=./tmp/
Explain results
Code for test scripts.
When contributing to this repository, please first discuss the change you wish to make via issue. If you are not familiar with pull requests, please read this documentation.
I want to thank Agustin Covarrubias for his collaboration in setting things in Google Cloud.