Adding additional layers to TFHubertModel throws OperatorNotAllowedInGraphError
Author: V-Sher
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.15.0
- Platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.12
- PyTorch version (GPU?): 1.10.0+cu111 (True)
- Tensorflow version (GPU?): 2.7.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help

@patrickvonplaten, @anton-l

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- ALBERT, BERT, XLM, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- T5, BART, Marian, Pegasus, EncoderDecoder: @patrickvonplaten
- Blenderbot, MBART: @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet, BigBird: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder, UniSpeech, UniSpeechSAT, SEW, SEW-D, Speech2Text: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten @narsil
- Tokenizers: @LysandreJik
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): [TFHubertModel](https://huggingface.co/docs/transformers/model_doc/hubert#transformers.TFHubertModel)

## To reproduce

Steps to reproduce the behavior:
Just a simple code snippet that loads a pretrained TFHubertModel and adds - (1) a lambda layer to sum over the hidden units obtained from the Hubert model (for example: from `(N, 120,1024)` -> `(N,1024)` and (2) dense & dropout layers

```import librosa
import tensorflow as tf
import torch
import numpy as np
from tensorflow.keras.optimizers import Adam
from transformers import TFHubertModel

def create_model(bert_model, dim):
  input_ids = tf.keras.Input(shape=(dim,),dtype='int32')
  attention_masks = tf.keras.Input(shape=(dim,),dtype='int32')
  
  output = bert_model([input_ids,attention_masks])
  output = output[0] 
  output = tf.keras.layers.Lambda(lambda x: tf.keras.backend.sum(x, axis=1), name = "Pooling_Embs")(output)
  output = tf.keras.layers.Dense(32,activation='relu')(output)
  output = tf.keras.layers.Dropout(0.2)(output)

  output = tf.keras.layers.Dense(1,activation='sigmoid')(output)
  model = tf.keras.models.Model(inputs = [input_ids,attention_masks],outputs = output)
  model.compile(Adam(learning_rate=1e-6), loss='binary_crossentropy', metrics=['accuracy'])
  return model

# custom model creation
hubert_model = TFHubertModel.from_pretrained('facebook/hubert-large-ls960-ft')
model = create_model(hubert_model, dim=38744)
model.summary()
```
The model compiles just fine without any error:
```
Model: "model_2"
__________________________________________________________________________________________________
 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 input_5 (InputLayer)           [(None, 38744)]      0           []                               
                                                                                                  
 input_6 (InputLayer)           [(None, 38744)]      0           []                               
                                                                                                  
 tf_hubert_model (TFHubertModel  TFBaseModelOutput(l  315438720  ['input_5[0][0]',                
 )                              ast_hidden_state=(N               'input_6[0][0]']                
                                one, 120, 1024),                                                  
                                 hidden_states=None                                               
                                , attentions=None)                                                
                                                                                                  
 Pooling_Embs (Lambda)          (None, 1024)         0           ['tf_hubert_model[1][0]']        
                                                                                                  
 dense_4 (Dense)                (None, 32)           32800       ['Pooling_Embs[0][0]']           
                                                                                                  
 dropout_173 (Dropout)          (None, 32)           0           ['dense_4[0][0]']                
                                                                                                  
 dense_5 (Dense)                (None, 1)            33          ['dropout_173[0][0]']            
                                                                                                  
==================================================================================================
Total params: 315,471,553
Trainable params: 315,471,553
Non-trainable params: 0
```
When I try to fit the model, the error is thrown (dummy inputs/attention masks used here are for demonstration purposes only, ideally they will come from passing audio through a feature extractor like `Wav2Vec2FeatureExtractor`):
```
# fit model
input_values = np.random.rand(5,38744)
attention_masks = np.random.randint(0,2, size=(5,38744))
labels = np.asarray([0, 1, 0, 0, 1])

model.fit([input_values,attention_masks], 
          labels,
          epochs=2,
          batch_size=2)
```
The error thrown:
```__________________________________________________________________________________________________
Epoch 1/2
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
<ipython-input-24-3a4ee842c0f2> in <module>()
      4 input_values = np.random.rand(5,38744)
      5 attention_masks = np.random.randint(0,2, size=(5,38744))
----> 6 history = model.fit([input_values,attention_masks],np.asarray(label),epochs=2,batch_size=2)

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py in autograph_handler(*args, **kwargs)
   1127           except Exception as e:  # pylint:disable=broad-except
   1128             if hasattr(e, "ag_error_metadata"):
-> 1129               raise e.ag_error_metadata.to_exception(e)
   1130             else:
   1131               raise

OperatorNotAllowedInGraphError: in user code:

    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 878, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 867, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 860, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 808, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
        raise e.with_traceback(filtered_tb) from None

    OperatorNotAllowedInGraphError: Exception encountered when calling layer "tf_hubert_model" (type TFHubertModel).
    
    in user code:
    
        File "/usr/local/lib/python3.7/dist-packages/transformers/models/hubert/modeling_tf_hubert.py", line 1453, in call  *
            outputs = self.hubert(
        File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
    
        OperatorNotAllowedInGraphError: Exception encountered when calling layer "hubert" (type TFHubertMainLayer).
        
        in user code:
        
            File "/usr/local/lib/python3.7/dist-packages/transformers/models/hubert/modeling_tf_hubert.py", line 1237, in call  *
                hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)
            File "/usr/local/lib/python3.7/dist-packages/transformers/models/hubert/modeling_tf_hubert.py", line 1168, in _mask_hidden_states  *
                mask_time_indices = _compute_mask_indices(
            File "/usr/local/lib/python3.7/dist-packages/transformers/models/hubert/modeling_tf_hubert.py", line 229, in _compute_mask_indices  *
                num_masked_spans = max(num_masked_spans, min_masks)
        
            OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.
        
        
        Call arguments received:
          • input_values=tf.Tensor(shape=(None, 38744), dtype=int32)
          • attention_mask=tf.Tensor(shape=(None, 38744), dtype=int32)
          • token_type_ids=None
          • position_ids=None
          • head_mask=None
          • inputs_embeds=None
          • output_attentions=False
          • output_hidden_states=False
          • return_dict=True
          • training=True
          • kwargs=<class 'inspect._empty'>
    
    
    Call arguments received:
      • input_values=['tf.Tensor(shape=(None, 38744), dtype=int32)', 'tf.Tensor(shape=(None, 38744), dtype=int32)']
      • attention_mask=None
      • token_type_ids=None
      • position_ids=None
      • head_mask=None
      • inputs_embeds=None
      • output_attentions=None
      • output_hidden_states=None
      • return_dict=None
      • training=True
```
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

The final aim is to get the model up and running so that a voice liveliness detection system (i.e. whether the audio is live or a replayed one) can be trained.
## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
The model should fit without any error, similar to the one in [this notebook](https://www.kaggle.com/dhruv1234/huggingface-tfbertmodel) where they did the same as above but for the TFBertModel.

## Additional Information
I have already tried removing the lambda layer just to see if that helps but the error persists.

Comments:
It seems that we have to change https://github.com/huggingface/transformers/blob/cc406da4debea3c2e7e93f42585339e79e18fd6b/src/transformers/models/hubert/modeling_tf_hubert.py#L231 

to use a tensorflow max function instead https://www.tensorflow.org/api_docs/python/tf/math/maximum

@gante , @Rocketknight1 - would you guys be interested in taking a look here?

@patrickvonplaten I can take a look, assigning to me

@patrickvonplaten replacing `max` by `tf.math.maximum` then exposes an error on `_scatter_values_on_batch_indices` due to an unknown batch size in TF's graph mode. However, forcing eager execution at train time (i.e. compiling the model with `run_eagerly=True`) bypasses all these errors, but then the model hits a shape-related issue in a Conv layer -- see below.

Running a modified version of the script kindly provided by @V-Sher 
```python
import tensorflow as tf
import numpy as np
from tensorflow.keras.optimizers import Adam
from transformers import TFHubertModel

def create_model(bert_model, dim):
  input_ids = tf.keras.Input(shape=(dim,),dtype='int32')
  attention_masks = tf.keras.Input(shape=(dim,),dtype='int32')
  output = bert_model([input_ids,attention_masks])
  output = output[0]
  output = tf.keras.layers.Lambda(lambda x: tf.keras.backend.sum(x, axis=1), name = "Pooling_Embs")(output)
  output = tf.keras.layers.Dense(32,activation='relu')(output)
  output = tf.keras.layers.Dropout(0.2)(output)
  output = tf.keras.layers.Dense(1, activation='sigmoid')(output)
  model = tf.keras.models.Model(inputs = [input_ids,attention_masks],outputs = output)
  model.compile(Adam(learning_rate=1e-6), loss='binary_crossentropy', metrics=['accuracy'], run_eagerly=True)
  return model

# custom model creation
hubert_model = TFHubertModel.from_pretrained('facebook/hubert-large-ls960-ft')
model = create_model(hubert_model, dim=38744)
model.summary()

# fit model
input_values = np.random.rand(5,38744)
attention_masks = np.random.randint(0,2, size=(5,38744))
labels = np.asarray([0, 1, 0, 0, 1])

print(model([input_values,attention_masks]))
model.fit([input_values,attention_masks],
          labels,
          epochs=2,
          batch_size=2)
```

We get
```
Epoch 1/2
Traceback (most recent call last):
  File "test.py", line 30, in <module>
    model.fit([input_values,attention_masks],
  File "/home/joao_huggingface_co/hf/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/joao_huggingface_co/hf/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 7107, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Computed input depth 1024 doesn't match filter input depth 64 [Op:Conv2DBackpropInput]
```

Before I dig deeper @patrickvonplaten @Rocketknight1, any clues?

@gante I don't have any immediate intuition about what the problem is here, sorry! If you get stuck or don't have time to keep working on it though, feel free to ping me and I'll dig in deeper.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

(not stale, next on my todo list)

Update: the script I shared above works on GPU, but fails on CPU due to upstream problems in Keras -- TF does not support backpropagation of grouped convolutions on CPU. (@V-Sher)

The PR mentioned above adds an informative error message for these situations.