From f93b578316c30c21e759194ca9d9186d75e49267 Mon Sep 17 00:00:00 2001
From: Tarun Karuturi <tkaruturi@meta.com>
Date: Tue, 29 Oct 2024 15:29:29 -0700
Subject: [PATCH] Add uint16 support in exir serde

Summary: Adding uint16 support in serialization after D64812253.

Differential Revision: D65175529
---
 exir/serde/export_serialize.py | 1 +
 exir/serde/schema.py           | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/exir/serde/export_serialize.py b/exir/serde/export_serialize.py
index 8ec864e489b..570af19baf8 100644
--- a/exir/serde/export_serialize.py
+++ b/exir/serde/export_serialize.py
@@ -148,6 +148,7 @@ def _reverse_map(d: Dict[Any, Enum]):
     torch.complex128: ScalarType.COMPLEXDOUBLE,
     torch.bool: ScalarType.BOOL,
     torch.bfloat16: ScalarType.BFLOAT16,
+    torch.uint16: ScalarType.UINT16
 }
 
 
diff --git a/exir/serde/schema.py b/exir/serde/schema.py
index 9bbdd5b1c02..16ac724917c 100644
--- a/exir/serde/schema.py
+++ b/exir/serde/schema.py
@@ -35,7 +35,7 @@ class ScalarType(IntEnum):
     COMPLEXDOUBLE = 11
     BOOL = 12
     BFLOAT16 = 13
-
+    UINT16 = 14
 
 class Layout(IntEnum):
     Unknown = 0