Performance issue in src/vegasflow/vflow.py

[DLPerf]

Hello! Our static bug checker has found a performance issue in src/vegasflow/vflow.py: refine_grid_per_dimension is repeatedly called in a for loop, but there are tf.function decorated functions while_check and while_body defined and called in refine_grid_per_dimension.

In that case, when refine_grid_per_dimension is called in a loop, the functions while_check and while_body will create new graphs every time, and that can trigger tf.function retracing warning.

Here is the tensorflow document to support it.

Briefly, for better efficiency, it's better to use:

@tf.function
def inner():
    pass

def outer():
    inner()  

than:

def outer():
    @tf.function
    def inner():
        pass
    inner()

Looking forward to your reply. Btw, I am glad to create a PR to fix it if you are too busy.

[scarlehoff]

Hi, if you create a pr i'll be happy.

It's not (or should not be) a performance bottleneck because the grid is only refined at the end and only during warmup so if the code gets much more complicated don't worry about it. If instead it is clean and faster then I'll be of course happy to merge it!