Add `jaxtyping` for torch/jax code.

[patel-zeel]

This blog explains the advanced usage of jaxtyping.

Examples

• Actively raise an error if the shapes of inputs do not match with defined shapes.
• Check if output_shape is input_shape - 2. It evaluates the math.