Memory leak with TF.function

[arthurpham]

https://colab.research.google.com/github/arthurpham/google_colab/blob/e20469359e82ecf96640837e2022f58de0df8314/AmericanOption_MC_MemoryLeak_TQF.ipynb

I've tried to reuse the jupyter notebook that show how to do a MC pricing.

for i in trange(1, 100000):
    strikes = tf.constant(np.random.rand(i) * (0.5*spot) + spot, dtype )
    price_eu_options(strikes, spot, sigma)

The test is slightly silly but if the application that calls the pricing library might chose to send an arbitrary strike size, so the batch size might not be known in advance, so i'm wondering what is the best way to handle that without asking quants to micromanage the shapes of each inputs (as a change in the shape of one input retrigger the graph construction) ?

Thank you

[cyrilchim]

Hi Arthur,
Sorry for the late reply. The issue you are having with memory is that tf.function in the colab is being compiled for static shapes. If you want to fix that, use input_signature arg:

price_eu_options = tf.function(
    set_up_pricer(expiries), 
    input_signature=[tf.TensorSpec([None], dtype=tf.float64),
                     tf.TensorSpec([], dtype=tf.float64),
                     tf.TensorSpec([], dtype=tf.float64)])

This ensures that strikes can be of any shape of rank 1. I tried this and it seems to work. It always makes sense to provide function signature and give names to TensorSpec if you are going to deploy the models.

Keep in mind, at the moment XLA compilation expects static shapes. The last time I checked TF recompiles your graph for every shape, meaning you are better off dealing with static shapes

Please let me know if this makes sense

[arthurpham]

Yes thank you, i didn't know there was a concept of None dimension to handle varying-size inputs.
I believe this is an important nuance that might mean that in the coding standard, an input signature should always be provided with a tf.function for a quant application as inputs can change in sizes and trigger multiple retracing if that input_signature is not given.