Support multi-input / multi-output models

[seanmor5]

No description provided.

[seanmor5]

We currently already support multi input models (just use Axon.input multiple times). To address multi output models, I propose a tuple combinator:

base = Axon.input({nil, 784}) |> Axon.dense(128)

out1 = base |> Axon.softmax()
out2 = base |> Axon.relu()

model = Axon.tuple([out1, out2])

[seanmor5]

Full support for multi-input models was added in #46.

In order to support multi-output models we need to add the tuple composite layer:

tuple(inputs, opts) :: %Axon{op: :tuple, parent: [input1, input2, ...]}

This layer will only be supported in other composite layers, as the last layer of the network, or in nx layers. In order to include in nx layers we need to also consider how to handle tuple output shapes in the nx layer. For example consider this trivial example:

{x1, x2} = {Axon.input({nil, 32}), Axon.input({nil, 32})}

Axon.tuple([x1, x2])
|> Axon.nx(fn {x1, x2} -> {Nx.cos(x1), Nx.sin(x2)} end)

We can by default wrap tuple output shapes in Axon.tuple, but the limitation is that the outputs then cannot be handled separately unless we introduce decomposition functions for working with the Axon.tuple layer.

[seanmor5]

I've added support for multi output models: 71fce20

Rather than introducing a new layer, you can pass tuples directly to Axon.predict, Axon.init, and Axon.compile. So:

def model do
  inp = Axon.input({nil, 784})
  x = Axon.dense(inp, 128)
  y = Axon.dense(inp, 10)
  {x, y} 
end

params = Axon.init(model())

Axon.predict(model(), params, Nx.random_uniform({32, 784}))

Currently missing support for multi-output / tuple models in the training API