elixir-nx / nx

Multi-dimensional arrays (tensors) and numerical definitions for Elixir
2.55k stars 190 forks source link

Add initial FP8 support #1507

Open seanmor5 opened 1 month ago

seanmor5 commented 1 month ago

This PR adds support for an emulated FP8 in Nx's binary backend and adds support for FP8 operations from XLA. The FP8 type in EXLA is just there: FP8E5M2 type. They have 3 of these FP8 types that we can configure if we want.

In the binary backend the logic is in shared.ex. We have a function for interpreting erlang float's to a binary representation, and then we convert the binary chunks to f16 chunks and interpret those to f64.