Tensorflow: Shape must be at least rank 3 but is rank 2

Question

I am training a bilstm model for text classification and running into the following error

ValueError: Shape must be at least rank 3 but is rank 2 for '{{node bidirectional_5/forward_lstm_6/lstm_cell_17/BiasAdd}} = BiasAdd[T=DT_HALF, data_format="NCHW"](bidirectional_5/forward_lstm_6/lstm_cell_17/MatMul, bidirectional_5/forward_lstm_6/lstm_cell_17/split_1)' with input shapes: [?,200], [200].

This is the code I am using

output_dim = 200
dropout = 0.2
recurrent_dropout = 0.2
lstm_units = 128
n_tags = 26
num_dense_neurons = 20
classifier_activation = time_activation = 'tanh'
learning_rate = 0.003
merge_mode = 'concat'
activation = 'softmax'
max_length = 20

model = Sequential()
# Add Embedding layer
e = Embedding(input_dim= 340862, output_dim=output_dim, 
          weights=[weight_matrix], 
          input_length=max_length, trainable=False)
model.add(e)
# Add bidirectional LSTM
b = Bidirectional(LSTM(units=output_dim, return_sequences=True, dropout=dropout, recurrent_dropout=recurrent_dropout), merge_mode = merge_mode)
model.add(b)
# Add LSTM
model.add(LSTM(units=lstm_units, return_sequences=True, dropout=dropout, recurrent_dropout=recurrent_dropout))
# Add timeDistributed Layer
model.add(TimeDistributed(Dense(n_tags, activation=time_activation)))
model.add(Flatten())
model.add(Dense(num_dense_neurons, activation=classifier_activation))
model.add(Dense(n_tags,activation='softmax'))
#Optimiser 
adam = Adam(learning_rate=learning_rate, beta_1=0.9, beta_2=0.999)
# Compile model
model.compile(loss='categorical_crossentropy', optimizer=adam, metrics=[f1_score])
model.summary()

The weight_matrix I am passing to the embedding layer is of shape (340862, 200)

Both keras and tensorflow versions are 2.6.0. I am running this on AWS Sagemaker. Not sure why this is happening.

Full error

InvalidArgumentError                      Traceback (most recent call last)
~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/framework/ops.py in _create_c_op(graph, node_def, inputs, control_inputs, op_def)
   1879   try:
-> 1880     c_op = pywrap_tf_session.TF_FinishOperation(op_desc)
   1881   except errors.InvalidArgumentError as e:

InvalidArgumentError: Shape must be at least rank 3 but is rank 2 for '{{node bidirectional_6/forward_lstm_7/lstm_cell_20/BiasAdd}} = BiasAdd[T=DT_HALF, data_format="NCHW"](bidirectional_6/forward_lstm_7/lstm_cell_20/MatMul, bidirectional_6/forward_lstm_7/lstm_cell_20/split_1)' with input shapes: [?,200], [200].

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-88-ddcbc617161c> in <module>
     20 # Add bidirectional LSTM
     21 b = Bidirectional(LSTM(units=output_dim, return_sequences=True, dropout=dropout, recurrent_dropout=recurrent_dropout), merge_mode = merge_mode)
---> 22 model.add(b)
     23 # Add LSTM
     24 model.add(LSTM(units=lstm_units, return_sequences=True, dropout=dropout, recurrent_dropout=recurrent_dropout))

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    528     self._self_setattr_tracking = False  # pylint: disable=protected-access
    529     try:
--> 530       result = method(self, *args, **kwargs)
    531     finally:
    532       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/sequential.py in add(self, layer)
    215       # If the model is being built continuously on top of an input layer:
    216       # refresh its output.
--> 217       output_tensor = layer(self.outputs[0])
    218       if len(tf.nest.flatten(output_tensor)) != 1:
    219         raise ValueError(SINGLE_LAYER_OUTPUT_ERROR_MSG)

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/layers/wrappers.py in __call__(self, inputs, initial_state, constants, **kwargs)
    581 
    582     if initial_state is None and constants is None:
--> 583       return super(Bidirectional, self).__call__(inputs, **kwargs)
    584 
    585     # Applies the same workaround as in `RNN.__call__`

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    975     if _in_functional_construction_mode(self, inputs, args, kwargs, input_list):
    976       return self._functional_construction_call(inputs, args, kwargs,
--> 977                                                 input_list)
    978 
    979     # Maintains info about the `Layer.call` stack.

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/base_layer.py in _functional_construction_call(self, inputs, args, kwargs, input_list)
   1113       # Check input assumptions set after layer building, e.g. input shape.
   1114       outputs = self._keras_tensor_symbolic_call(
-> 1115           inputs, input_masks, args, kwargs)
   1116 
   1117       if outputs is None:

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/base_layer.py in _keras_tensor_symbolic_call(self, inputs, input_masks, args, kwargs)
    846       return tf.nest.map_structure(keras_tensor.KerasTensor, output_signature)
    847     else:
--> 848       return self._infer_output_signature(inputs, args, kwargs, input_masks)
    849 
    850   def _infer_output_signature(self, inputs, args, kwargs, input_masks):

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/base_layer.py in _infer_output_signature(self, inputs, args, kwargs, input_masks)
    886           self._maybe_build(inputs)
    887           inputs = self._maybe_cast_inputs(inputs)
--> 888           outputs = call_fn(inputs, *args, **kwargs)
    889 
    890         self._handle_activity_regularization(inputs, outputs)

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/layers/wrappers.py in call(self, inputs, training, mask, initial_state, constants)
    697 
    698       y = self.forward_layer(forward_inputs,
--> 699                              initial_state=forward_state, **kwargs)
    700       y_rev = self.backward_layer(backward_inputs,
    701                                   initial_state=backward_state, **kwargs)

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/layers/recurrent.py in __call__(self, inputs, initial_state, constants, **kwargs)
    657 
    658     if initial_state is None and constants is None:
