Support complex types

[elixir-nx]

Support complex types

[seanmor5]

One proposal that would resolve #117 (before the merge), #118, #119 and possibly allow for more flexible handling of numeric types would be to introduce an Nx.Type behaviour and implement our current types in terms of the behaviour. Right now we use the match!, read!, and write! macros in pretty much every Nx.BinaryTensor method. A behaviour would look something like:

defmodule Nx.Type.Float16 do

   @impl true
   def pack(scalar), do: ...pack to binary...
   
   @impl true
   def unpack(binary), do: ...unpack to scalar...

So pack and unpack are all about encoding/decoding terms to binaries and vice versa. Then, in order to extend our current implementations to any Nx.Type module, we'd need to override the Kernel methods to work on special scalar terms. Signed/Unsigned integers need no overrides. Floats only need to override in the case of handling NaN, -Inf, Inf, -0, +0, etc. Complex numbers need to override everything. For example, a Complex128 would unpack into a tuple of {real term, imag term} and then :erlang.+, :erlang.*, etc. would be overridden to operate on those tuples.

I think a big advantage of this approach is in flexibility and rapid prototyping of research ideas. While nobody would test/research their own floating point representations in Elixir, they could test them in a native backend (like EXLA) and then implement this behavior to get full introspection from the Elixir side of things. In addition to this, backend developers aren't hamstrung to a fixed set of types. While I'm sure our current implementation covers a pretty complete set of types, it doesn't mean something won't pop up in the future (e.g. bfloat16).

A really big disadvantage is that this approach comes with a lot of overhead initially, but I think in many ways it will keep our current implementations type-agnostic.

[josevalim]

I think this is a good direction. The concern is indeed the overhead. One of the ways we can tackle the overhead is by compiling those clauses at compilation time: instead of invoking the behaviour at runtime, we call it at compilation time to unpack the relevant clauses, similar to how we do it today.

Unfortunately complex numbers are a bit more complex (ha) because we need to introduce a runtime representation for them too. A struct or a two-element tuple. And we likely want to make things like Nx.tensor([some_complex_number]) to work with them out of the box too - so we will need to look at every is_number clause and rewrite them to be is_number or is_complex.

But in a nutshell, if we do add a new type, we should probably consider moving the approach that you described. 👍