cast_bool_to_primitive breaks TensorFlow graph support.

[AndreasMadsen]

🐛 Bug

To reproduce

import transformers
bert = tf.function(transformers.TFBertForMaskedLM.from_pretrained('bert-base-uncased'))

for i in range(2):
    (_, hidden_state) = bert(tf.constant([[10,11,12]]), output_hidden_states=True)
    print(f'computed {i}')

Errors with

ValueError: not enough values to unpack (expected 2, got 1)

Expected behavior

computed 1
computed 2

Same result as if tf.function was not used.

Environment info

Example environment : https://colab.research.google.com/gist/AndreasMadsen/593df94a3319dee58bba33a26efedeb3/untitled6.ipynb

• transformers version: 3.0.2
• Platform: Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.6.9
• PyTorch version (GPU?): 1.5.1+cu101 (False)
• Tensorflow version (GPU?): 2.2.0 (False)
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

---

Details

The bug happens due to cast_bool_to_primitive, that was introduced in 6e603cb. Before that, it was possible to get the hidden_states from Bert in TensorFlow graph/function mode.

Generally speaking, casting TensorFlow tensors to primitives is not a good practice, as it only works in eager mode. It is also completely unnecessary in this case, as using if bool_tensor_scalar: works perfectly fine.

def print_bool(x):
   if x:
       print('True')
   else:
       print('False')

print_bool_graph = tf.function(print_bool)

print('eager:')
print_bool(True) # Prints True
print_bool(False) # Prints False
print_bool(tf.constant(True)) # Prints True
print_bool(tf.constant(False)) # Prints False

print('')
print('graph:')
print_bool_graph(True) # Prints True
print_bool_graph(False) # Prints False
print_bool_graph(tf.constant(True)) # Prints True
print_bool_graph(tf.constant(False)) # Prints False

I can see there are some cases where defaults are used. The right way to handle that is to implement the default handling upstream in the first call() method. A lesser way would be to implement it as:

def cast_bool_to_primitive(x, default_value=False):
  if x is None:
    return default_value
  return x

[patrickvonplaten]

Hey @AndreasMadsen,

Thanks a lot for the detailed error description! I added this function, so I will try to take a look as soon as possible. I will be off for the next two weeks though (but will not this issue down). If by change @jplu you find some time, feel free to take a look :-)

[jplu]

Should be fixed in the PR #5468

[AndreasMadsen]

Hi @jplu, I'm sorry, but I doubt #5468 will fix the issue. Fundamentally speaking casting to primitives is not a good practice in TensorFlow, as it invalidates the use of @tf.function and is generally unnecessary as described above. Casting to primitives is, in my experience, just never the correct solution in TensorFlow.

I do think #5468 mitigates the issue, which is maybe where the confusion is coming from. This is because, the models will now correctly default to the config object when output_hidden_states=True is not specified as an input. In those cases object property is never cast to a tensor to begin with, therefore the @tf.function graph will be statically compiled to always output the hidden_states, as intended.

However, the behavior is different when output_hidden_states=True is specified as an input, as it will be cast to a Tensor when it becomes part of the inputs argument in call(). After that, it is not possible to convert back to a primitive, as that invalidates @tf.function.

If you insist on keeping it as a primitive, the best solution might be to specify it as an aux-input, similar to training and mask in a keras.layers.Layer, as they don't get converted the same way. I'm not familiar enough with the Keras internals to know the details here, and I think it might also be incompatible with compute_output_shape etc.

BTW, in the keras RNN layers, hidden_state is only specified in the constructor, properly because it can get a bit messy having to specify it in the inputs, but I don't see anything fundamentally wrong with specifying it in inputs.

[jplu]

The PR fix your issue at least for the piece of code you are providing, here the result I have:

>>> import transformers
>>> import tensorflow as tf
>>> bert = tf.function(transformers.TFBertForMaskedLM.from_pretrained('bert-base-uncased'))
>>> for i in range(2):
...   (_, hidden_state) = bert(tf.constant([[10,11,12]]), output_hidden_states=True)
...   print(f'computed {i}')
... 
computed 0
computed 1
>>> 

[foxik]

I believe the solution of @AndreasMadsen is correct -- output_hidden_states and also output_attentions should be passed as named arguments of the call method of TFBertEncoder.call -- that way they are not converted to Tensors and can be just constants.

[AndreasMadsen]

@jplu Sorry for the misunderstanding. The test now works because output_hidden_states is now an auxiliary-input, thus is stays a primitive, thus the casting is no longer involved. However, casting to primitives is still not good practice in TensorFlow, so it doesn't take much to break it.

I read your code more thoroughly, and have the following two failure cases for you.

Case 1:

bert = tf.function(transformers.TFBertForMaskedLM.from_pretrained('bert-base-uncased'))
outputs = bert(tf.constant([[10,11,12]]), output_attentions=True)
assert len(outputs) == 2

