Fix encoding of bf16

Author: josevalim

exla/lib/exla/mlir/value.ex

@@ -921,13 +921,14 @@ defmodule EXLA.MLIR.Value do
     end
   end
 
-  defp float_hex(value, {_, size} = type) do
+  defp float_hex(value, {mod, size} = type) do
     data =
       case value do
         :nan -> type |> Nx.Type.nan_binary() |> native_to_big()
         :infinity -> type |> Nx.Type.infinity_binary() |> native_to_big()
         :neg_infinity -> type |> Nx.Type.neg_infinity_binary() |> native_to_big()
         value when size == 8 -> f8E5M2_to_big(value)
+        value when mod == :bf and size == 16 -> bf16_to_big(value)
         value -> <<value::float-size(size)-big>>
       end
 
@@ -938,6 +939,10 @@ defmodule EXLA.MLIR.Value do
     binary_part(<<x::float-big-16>>, 0, 1)
   end
 
+  defp bf16_to_big(x) do
+    binary_part(<<x::float-big-32>>, 0, 2)
+  end
+
   defp native_to_big(binary) do
     size = byte_size(binary) * 8
     <<value::size(size)-native>> = binary

exla/test/exla/defn/expr_test.exs

@@ -86,19 +86,23 @@ defmodule EXLA.Defn.ExprTest do
     end
   end
 
-  describe "float8" do
-    defn return_float8, do: Nx.tensor(1, type: {:f, 8})
+  describe "types" do
+    defn return_f8, do: Nx.tensor(1, type: {:f, 8})
 
-    test "supports float8 return types" do
-      assert_equal(return_float8(), Nx.tensor(1, type: {:f, 8}))
+    test "f8" do
+      assert_equal(return_f8(), Nx.tensor(1, type: {:f, 8}))
+    end
+
+    defn return_f16, do: Nx.tensor(1, type: {:f, 16})
+
+    test "f16" do
+      assert_equal(return_f16(), Nx.tensor(1, type: {:f, 16}))
     end
-  end
 
-  describe "float16" do
-    defn return_float, do: Nx.tensor(1, type: {:f, 16})
+    defn return_bf16, do: Nx.tensor(1, type: {:bf, 16})
 
-    test "supports float16 return types" do
-      assert_equal(return_float(), Nx.tensor(1, type: {:f, 16}))
+    test "bf16" do
+      assert_equal(return_bf16(), Nx.tensor(1, type: {:bf, 16}))
     end
   end