Performance optimisation for JAX compilation?

[lm1909]

I am using autodiff to optimise the parameters in PTensors via TNOptimizer. In particular I am using JAX as a backend. The problem is that the compilation step of JAX takes quite long. Of course I am aware that there is an inherent tradeoff between compilation time and runtime, but I think the structure of the loss functions (which in the end should evaluate the expectation value of a Hamiltonian for a state created by a circuit) is probably something that is hard to compile for JAX. I am not an expert with JAX so this is something a bit hard to think about for me. JAX occasionally prints warnings that constant folding is slow - my loss functions look very similar to issue #81, in particular this comment:

def loss_for_term(term_key: tuple[int], hamiltonian_term, psi) -> float:
    ket = psi.select(lightcone_tags[term_key], 'any')
    energy_network = ket.gate(hamiltonian_term, term_key, tags='...') | ket.H
    return do('real', energy_network.contract(all, optimize=cotengra_path_optimiser, backend='jax'))
loss_functions = [partial(loss_for_term, k, t) for k, t in hamiltonian_terms.items()]

Any ideas for improving the performance would be much appreciated, thanks in advance!

[jcmgray]

Hi @lm1909. Sadly I don't really have much in the way of tips. I can't find a link, but I recall that jax/XLA has a non-linear compilation time with respect to the number of nodes in the computational graph. That is particularly bad for tensor network stuff, since the computational graph is dynamically generated with often orders of magnitude more nodes than 'normal machine learning model stuff'.

I haven't looked into whether jit'ing blocks of functions at a lower level helps overall. It may be the case for example, that if you have a complicated parametrizing function, that pre compiling it on its own, rather than tracing through it repeatedly helps, but I am not really sure.

Probably one just has to try these out! This is definitely a useful discussion to keep open, in case anyone else has tips though.

[lm1909]

Hi @jcmgray - thanks a lot for the pointer, I will try that out! (and thanks in general for all your work on quimb :) )

[nikitn2]

Is it possible to pre-compute the environments around a tensor or small group of tensors, and just optimise them at a time, one group after another in a sort of DMRG fashion? I'd guess this should probably reduce the number of “computational nodes„ and hence speed up compilation?