tf-2.6.0rc2 (#1662)

* tf-2.6.0rc2

Signed-off-by: Guenther Schmuelling <[email protected]>

* fix large model support for tf-2.6

Signed-off-by: Guenther Schmuelling <[email protected]>

* pylint

Signed-off-by: Guenther Schmuelling <[email protected]>

Author: guschmue

ci_build/azure_pipelines/templates/setup.yml

@@ -41,6 +41,11 @@ steps:
       then
         pip install tensorflow-text>=2.5
       fi
+      if [[ $CI_TF_VERSION == 2.6* ]] ;
+      then
+        # FIXME: make it >= 2.6 after offical tensorflow-text was released
+        pip install tensorflow-text==2.6.0rc0
+      fi
     fi
 
     python setup.py install

ci_build/azure_pipelines/unit_test.yml

@@ -3,6 +3,17 @@
 stages:
   - stage:
     jobs:
+    - template: 'templates/job_generator.yml'
+      parameters:
+        # tf 2.5
+        python_versions: ['3.8']
+        tf_versions: ['2.6.0rc2']
+        onnx_opsets: ['14']
+        job:
+          steps:
+          - template: 'unit_test.yml'
+        report_coverage: 'True'
+
     - template: 'templates/job_generator.yml'
       parameters:
         # TFJS tf 2.5

tf2onnx/tf_loader.py

@@ -161,7 +161,18 @@ def make_tensor_proto_wrapped(values, dtype=None, shape=None, verify_shape=False
     tensor_util.make_tensor_proto = make_tensor_proto_wrapped
 
     try:
-        converter_data = _FunctionConverterData(func=func, lower_control_flow=False, aggressive_inlining=True)
+        function_converter = _FunctionConverterData
+        if LooseVersion(tf.__version__) >= "2.6.0":
+            from tensorflow.python.eager import context
+            from tensorflow.python.framework.convert_to_constants import _FunctionConverterDataInEager, \
+                _FunctionConverterDataInGraph
+            if context.executing_eagerly():
+                function_converter = _FunctionConverterDataInEager
+            else:
+                function_converter = _FunctionConverterDataInGraph
+        else:
+            function_converter = _FunctionConverterData
+        converter_data = function_converter(func=func, lower_control_flow=False, aggressive_inlining=True)
         frozen_graph_def, _ = _replace_variables_by_constants(converter_data=converter_data)
     finally:
         tensor_util.make_tensor_proto = make_tensor_proto_original