From b2b7c738f421b9c0c52023637e6a8a111d6a7b2a Mon Sep 17 00:00:00 2001
From: Kyle Herndon <kyle.herndon@amd.com>
Date: Thu, 13 Feb 2025 01:02:11 +0000
Subject: [PATCH] Fix bfloat16_t rounding functions

---
 shortfin/python/array_host_ops.cc | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/shortfin/python/array_host_ops.cc b/shortfin/python/array_host_ops.cc
index c73ad786e..55e0546a3 100644
--- a/shortfin/python/array_host_ops.cc
+++ b/shortfin/python/array_host_ops.cc
@@ -103,6 +103,23 @@ template <>
 struct is_trivially_copyable<bfloat16_t> : std::true_type {};
 }  // namespace std
 
+// Math functions needed by xtensor for bfloat16_t
+inline constexpr bfloat16_t round(bfloat16_t x) noexcept {
+  return bfloat16_t(std::round(float(x)));
+}
+
+inline constexpr bfloat16_t ceil(bfloat16_t x) noexcept {
+  return bfloat16_t(std::ceil(float(x)));
+}
+
+inline constexpr bfloat16_t floor(bfloat16_t x) noexcept {
+  return bfloat16_t(std::floor(float(x)));
+}
+
+inline constexpr bfloat16_t trunc(bfloat16_t x) noexcept {
+  return bfloat16_t(std::trunc(float(x)));
+}
+
 #endif  // BFLOAT16_HPP
 
 using namespace shortfin::array;