August 26, 2024MicroJaxJAXmicrojaxTPUpytorchflaxJAX from scratchMicroJax is a minimal implementation of JAX in pure Python. It is a fun project to learn about JAX and how it works under the hood.