[markov_asset] Update to JAX and runtime comparison#654
[markov_asset] Update to JAX and runtime comparison#654xuanguang-li wants to merge 14 commits intomainfrom
Conversation
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (0317e4a) 📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions |
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (76c64b2) 📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions |
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (02dfe6e) 📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions |
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (211ec24) 📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions |
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (cb2995a) 📚 Changed Lecture Pages: markov_asset |
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (4adc2f5) 📚 Changed Lecture Pages: markov_asset |
|
📖 Netlify Preview Ready! Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (e47f942) 📚 Changed Lecture Pages: markov_asset, markov_asset_numpy |
This pull request modernizes the
markov_asset.mdnotebook by refactoring the codebase to use JAX and JAX NumPy.numpyandnumpy.linalgtojax.numpyandjax.numpy.linalg.AssetPriceModelclass with an immutable, type-safeNamedTupleversion and introduced acreate_ap_modelfactory function for consistent construction.jax.lax.while_loopandjax.lax.fori_loop.Update: Runtime Comparison Between
JAX (GPU),JAX (CPU), andNumPyMethodology:
JAXversion uses the code in this PR, while theNumPyversion uses the code inmain.JAX (GPU)is measured using Google Colab.qe.timeitover 1,000 iterations.Results:
JAX (CPU)>NumPy>JAX (GPU).