Question about indexing using numpy vs jax.numpy inside a JIT compiled function

[fishjojo]

Dear developers,

I have a function as follows

@jit
def amplitudes_to_vector(t1, t2, out=None):
    nocc, nvir = t1.shape
    nov = nocc * nvir
    vector_t1 = t1.ravel()
    vector_t2 = t2.transpose(0,2,1,3).reshape(nov,nov)[numpy.tril_indices(nov)]
    vector = jnp.concatenate((vector_t1, vector_t2), axis=None)
    return vector

where t1.shape=(8,38) and t2.shape=(8,8,38,38).
The JIT compiled function is quite simple and uses a buffer with the size of ~1MB, which is as expected.
However, if I replace numpy.tril_indices with jax.numpy.tril_indices, the complied function becomes very complicated and the buffer size is doubled (see module_0309.jit_amplitudes_to_vector.cpu_after_optimizations-buffer-assignment.txt). What is the cause of this, and should I just use numpy for indexing purposes instead of jax.numpy?

Thank you in advance.

[jakevdp]

Thanks for the question! The difference here is that if you use np.tril_indices, the indices are computed at trace time and embedded in the comutation as a compile-time constant. When you use jnp.tril_indices, the indices are computed from scratch at runtime. This trades off runtime and compile-time performance; also you'll find that embedding large compile-time constants can have other performance implications due to the size of the serialized program.