enhance const_fold_optimizer: supports to fold const with Cast

Author: zhijxu-MS

tests/test_optimizers.py

@@ -503,6 +503,25 @@ def test_const_fold_unsqueeze_with_const(self):
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
         self.run_and_compare(["res"], {"X": np.random.randn(1).astype(np.float32)}, model_proto,
                              "Unsqueeze", 0)
+
+    def test_const_fold_cast_with_const(self):
+        shape = (6, 6)
+        const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
+                                          vals=np.random.randn(*shape).flatten().astype(np.float32))
+        node1 = helper.make_node("Constant", [], ["const"], value=const_tensor)
+        node2 = helper.make_node("Cast", ["const"], ["value1"], to=TensorProto.INT64)
+        node3 = helper.make_node("Add", ["value1", "X"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "test_const_fold_cast_with_const",
+            [helper.make_tensor_value_info("X", TensorProto.INT64, shape)],
+            [helper.make_tensor_value_info("res", TensorProto.INT64, shape)],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["res"], {"X": np.random.randn(*shape).astype(np.int64)}, model_proto,
+                             "Cast", 0)
     # Const Fold Optimizer Tests End
 
 

tf2onnx/optimizer/const_fold_optimizer.py

@@ -92,6 +92,14 @@ def _replace_node_with_const(node, graph, vals):
             graph.replace_all_inputs(graph.get_nodes(), old_input, const_node.output[0])
         graph.remove_node(node.name)
 
+    @staticmethod
+    @_register_func("Cast")
+    def _fold_cast(node, graph):
+        const_val = node.inputs[0].get_tensor_value(as_list=False)
+        np_dtype = utils.ONNX_TO_NUMPY_DTYPE[node.get_attr("to").i]
+        const_val_after_cast = const_val.astype(np_dtype)
+        return [const_val_after_cast]
+
     @staticmethod
     @_register_func("Transpose")
     def _fold_transpose(node, graph) -> list: