Cache misses in tensordot

[ngoettin]

So, I have this small test script:

import time
import functools

import jax
import numpy as np
import jax.numpy as jnp

@functools.partial(jax.jit, static_argnames=("inds",))
def test(a, b, inds):
    return jnp.tensordot(a, b, axes=inds)

a = jnp.array(np.random.rand(2,2,2,2))
b = jnp.array(np.random.rand(2,2,2,2))

for qubit_inds in ((0,1), (2,3), (1,3)):
    time.sleep(0.1)
    test(a, b, (tuple(range(-len(qubit_inds), 0)), qubit_inds))

When I profile it (for example with python -m pyinstrument -r html -o res.html test.py) I see that most of the computation time is lost to cache misses. It kind of makes sense, because I give three different qubit_inds, but given that all dimensions are always the same, is there a possibility to speed this up?

[jakevdp]

This is working as expected: each time you change qubit_inds you end up with a new compilation, because you're effectively running a different program that contracts over different axes.

Since all the dimensions are the same, you might be able to reduce the required compilations by using the dynamic_transpose recipe from #8076 (comment) to dynamically transpose the array such that the tensordot always occurs over the same dimensions. This would prevent multiple compilations, but I suspect the runtime will be quite a bit slower, so you'd have to balance that tradeoff.

[ngoettin]

Thanks a lot for the very fast answer! This would explain the behavior I saw in qujax, where I constantly received cache misses I think due to this function, which is executed very often https://github.com/CQCL/qujax/blob/develop/qujax/statetensor.py#L18

I'll have a look at the dynamic_transpose that seems very promising, thanks!