[IR] Introduce short name for dtypes (#2249)

Introduce short name for dtypes as a more compact way of describing the
data types in strings. Users can already access the enums by name with
e.g. `ir.DataType["DOUBLE"]`.

Author: justinchuby

onnxscript/ir/_enums.py

@@ -100,6 +100,17 @@ def from_numpy(cls, dtype: np.dtype) -> DataType:
                 return DataType.FLOAT4E2M1
         raise TypeError(f"Unsupported numpy data type: {dtype}")
 
+    @classmethod
+    def from_short_name(cls, short_name: str) -> DataType:
+        """Returns the ONNX data type for the short name.
+
+        Raises:
+            TypeError: If the short name is not available for the data type.
+        """
+        if short_name not in _SHORT_NAME_TO_DATA_TYPE:
+            raise TypeError(f"Unknown short name: {short_name}")
+        return cls(_SHORT_NAME_TO_DATA_TYPE[short_name])
+
     @property
     def itemsize(self) -> float:
         """Returns the size of the data type in bytes."""
@@ -115,6 +126,22 @@ def numpy(self) -> np.dtype:
             raise TypeError(f"Numpy does not support ONNX data type: {self}")
         return _DATA_TYPE_TO_NP_TYPE[self]
 
+    def short_name(self) -> str:
+        """Returns the short name of the data type.
+
+        The short name is a string that is used to represent the data type in a more
+        compact form. For example, the short name for `DataType.FLOAT` is "f32".
+        To get the corresponding data type back, call ``from_short_name`` on a string.
+
+        Naming reference: https://github.com/pytorch/pytorch/blob/4bead7b85ea4160243c74109e0ce9bb80686d016/torch/utils/_dtype_abbrs.py
+
+        Raises:
+            TypeError: If the short name is not available for the data type.
+        """
+        if self not in _DATA_TYPE_TO_SHORT_NAME:
+            raise TypeError(f"Short name not available for ONNX data type: {self}")
+        return _DATA_TYPE_TO_SHORT_NAME[self]
+
     def __repr__(self) -> str:
         return self.name
 
@@ -184,3 +211,32 @@ def __str__(self) -> str:
 
 # ONNX DataType to Numpy dtype.
 _DATA_TYPE_TO_NP_TYPE = {v: k for k, v in _NP_TYPE_TO_DATA_TYPE.items()}
+
+_DATA_TYPE_TO_SHORT_NAME = {
+    DataType.UNDEFINED: "undefined",
+    DataType.BFLOAT16: "bf16",
+    DataType.DOUBLE: "f64",
+    DataType.FLOAT: "f32",
+    DataType.FLOAT16: "f16",
+    DataType.FLOAT8E4M3FN: "f8e4m3fn",
+    DataType.FLOAT8E5M2: "f8e5m2",
+    DataType.FLOAT8E4M3FNUZ: "f8e4m3fnuz",
+    DataType.FLOAT8E5M2FNUZ: "f8e5m2fnuz",
+    DataType.FLOAT4E2M1: "f4e2m1",
+    DataType.COMPLEX64: "c64",
+    DataType.COMPLEX128: "c128",
+    DataType.INT4: "i4",
+    DataType.INT8: "i8",
+    DataType.INT16: "i16",
+    DataType.INT32: "i32",
+    DataType.INT64: "i64",
+    DataType.BOOL: "b8",
+    DataType.UINT4: "u4",
+    DataType.UINT8: "u8",
+    DataType.UINT16: "u16",
+    DataType.UINT32: "u32",
+    DataType.UINT64: "u64",
+    DataType.STRING: "s",
+}
+
+_SHORT_NAME_TO_DATA_TYPE = {v: k for k, v in _DATA_TYPE_TO_SHORT_NAME.items()}

onnxscript/ir/_enums_test.py

@@ -122,6 +122,37 @@ def test_repr_and_str_return_name(self):
         self.assertEqual(str(_enums.DataType.DOUBLE), "DOUBLE")
         self.assertEqual(repr(_enums.DataType.DOUBLE), "DOUBLE")
 
+    def test_short_name_conversion(self):
+        for dtype in _enums.DataType:
+            short_name = dtype.short_name()
+            self.assertEqual(_enums.DataType.from_short_name(short_name), dtype)
+
+    def test_access_by_name(self):
+        self.assertEqual(_enums.DataType["FLOAT"], _enums.DataType.FLOAT)
+        self.assertEqual(_enums.DataType["UINT8"], _enums.DataType.UINT8)
+        self.assertEqual(_enums.DataType["INT8"], _enums.DataType.INT8)
+        self.assertEqual(_enums.DataType["UINT16"], _enums.DataType.UINT16)
+        self.assertEqual(_enums.DataType["INT16"], _enums.DataType.INT16)
+        self.assertEqual(_enums.DataType["INT32"], _enums.DataType.INT32)
+        self.assertEqual(_enums.DataType["INT64"], _enums.DataType.INT64)
+        self.assertEqual(_enums.DataType["STRING"], _enums.DataType.STRING)
+        self.assertEqual(_enums.DataType["BOOL"], _enums.DataType.BOOL)
+        self.assertEqual(_enums.DataType["FLOAT16"], _enums.DataType.FLOAT16)
+        self.assertEqual(_enums.DataType["DOUBLE"], _enums.DataType.DOUBLE)
+        self.assertEqual(_enums.DataType["UINT32"], _enums.DataType.UINT32)
+        self.assertEqual(_enums.DataType["UINT64"], _enums.DataType.UINT64)
+        self.assertEqual(_enums.DataType["COMPLEX64"], _enums.DataType.COMPLEX64)
+        self.assertEqual(_enums.DataType["COMPLEX128"], _enums.DataType.COMPLEX128)
+        self.assertEqual(_enums.DataType["BFLOAT16"], _enums.DataType.BFLOAT16)
+        self.assertEqual(_enums.DataType["FLOAT8E4M3FN"], _enums.DataType.FLOAT8E4M3FN)
+        self.assertEqual(_enums.DataType["FLOAT8E4M3FNUZ"], _enums.DataType.FLOAT8E4M3FNUZ)
+        self.assertEqual(_enums.DataType["FLOAT8E5M2"], _enums.DataType.FLOAT8E5M2)
+        self.assertEqual(_enums.DataType["FLOAT8E5M2FNUZ"], _enums.DataType.FLOAT8E5M2FNUZ)
+        self.assertEqual(_enums.DataType["UINT4"], _enums.DataType.UINT4)
+        self.assertEqual(_enums.DataType["INT4"], _enums.DataType.INT4)
+        self.assertEqual(_enums.DataType["FLOAT4E2M1"], _enums.DataType.FLOAT4E2M1)
+        self.assertEqual(_enums.DataType["UNDEFINED"], _enums.DataType.UNDEFINED)
+
 
 class AttributeTypeTest(unittest.TestCase):
     def test_enums_are_the_same_as_spec(self):