From 84e3e068bc24878a549c2c0a5031b720c56fa4ed Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 28 May 2025 09:02:34 -0700
Subject: [PATCH] Implement `is_floating_point` on dtypes

Migration of changes in https://github.com/microsoft/onnxscript/pull/2335

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 src/onnx_ir/_enums.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/onnx_ir/_enums.py b/src/onnx_ir/_enums.py
index 26ecaa47..19fa5920 100644
--- a/src/onnx_ir/_enums.py
+++ b/src/onnx_ir/_enums.py
@@ -142,6 +142,20 @@ def short_name(self) -> str:
             raise TypeError(f"Short name not available for ONNX data type: {self}")
         return _DATA_TYPE_TO_SHORT_NAME[self]
 
+    def is_floating_point(self) -> bool:
+        """Returns True if the data type is a floating point type."""
+        return self in {
+            DataType.FLOAT,
+            DataType.FLOAT16,
+            DataType.DOUBLE,
+            DataType.BFLOAT16,
+            DataType.FLOAT8E4M3FN,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+            DataType.FLOAT4E2M1,
+        }
+
     def __repr__(self) -> str:
         return self.name