Demonstrate how to add JIT using MLIR to micrograd

[fzakaria]

Hi @karpathy !

I'm not expecting you to merge this (although I'd very much welcome it!) -- but I wanted to contribute publicly work myself and @alexander-shaposhnikov have done to demonstrate adding a JIT Just In Time compiler for micrograd.

The main change here is the introduction of a new jit.py module which can take various micrograd computation graphs: Value, Neuron, Layer etc.. and produce MLIR using the arithmetic dialect. The IR is then lowered to LLVM IR which can then be executed directly via a provided CPU execution engine.

test_jit.py has some great examples but the API is straightforward

def test_mlp_complex_multiple_out():
    random.seed(10)
    nn = MLP(nin=2, nouts=[2, 2])
    jnn = jit(nn)
    args = [-30., -20.]
    for r, jr in zip(nn(args), jnn(args)):
        assert math.isclose(r.data, jr, abs_tol=1e-04)

Follow-ups:

• There's more opportunity to improve the JIT by maybe reproducing Layer & MLP using the linalg & tensor dialects.
• Demonstrate how MLIR can be lowered and executed via a accelerator (i.e. CUDA) using the gpu dialect.

Changes done to the repository:

• added more to .gitignore
• added support for direnv which automatically sets up venv
• created a requirements.txt file -- to help bring in MLIR dependencies
• added init.py to the test directory so that you can run pytest by itself
• introduced a new jit.py module
• added documentation
• added a similar to cell to demo.ipynb to show the example using JIT