Fails with an IndexError. Essentially because casting in a tf.function can not be done. Edit: I think it is an inconsistent casting, in some functions the output_attentions is a primitive and the "casting" works by being a no-op, it other functions the casting can't be done as it is in a tf.function.

Case 2:

bert = tf.function(transformers.TFBertForMaskedLM.from_pretrained('bert-base-uncased'))
outputs = bert(tf.constant([[10,11,12]]), output_hidden_states=tf.constant(True))
assert len(outputs) == 2

outputs only one output (two is exected), because tf.constant can not be casted in a tf.function.

---

However, I would like that, instead of working around these issues, you read the documentation on AutoGraph.

I really think there is a misunderstanding here, about what TensorFlow AutoGraph can do for you and why casting to primitives is really not necessary at all. I would suggest you read https://www.tensorflow.org/guide/function#conditionals and also check out the hidden docs, such as https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#effects-of-the-tracing-process which explains it in more details.

What @foxik says is true, but I think depending on the auxiliary-inputs just avoids the misunderstanding. Truly, casting to primitives is just not a good idea.

[jplu]

Ok, thanks for the hints, I will review that part! I should be able to start working on it from tomorrow or Thursday.

[jplu]

I have started to work on this today. I basically just removed the casting to primitive function and updated the code accordingly. Nevertheless it doesn't work when booleans are tensors (your Case 2) because AutoGraph do not let you return an output of different size on each branch. A simple example to simply the use case can be:

import tensorflow as tf

@tf.function
def output(output1, output2):
    first_conditional_output = ()
    second_conditional_output = ()

    for i in range(2):
        if output1:
            first_conditional_output = first_conditional_output + (i,)

        if output2:
            second_conditional_output = second_conditional_output + (i,)

    outputs = (0,)
    if output1:
        outputs = outputs + (first_conditional_output,)
    if output2:
        outputs = outputs + (second_conditional_output,)

    return outputs

This piece of code works fine as long as the parameters are primitives, but when using tf.constant(True) for at least one of the two makes it fail with:

/opt/anaconda3/envs/transformers/lib/python3.7/site-packages/tensorflow/python/ops/cond_v2.py:633 _make_indexed_slices_indices_types_match
    assert len(set(outs_per_branch)) == 1, outs_per_branch

AssertionError: [1, 0]

I currently struggling on this, and try to find a workaround but if it is not possible to fix this, I think we will just skip this functionality.

[foxik]

@jplu TensorFlow guys internally use tf.cond to handle this case, for example here: https://github.com/tensorflow/tensorflow/blob/2b96f3662bd776e277f86997659e61046b56c315/tensorflow/python/framework/smart_cond.py#L27-L59 . Keras layers like Dropout have another special case for Variables here (the smart_module.smart_cond is the above method): https://github.com/tensorflow/tensorflow/blob/2b96f3662bd776e277f86997659e61046b56c315/tensorflow/python/keras/utils/tf_utils.py#L42-L65

[jplu]

I already came across this, but I don't see how it could solve this issue. Each if in this piece of code is automatically translated to a tf.cond. The problem is that there is no else and then the branches true_fn and false_fn will have a different output size from each other. This is not allowed and this is the problem :(

Unless I'm missing something in what you said.

I'm still trying to figure out how to find a workaround anyway.

[foxik]

@jplu I thought that tf.cond can return arbitrary results from the true_fn and false_fn (contrary to Autograph), but you are right, it is not allowed during graph construction -- my bad.

Personally I think it is enough for the output_hidden_size to be a Python constant (and fail otherwise). It is in fact not obvious what type should the output have in case computation graph is used and the output_hidden_size is undefined Tensor value.

So I think you were right with skipping the functionality ;-)

[AndreasMadsen]

@jplu Good point, regarding the variable output-signature. I think it is perfectly acceptable to assert for a primitive input, at least for now.

Alternatively, the solution would be to return an empty tensor when tf.constant([False]) is used. Such an approach could look like this:

import tensorflow as tf

@tf.function
def output(output1, output2):
    first_conditional_output = tf.TensorArray(tf.int64, size=0, dynamic_size=True, clear_after_read=True)
    second_conditional_output = tf.TensorArray(tf.int64, size=0, dynamic_size=True, clear_after_read=True)

    for i in range(2):
        if output1:
            first_conditional_output = first_conditional_output.write(i, i)
        if output2:
            second_conditional_output = second_conditional_output.write(i, i)

    outputs = (0,)
    if isinstance(output1, tf.Tensor) or output1:
        outputs = outputs + (first_conditional_output.stack(),)

    if isinstance(output2, tf.Tensor) or output2:
        outputs = outputs + (second_conditional_output.stack(),)

    return outputs

edit: PS: there is more documentation here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md

[jplu]

@AndreasMadsen Thanks!! I basically already came across the same solution (which is I think the only one 😄), but the issue here is that we will always have three outputs, which is not really what is wanted. I have to talk to the whole team to see if they are ok with this or not, If yes, we will update the code base accordingly otherwise we will skip this.

[jplu]

I'm trying to figure out if using tf.boolean_mask could work.