Support all tf types

smdebug/core/tfevent/util.py

@@ -27,13 +27,22 @@
     np.dtype(np.complex64): "DT_COMPLEX64",
     np.dtype(np.complex128): "DT_COMPLEX128",
     np.dtype(np.bool): "DT_BOOL",
+    np.dtype([("qint8", "i1")]): "DT_QINT8",
+    np.dtype([("quint8", "u1")]): "DT_QUINT8",
+    np.dtype([("qint16", "<i2")]): "DT_QINT16",
+    np.dtype([("quint16", "<u2")]): "DT_UINT16",
+    np.dtype([("qint32", "<i4")]): "DT_INT32",
 }
 
 
 def _get_proto_dtype(npdtype):
-    if npdtype.kind == "U":
-        return (False, "DT_STRING")
-    return (True, _NP_DATATYPE_TO_PROTO_DATATYPE[npdtype])
+    if hasattr(npdtype, "kind"):
+        if npdtype.kind == "U" or npdtype.kind == "O" or npdtype.kind == "S":
+            return False, "DT_STRING"
+    try:
+        return True, _NP_DATATYPE_TO_PROTO_DATATYPE[npdtype]
+    except KeyError:
+        raise TypeError(f"Numpy Datatype: {np.dtype(npdtype)} is currently not supported")
 
 
 def make_tensor_proto(nparray_data, tag):

tests/tensorflow2/test_tensorflow2_datatypes.py

@@ -0,0 +1,31 @@
+# Third Party
+import numpy as np
+from packaging import version
+from tensorflow.python.framework.dtypes import _NP_TO_TF
+
+# First Party
+from smdebug.core.tfevent.util import _get_proto_dtype
+
+
+def test_tensorflow2_datatypes():
+    # _NP_TO_TF contains all the mappings
+    # of numpy to tf types
+    try:
+        from tensorflow import __version__ as tf_version
+
+        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)
+    except (ModuleNotFoundError, ValueError, ImportError):
+        pass
+
+    for _type in _NP_TO_TF:
+        try:
+            _get_proto_dtype(np.dtype(_type))
+        except Exception:
+            assert False
+    assert True