Closely related to this if you're interested in the topic if Deep Mind's guide on how to scale your model.

https://jax-ml.github.io/scaling-book/roofline/