Modify Asserts to Work with TF 2.1.0 and TF 2.0.0

tests/tensorflow2/test_keras.py

@@ -558,7 +558,7 @@ def test_include_regex(out_dir, tf_eager_mode):
 
     tr = create_trial_fast_refresh(out_dir)
     tnames = tr.tensor_names(collection="custom_coll")
-    assert len(tnames) == 12
+    assert len(tnames) == (12 if is_tf_2_2() else 4)
     for tname in tnames:
         assert tr.tensor(tname).value(0) is not None
 
@@ -729,10 +729,7 @@ def test_keras_fit_pure_eager(out_dir, tf_eager_mode):
     helper_keras_fit(trial_dir=out_dir, hook=hook, eager=tf_eager_mode, run_eagerly=True)
 
     trial = smd.create_trial(path=out_dir)
-    if is_tf_2_2():
-        assert len(trial.tensor_names()) == 27
-    else:
-        assert len(trial.tensor_names()) == (20 if is_tf_2_3() else 21)
+    assert len(trial.tensor_names()) == (27 if is_tf_2_2() else 13)
     assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
     assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
     assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5

tests/tensorflow2/test_model_subclassing.py

@@ -2,6 +2,7 @@
 import tensorflow as tf
 from tensorflow.keras.layers import BatchNormalization, Conv2D, Dense, Flatten
 from tensorflow.keras.models import Model
+from tests.tensorflow2.utils import is_tf_2_2
 
 # First Party
 import smdebug.tensorflow as smd
@@ -78,7 +79,12 @@ def test_subclassed_model(out_dir):
     trial = smd.create_trial(out_dir)
     assert len(trial.tensor_names(collection=smd.CollectionKeys.LAYERS)) == 8
 
-    assert trial.tensor_names(collection=smd.CollectionKeys.INPUTS) == ["model_input"]
-    assert trial.tensor_names(collection=smd.CollectionKeys.OUTPUTS) == ["labels", "predictions"]
     assert trial.tensor_names(collection=smd.CollectionKeys.LOSSES) == ["loss"]
-    assert len(trial.tensor_names(collection=smd.CollectionKeys.GRADIENTS)) == 6
+    if is_tf_2_2():
+        # Feature to save model inputs and outputs was first added for TF 2.2.0
+        assert trial.tensor_names(collection=smd.CollectionKeys.INPUTS) == ["model_input"]
+        assert trial.tensor_names(collection=smd.CollectionKeys.OUTPUTS) == [
+            "labels",
+            "predictions",
+        ]
+        assert len(trial.tensor_names(collection=smd.CollectionKeys.GRADIENTS)) == 6

tests/tensorflow2/test_support_dicts.py

@@ -1,6 +1,8 @@
 # Third Party
 import numpy as np
+import pytest
 import tensorflow as tf
+from tests.tensorflow2.utils import is_tf_2_2
 
 # First Party
 import smdebug.tensorflow as smd
@@ -29,6 +31,10 @@ def create_model():
     return model
 
 
+@pytest.mark.skipif(
+    is_tf_2_2() is False,
+    reason="Feature to save model inputs and outputs was first added for TF 2.2.0",
+)
 def test_support_dicts(out_dir):
     model = create_model()
     optimizer = tf.keras.optimizers.Adadelta(lr=1.0, rho=0.95, epsilon=None, decay=0.0)

tests/tensorflow2/test_tensorflow2_datatypes.py

@@ -1,31 +1,32 @@
 # Third Party
 import numpy as np
-from packaging import version
+import pytest
 from tensorflow.python.framework.dtypes import _NP_TO_TF
+from tests.tensorflow2.utils import is_tf_2_2
 
 # First Party
 from smdebug.core.tfevent.util import _get_proto_dtype
 
 
+@pytest.mark.skipif(
+    is_tf_2_2() is False, reason="Brain Float Is Unavailable in lower versions of TF"
+)
 def test_tensorflow2_datatypes():
     # _NP_TO_TF contains all the mappings
     # of numpy to tf types
     try:
-        from tensorflow import __version__ as tf_version
+        from tensorflow.python import _pywrap_bfloat16
 
-        if version.parse(tf_version) >= version.parse("2.0.0"):
-            from tensorflow.python import _pywrap_bfloat16
-
-            # TF 2.x.x Implements a Custom Numpy Datatype for Brain Floating Type
-            # Which is currently only supported on TPUs
-            _np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
-            _NP_TO_TF.pop(_np_bfloat16)
+        # TF 2.x.x Implements a Custom Numpy Datatype for Brain Floating Type
+        # Which is currently only supported on TPUs
+        _np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
+        _NP_TO_TF.pop(_np_bfloat16)
     except (ModuleNotFoundError, ValueError, ImportError):
         pass
 
     for _type in _NP_TO_TF:
         try:
             _get_proto_dtype(np.dtype(_type))
         except Exception:
-            assert False
+            assert False, f"{_type} not supported"
     assert True