Flash attention puzzle

This PR adds a new notebook with a guided puzzle to implement the forward pass of FlashAttention2. It also makes some modifications to the visualiser to allow it to render the solution (example solution for testing can be found [here](https://github.com/MPKonst/GPU-Puzzles/blob/attention/Flash_attention_solution.ipynb)). The main modifications needed were: - adding support for more binary operations on `Scalar` and `ScalarHistory` - monkeypatching math.exp - allowing for more colours, so examples with more TPB can be plotted. The other changes to lib.py are mostly cosmetic.