Added unit test

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1213,6 +1213,22 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val1})
 
+    def test_slice_from_shape_const_fold(self):
+        x_val = np.array([4, 3], dtype=np.int64)
+        x_shape = np.array([-1, 3], dtype=np.int64)
+        def func(x):
+            z = tf.zeros(x)
+            x = tf.reshape(z, tf.constant(x_shape))
+            s = tf.shape(x)
+            t1 = tf.constant([1], dtype=tf.int32)
+            t2 = tf.constant([2], dtype=tf.int32)
+            y = tf.strided_slice(s, t1, t2, shrink_axis_mask=1)
+            return tf.identity(y, name=_TFOUTPUT)
+        def graph_validator(g):
+            # After constant folding just an input and const output node remain
+            return len(g.get_nodes()) == 2
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, graph_validator=graph_validator)
+
     def test_slice(self):
         x_val = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=np.float32)
         def func(x):

tf2onnx/tf_utils.py

@@ -20,7 +20,7 @@
 
 from onnx import helper, onnx_pb, numpy_helper
 
-from tf2onnx.utils import make_sure, is_tf_const_op, port_name
+from tf2onnx.utils import make_sure, is_tf_const_op, port_name, map_onnx_to_numpy_type
 from . import logging
 
 logger = logging.getLogger(__name__)
@@ -166,10 +166,11 @@ def get_index_from_strided_slice_of_shape(node, outputs_to_values):
         return None
     return i1
 
-def compute_const_folding_using_tf(g, const_node_values):
+def compute_const_folding_using_tf(g, const_node_values, graph_outputs):
     """Find nodes with constant inputs and compute their values using TF"""
     if const_node_values is None:
         const_node_values = {}
+    graph_outputs = set(graph_outputs)
     from tf2onnx.tf_loader import tf_session, tf_placeholder  # pylint: disable=import-outside-toplevel
 
     ops = g.get_operations()
@@ -208,15 +209,16 @@ def compute_const_folding_using_tf(g, const_node_values):
                 shape = shape_node_outputs[input_names[0]]
                 i = get_index_from_strided_slice_of_shape(node, outputs_to_values)
                 if i is not None and 0 <= i < len(shape) and shape[i] is not None:
-                    outputs_to_values[output_names[0]] = np.array(shape[i])
+                    np_dtype = map_onnx_to_numpy_type(map_tf_dtype(node.outputs[0].dtype))
+                    outputs_to_values[output_names[0]] = np.array(shape[i], dtype=np_dtype)
                     outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
                     progress = True
             can_fold = node.type not in ['Enter']
             can_fold = can_fold and len(input_names) > 0 and all(inp in outputs_to_values for inp in input_names)
             # We can only fold nodes with a single output
             can_fold = can_fold and len(output_names) == 1 and output_names[0] not in outputs_to_values
             # Skip if value already computed, used, and discarded
-            can_fold = can_fold and output_names[0] not in unneeded_outputs
+            can_fold = can_fold and output_names[0] not in unneeded_outputs and output_names[0] not in graph_outputs
             if can_fold:
                 # Make a mini graph containing just the node to fold
                 g2 = tf.Graph()

tf2onnx/tfonnx.py

@@ -409,7 +409,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if target is None:
         target = constants.DEFAULT_TARGET
 
-    outputs_to_values, outputs_to_dtypes = compute_const_folding_using_tf(tf_graph, const_node_values)
+    outputs_to_values, outputs_to_dtypes = compute_const_folding_using_tf(tf_graph, const_node_values, output_names)
 
     onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, _ = \
         tensorflow_to_onnx(tf_graph, shape_override, const_node_values)