search

elixir-nx / axon

Nx-powered Neural Networks
Apache License 2.0
1.55k stars 104 forks source link

Integrate named tensors #57

Open seanmor5 opened 3 years ago

seanmor5 commented 3 years ago

On the high-level API, we can integrate named tensors by specifying the expected names on input:

Axon.input(batch: nil, channels: 3, height: 224, width: 224)

Then we'll need to consider how these are transformed through the network. For other layers, we could consider adding an option to specify output features as keywords:

Axon.input(batch: nil, pixels: 784)
|> Axon.dense(features: 128, activation: :relu)
|> Axon.dense(label: 10, activation: :softmax)
josevalim commented 3 years ago

I am not sure I like the mixing of keyword lists vs tuples - in a way I would prefer a separate :names option, as in Nx, but maybe we should be allowing keyword lists in Nx too....