Hi all, I’ve got a question regarding memoization. I’m setting up a Monte Carlo simulation that involves computing some information about a design matrix which returns a TypedDict, which is then passed to my function that does the random sampling.
The computation on the design matrix takes much longer than the random sampling (depending on the size of the design matrix, but even in small samples it’s 10x slower), so my current solution is to have two separate jitted functions, one for making the TypedDict (and has some form of memoization) which is returned to Python, which is then immediately passed into the other function. I expect that this all would run faster if I could combine all of these things into a single jitted function. Given that the precision of the simulation is increased with more resamples, I’m really looking for any speed gains possible.
My question is - is there a way to memoize a jitted function that is called from another jitted function?
Here’s a minimal example.
#Current solution: @nb.jit def make_dict(data): #do some stuff return dict #this is a dict of (X,Y) tuples to 1D arrays of ints, but the arrays are not necessarily the same length @nb.jit def random_sampling(data, dict): #do some more stuff return resampled_data def monte_carlo(data): dict = make_dict(data) #cache this result resampled_data = random_sampling(data, dict) return resampled_data #what i would like to do @nb.jit def random_sampling(data): dict = make_dict(data) ###with memoization #do some stuff return resampled_data