Compilation with jax backend and dynamic network structures

[MSRudolph]

Hello TensorNetwork Team and community 👋

I have a question regarding compilation using the jax backend. In my particular use-case, I have a tensor network that I connect differently for every contraction. I have noticed that this makes runtime slower with the jax backend than with the numpy backend. Using the snakeviz profiler, I can see that the library compiles with every new step which makes me think that you are jit compiling the contraction of a specific network. The result is that even when I jit compile the entire sequence of different network connectivities and contractions on a higher level, the resulting performance is only on par with the numpy backend. Is there a way to turn off the native compilation, or do you have an idea how I can still leverage it?

Manuel