tf.debugging & tests for tf.function

[hstojic]

This is a follow up on discussion in PR #299 - at the moment we use tf.function for optimizing the model, but users might pass any other piece of code through it e.g. whole BO loop through it. Hence, it would be good to have some tests that (all of) our code can be used within tf.function without errors. Most likely we would have errors, since there is a lot of python code can cause errors within tf.function. This is not a high priority probably, but it could get out of hand with time so perhaps at least we could start testing any new piece of code that we want to merge into develop?

Related to this, we should use tf.debugging everywhere in the code instead of if-else statements with some raise statement that we dominantly use at the moment.

[joelberkeley]

the situation is actually much worse than I realised re raise and tf.function. The following test passes

def test_raise_is_incompatible_with_tf_function() -> None:
    @tf.function
    def f(a: tf.Tensor) -> tf.Tensor:
        if a <= tf.constant(0):
            raise ValueError

        return a

    with pytest.raises(ValueError):
        f(tf.constant(1))

which suggests that the ValueError is raised even though the value shouldn't execute the error branch (perhaps due to tracing covering every branch). This suggests that any function that has any standard python exceptions in it cannot work for any arguments with tf.function.

[hstojic]

thanks for making this concrete, this seems quite indicative about exceptions

[uri-granta]

(connected` to #722)

[uri-granta]

Bumping this ticket in light of #880.