Microjax: JAX em duas classes e seis funções

Inspirado no Micrograd de Andrej Karpathy, o Microjax é uma biblioteca que replica a funcionalidade do JAX usando apenas duas classes e seis funções. Ao contrário do popular PyTorch, o Microjax adota o estilo de programação mais funcional do JAX. Este tutorial toma emprestado muito da excelente apresentação de Matthew J Johnson de 2017 sobre o autograd, o predecessor do JAX, simplificando-o e empacotando-o como um notebook.