Fix CI test errors introduced by TF 2.4.0 (onnx#842)

Fix CI test errors introduced by TF 2.4.0 including
lstm: test_model.py
gru: test_model.py
rnn: test_model.py
mod: mod.py

Signed-off-by: Chin Huang <[email protected]>

Author: chinhuang007

onnx_tf/handlers/backend/mod.py

@@ -15,10 +15,11 @@ class Mod(ArithmeticMixin, BackendHandler):
       tf.uint16: tf.int32,
       tf.uint32: tf.int64,
       tf.int8: tf.int32,
-      tf.int16: tf.int32
+      tf.int16: tf.int32,
+      tf.float16: tf.float32
   }
   supported_types = [
-      tf.int32, tf.int64, tf.float16, tf.float32, tf.float64, tf.bfloat16
+      tf.int32, tf.int64, tf.float32, tf.float64, tf.bfloat16
   ]
 
   @classmethod

test/backend/test_model.py

@@ -55,11 +55,6 @@ def test_lstm_savedmodel(self):
             helper.make_tensor_value_info("Y_c", TensorProto.FLOAT, [1, 2, 3])
         ])
 
-    # prepare the ONNX model and save it as a Tensorflow SavedModel
-    tf_rep = prepare(helper.make_model(graph_def))
-    model_path = "lstm_savedmodel"
-    tf_rep.export_graph(model_path)
-
     # Initializing Inputs
     X = np.array([[[1., 2., 3., 4.], [5., 6., 7., 8.]]]).astype(np.float32)
     W = weight_scale * np.ones((1, number_of_gates * hidden_size, input_size)).astype(np.float32)
@@ -68,6 +63,12 @@ def test_lstm_savedmodel(self):
     seq_lens = np.repeat(X.shape[0], X.shape[1]).astype(np.int32)
     init_h = weight_scale * np.ones((1, X.shape[1], hidden_size)).astype(np.float32)
 
+    # prepare the ONNX model and save it as a Tensorflow SavedModel
+    tf_rep = prepare(helper.make_model(graph_def))
+    tf_rep.run({"X": X, "W": W, "R": R, "B": B, "sequence_lens": seq_lens, "initial_h": init_h })
+    model_path = "lstm_savedmodel"
+    tf_rep.export_graph(model_path)
+
     # use the ONNX reference implementation to get expected output
     lstm = LSTM_Helper(X=X, W=W, R=R, B=B, initial_h=init_h)
     _, Y_ref = lstm.step()
@@ -108,10 +109,6 @@ def test_gru_savedmodel(self):
             helper.make_tensor_value_info("Y_h", TensorProto.FLOAT, [1, 3, 3])
         ])
 
-    tf_rep = prepare(helper.make_model(graph_def))
-    model_path = "gru_savedmodel"
-    tf_rep.export_graph(model_path)
-
     # initializing Inputs
     X = np.array([[[1., 2., 3.], [4., 5., 6.], [7., 8.,
                                                 9.]]]).astype(np.float32)
@@ -124,6 +121,12 @@ def test_gru_savedmodel(self):
     R_B = np.zeros((1, number_of_gates * hidden_size)).astype(np.float32)
     B = np.concatenate((W_B, R_B), axis=1)
 
+    # prepare the ONNX model and save it as a Tensorflow SavedModel
+    tf_rep = prepare(helper.make_model(graph_def))
+    tf_rep.run({"X": X, "W": W, "R": R, "B": B})
+    model_path = "gru_savedmodel"
+    tf_rep.export_graph(model_path)
+
     # use the ONNX reference implementation to get the expected output
     gru = GRU_Helper(X=X, W=W, R=R, B=B)
     _, Y_ref = gru.step()
@@ -161,15 +164,17 @@ def test_rnn_savedmodel(self):
             helper.make_tensor_value_info("Y_h", TensorProto.FLOAT, [1, 2, 4])
         ])
 
-    tf_rep = prepare(helper.make_model(graph_def))
-    model_path = "rnn_savedmodel"
-    tf_rep.export_graph(model_path)
-
     # initializing Inputs
     X = np.array([[[1., 2.], [3., 4.], [5., 6.]]]).astype(np.float32)
     W = weight_scale * np.ones((1, hidden_size, input_size)).astype(np.float32)
     R = weight_scale * np.ones((1, hidden_size, hidden_size)).astype(np.float32)
 
+    # prepare the ONNX model and save it as a Tensorflow SavedModel
+    tf_rep = prepare(helper.make_model(graph_def))
+    tf_rep.run({"X": X, "W": W, "R": R})
+    model_path = "rnn_savedmodel"
+    tf_rep.export_graph(model_path)
+
     # use the ONNX reference implementation to get the expected output
     rnn = RNN_Helper(X=X, W=W, R=R)
     _, Y_ref = rnn.step()