This is a simple JAX pseudo-framework built on top of Haiku's amazing tutorial. It was built as a project to begin my studies with JAX and implementing several DL architectures/blocks from scratch. It is not a serious framework, so please, don't try to use it on a real project.
Feel free to clone this repo, try it out and maybe add any additional functionality :)