AlexanderMath

Hmm, it might be that the ```@static``` part of SpectralNorm breaks somehow.

I tried with tree different checkpoints at 10k and 20k iterations, the issue persists. The only difference I found is that sometimes the generator loss becomes very large (+60) instead...

Do you think it would be sufficient to remove SpectralNorm class and add a line ```SpectralNorm = torch.nn.utils.spectral_norm```?

**Note.** For [water STO-3G test-case ](https://github.com/graphcore-research/pyscf-ipu/blob/023e11487762731117d73d14bc98b2f74226f800/test/test_experimental.py#L210) LMAX is 1 instead of 4. This means the sizes of [Ci Cj Ck](https://github.com/graphcore-research/pyscf-ipu/blob/023e11487762731117d73d14bc98b2f74226f800/pyscf_ipu/experimental/integrals.py#L229) become (4*4+1)^3=3375 instead of (1+1)^3=8. **As we discussed, I think...

This (jax trace time not IPU compile time) is what I was looking at for water sparse_eri. If you comment out `factorial`, `binomial` and some of the gamma functions it...

Looks like culprit is here : https://github.com/graphcore-research/pyscf-ipu/blob/023e11487762731117d73d14bc98b2f74226f800/pyscf_ipu/experimental/special.py#L80 TODO: write vertex which implements this function and see if it works.

**Note:** This also reduces FLOPs **Note:** Consider GD wrt density_matrix as in D4FT. If we instead do GD wrt NN(mol)=dm and then DFT_energy(NN(mol)) a single DFT iteration is sufficient to...

Could you add a list of the needed functions? Assuming they are things like `erfc` and `inv_gamma` etc. Do you suspect we'll have to implement them in C?

Write a "--benchmark_tracking True" argparse argument. Goals: - Report cycle count of different molecules (@balancap ). - Keep compile time < 30min. Molecules/hw configuration - Do C20 (single IPU) -eri_threshold...

> I had the idea to introduce a function decorator that will run an operation twice: Do "operation=nanoDFT" or e.g. "operation=einsum(eri, dm)"? I was thinking that we' run nanoDFT twice,...