I love this. I wish every big library had something like this. It helped me contribute to JAX in the past, and is a great educational resource and source of inspiration for my own tools.
I’ve tried to find something similar for pytorch and numpy in the past and was let down.