import numpy as np, matplotlib.pyplot as plt
E, A, B, a, b = 1.69, 406.4, 410.7, 0.34, 0.28
plt.figure(figsize=(5.4, 3.0))
opt_N, opt_L = [], []
for C in [1e18, 1e20, 1e22, 1e24]:
N = np.logspace(7, 12, 300); D = C/(6*N) # fixed compute: C = 6ND
L = E + A/N**a + B/D**b
plt.loglog(N, L, alpha=0.8, label=f"C=10$^{{{int(np.log10(C))}}}$")
i = np.argmin(L); opt_N.append(N[i]); opt_L.append(L[i])
plt.loglog(opt_N, opt_L, "ko--", lw=1, ms=5, label="optimal")
plt.xlabel("parameters $N$ (log)"); plt.ylabel("loss (log)")
plt.title("Each compute budget has an optimal model size"); plt.legend(fontsize=6)
plt.tight_layout(); plt.show()