search

secondmind-labs / trieste

A Bayesian optimization toolbox built on TensorFlow
Apache License 2.0
212 stars 42 forks source link

Avoid use of python bool in qmc_normal_samples #792

Closed ChrisMorter closed 5 months ago

ChrisMorter commented 9 months ago

Avoid use of python if statement in qmc_normal_samples which prevents it, or any code which calls it, from being saved as a tf.module.

Summary

qmc_normal_samples contains a precondition check to handle the edge cases where the number of samples is zero, or the sample dim is zero. This was implemented using a python if statement, but this means that the function can't be saved as a tf.function.

The error raised if you try this is:

tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: Using a symbolic `tf.Tensor` as a Python `bool` is not allowed: AutoGraph is disabled in this function. Try decorating it directly with @tf.function.

As the error message suggests, this can be worked around by ensuring that any usages of qmc_normal_samples are wrapped in tf.function, but this seems like an unnecessary burden to place on calling code when we can just replace the problematic if statement with tf.cond instead.

...

Fully backwards compatible: yes

PR checklist

khurram-ghani commented 9 months ago

I would like to understand a bit more about the context in which the error occurs please. Are we not using auto-graph when saving the function? If we are, I am a bit surprised that it didn't convert the code to effectively what you have done manually. There are plenty of other places in trieste where we use python if expressions and rely on auto-graph to convert them. So it would be good to understand what is different here. It is concerning if we had to make this kind of change repeatedly.

ChrisMorter commented 9 months ago

I would like to understand a bit more about the context in which the error occurs please. Are we not using auto-graph when saving the function? If we are, I am a bit surprised that it didn't convert the code to effectively what you have done manually. There are plenty of other places in trieste where we use python if expressions and rely on auto-graph to convert them. So it would be good to understand what is different here. It is concerning if we had to make this kind of change repeatedly.

So the issue occurs when code outside of Trieste attempts to save models which call into Trieste with autograph=False. I'm slightly surprised we haven't hit this issue sooner, because it seems like it would happen whenever the calling code encounters an if expression. Maybe we've just been lucky so far due to mainly calling prediction methods which don't tend to have much branching logic.

In this case, I'm not sure why the calling code doesn't use autograph. Perhaps the resolution should be to leave Trieste as is it and to update the callers to use autograph. I'll look into how feasible that is... I agree with your point that if we wish for Trieste functions to generally support saving without using autograph then we'd need to update many places, which would be a pain.

hstojic commented 8 months ago

@ChrisMorter and @khurram-ghani do we have a resolution on this topic?

khurram-ghani commented 8 months ago

@ChrisMorter has updated the calling code to always use autograph and that has resolved the original issue; so I believe this trieste change is not needed anymore. @ChrisMorter are you okay with closing this PR?

ChrisMorter commented 5 months ago

@ChrisMorter has updated the calling code to always use autograph and that has resolved the original issue; so I believe this trieste change is not needed anymore. @ChrisMorter are you okay with closing this PR?

@khurram-ghani apologies for the slow reply, I must have missed the notification for your message. Yes, that's correct, I'm happy to close this PR.