Phil Wang/lucidrains has phenomenal implementations of stuff, I'd recommend checking them out and reading their code.

Furthermore, I'd recommend simply reading more code and tackling complex problems, e.g. try building a DL framework from "scratch" ontop of jax. Read the Haiku codebase, and compare it to say Equinox (I am a big fan of this one). Go through the huggingface code bases, e.g. transformers. Choose a model and build it from scratch and make it compatible with their API.