From 03f07e2a1e375252b34c0e232da670f13e68836c Mon Sep 17 00:00:00 2001
From: Nachiket Mokashi <nmokashi@amazon.com>
Date: Thu, 7 Nov 2024 11:51:29 -0800
Subject: [PATCH] XLA downcast for S64 and U64 for Neuron (#8347)

---
 torch_xla/csrc/dtype.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/torch_xla/csrc/dtype.cpp b/torch_xla/csrc/dtype.cpp
index 923f1152c9d..8f7be2eacb2 100644
--- a/torch_xla/csrc/dtype.cpp
+++ b/torch_xla/csrc/dtype.cpp
@@ -143,9 +143,11 @@ xla::PrimitiveType MaybeDowncastToXlaDeviceType(
       return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::S32
                                         : xla::PrimitiveType::S16;
     case xla::PrimitiveType::S64:
-      return xla::PrimitiveType::S64;
+      return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::S32
+                                        : xla::PrimitiveType::S64;
     case xla::PrimitiveType::U64:
-      return xla::PrimitiveType::U64;
+      return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::U32
+                                        : xla::PrimitiveType::U64;
     case xla::PrimitiveType::C128:
       return xla::PrimitiveType::C128;
     default: