[Python] Function initialization issue using jax for flexible variables

The first run of the function using jax requires significantly more time that subsequent calls. It appears because during the first time this function is initialized. The tricky thing is if we change the size of the variable in the function, it will be compiled again.

I am writing a back testing tool with jax and during the calculations I use different size matrixes (because of the different amount of business day for each specific month). Do you know best practices how to handle with it?

Maybe I need to determine the length of the variable and if the current matrix is smaller, just fill it with NaNs?

Continue reading...