Upgrade Tensor Puzzlers from torchtyping to jaxtyping.

[davideger]

jaxtyping is a new, improved, and maintained version of torchtyping by the same author (patrick kidger). When used with beartype, jaxtyping can inform the user of shape mismatch errors at run time.

Other minor formatting issues were fixed in the Tensor Puzzles notebook.

[srush]

Amazing! I was just planning on doing this.

[srush]

Just curious, Why do you need "{j j}"?

[davideger]

My understanding (and @patrick-kidger correct me if I am wrong) is that if you want jaxtyping to ensure that a return type dimension matches a scalar function parameter, you need to use f-string escaping to reference the function parameter. At least that's what I get from interpreting https://github.com/google/jaxtyping/blob/main/docs/api/array.md.