From ff50fe82ef262ac6d3fb7d81eb58a56dd548b67e Mon Sep 17 00:00:00 2001 From: hassanaymanelprens58-maker Date: Thu, 30 Apr 2026 14:44:29 +0300 Subject: [PATCH] fix: add missing Float8 mappings to StableHLO to prevent abort This patch adds missing Float8 tensor types (E4M3FN, E5M2, etc.) to the StableHLO conversion maps. This prevents an XLA_ERROR() and subsequent process-wide crash (Issue #506791409). --- torch_xla/csrc/runtime/stablehlo_helper.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/torch_xla/csrc/runtime/stablehlo_helper.cpp b/torch_xla/csrc/runtime/stablehlo_helper.cpp index d2bb31aa5d7..44d79bb2f7e 100644 --- a/torch_xla/csrc/runtime/stablehlo_helper.cpp +++ b/torch_xla/csrc/runtime/stablehlo_helper.cpp @@ -186,6 +186,10 @@ const std::string GetTorchDtypeToStablehloDtype(const std::string& dtype) { if (dtype == "torch.int16") return "i16"; if (dtype == "torch.int32") return "i32"; if (dtype == "torch.int64") return "i64"; + if (dtype == "torch.float8_e4m3fn") return "f8E4M3FN"; + if (dtype == "torch.float8_e4m3fnuz") return "f8E4M3FNUZ"; + if (dtype == "torch.float8_e5m2") return "f8E5M2"; + if (dtype == "torch.float8_e5m2fnuz") return "f8E5M2FNUZ"; XLA_ERROR() << "Unsupported dtype for conversion to Stablehlo type: " << dtype; } @@ -200,6 +204,10 @@ GetHloDtypeToStablehloDtypeMap() { {xla::PrimitiveType::U32, "ui32"}, {xla::PrimitiveType::U64, "ui64"}, {xla::PrimitiveType::F16, "f16"}, {xla::PrimitiveType::BF16, "bf16"}, {xla::PrimitiveType::F32, "f32"}, {xla::PrimitiveType::F64, "f64"}, + {xla::PrimitiveType::F8E4M3FN, "f8E4M3FN"}, + {xla::PrimitiveType::F8E4M3FNUZ, "f8E4M3FNUZ"}, + {xla::PrimitiveType::F8E5M2, "f8E5M2"}, + {xla::PrimitiveType::F8E5M2FNUZ, "f8E5M2FNUZ"}, }; return m_; }