--> 659       return super(RNN, self).__call__(inputs, **kwargs)
    660 
    661     # If any of `initial_state` or `constants` are specified and are Keras

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1035         with autocast_variable.enable_auto_cast_variables(
   1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
   1038 
   1039         if self._activity_regularizer:

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/layers/recurrent_v2.py in call(self, inputs, mask, training, initial_state)
   1165           input_length=row_lengths if row_lengths is not None else timesteps,
   1166           time_major=self.time_major,
-> 1167           zero_output_for_mask=self.zero_output_for_mask)
   1168       runtime = _runtime(_RUNTIME_UNKNOWN)
   1169     else:

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    204     """Call target, and fall back on dispatchers if there is a TypeError."""
    205     try:
--> 206       return target(*args, **kwargs)
    207     except (TypeError, ValueError):
    208       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/backend.py in rnn(step_function, inputs, initial_states, go_backwards, mask, constants, unroll, input_length, time_major, zero_output_for_mask)
   4353     # the value is discarded.
   4354     output_time_zero, _ = step_function(
-> 4355         input_time_zero, tuple(initial_states) + tuple(constants))
   4356     output_ta = tuple(
   4357         tf.TensorArray(

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/layers/recurrent_v2.py in step(inputs, states)
   1153 
   1154       def step(inputs, states):
-> 1155         return self.cell(inputs, states, **kwargs)
   1156 
   1157       last_output, outputs, states = backend.rnn(

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1035         with autocast_variable.enable_auto_cast_variables(
   1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
   1038 
   1039         if self._activity_regularizer:

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/layers/recurrent.py in call(self, inputs, states, training)
   2440         b_i, b_f, b_c, b_o = tf.split(
   2441             self.bias, num_or_size_splits=4, axis=0)
-> 2442         x_i = backend.bias_add(x_i, b_i)
   2443         x_f = backend.bias_add(x_f, b_f)
   2444         x_c = backend.bias_add(x_c, b_c)

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    204     """Call target, and fall back on dispatchers if there is a TypeError."""
    205     try:
--> 206       return target(*args, **kwargs)
    207     except (TypeError, ValueError):
    208       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/keras/backend.py in bias_add(x, bias, data_format)
   5973   if len(bias_shape) == 1:
   5974     if data_format == 'channels_first':
-> 5975       return tf.nn.bias_add(x, bias, data_format='NCHW')
   5976     return tf.nn.bias_add(x, bias, data_format='NHWC')
   5977   if ndim(x) in (3, 4, 5):

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    204     """Call target, and fall back on dispatchers if there is a TypeError."""
    205     try:
--> 206       return target(*args, **kwargs)
    207     except (TypeError, ValueError):
    208       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/ops/nn_ops.py in bias_add(value, bias, data_format, name)
   3495     else:
   3496       return gen_nn_ops.bias_add(
-> 3497           value, bias, data_format=data_format, name=name)
   3498 
   3499 

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/ops/gen_nn_ops.py in bias_add(value, bias, data_format, name)
    689   data_format = _execute.make_str(data_format, "data_format")
    690   _, _, _op, _outputs = _op_def_library._apply_op_helper(
--> 691         "BiasAdd", value=value, bias=bias, data_format=data_format, name=name)
    692   _result = _outputs[:]
    693   if _execute.must_record_gradient():

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py in _apply_op_helper(op_type_name, name, **keywords)
    748       op = g._create_op_internal(op_type_name, inputs, dtypes=None,
    749                                  name=scope, input_types=input_types,
--> 750                                  attrs=attr_protos, op_def=op_def)
    751 
    752     # `outputs` is returned as a separate return value so that the output

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in _create_op_internal(self, op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_device)
    599     return super(FuncGraph, self)._create_op_internal(  # pylint: disable=protected-access
    600         op_type, captured_inputs, dtypes, input_types, name, attrs, op_def,
--> 601         compute_device)
    602 
    603   def capture(self, tensor, name=None, shape=None):

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/framework/ops.py in _create_op_internal(self, op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_device)
   3567           input_types=input_types,
   3568           original_op=self._default_original_op,
-> 3569           op_def=op_def)
   3570       self._create_op_helper(ret, compute_device=compute_device)
   3571     return ret

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/framework/ops.py in __init__(self, node_def, g, inputs, output_types, control_inputs, input_types, original_op, op_def)
   2040         op_def = self._graph._get_op_def(node_def.op)
   2041       self._c_op = _create_c_op(self._graph, node_def, inputs,
-> 2042                                 control_input_ops, op_def)
   2043       name = compat.as_str(node_def.name)
   2044 

~/anaconda3/envs/JupyterSystemEnv/lib/python3.7/site-packages/tensorflow/python/framework/ops.py in _create_c_op(graph, node_def, inputs, control_inputs, op_def)
   1881   except errors.InvalidArgumentError as e:
   1882     # Convert to ValueError for backwards compatibility.
-> 1883     raise ValueError(str(e))
   1884 
   1885   return c_op

ValueError: Shape must be at least rank 3 but is rank 2 for '{{node bidirectional_6/forward_lstm_7/lstm_cell_20/BiasAdd}} = BiasAdd[T=DT_HALF, data_format="NCHW"](bidirectional_6/forward_lstm_7/lstm_cell_20/MatMul, bidirectional_6/forward_lstm_7/lstm_cell_20/split_1)' with input shapes: [?,200], [200].