View Source Bumblebee.ModelSpec behaviour (Bumblebee v0.1.2)

An interface for configuring and building models based on the same architecture.

Every module implementing this behaviour is expected to also define a configuration struct.

Link to this section Summary

Callbacks

Returns the list of supported model architectures.

Builds a template input for the model.

Builds an Axon model according to the given configuration.

Link to this section Types

Link to this section Callbacks

@callback architectures() :: [atom()]

Returns the list of supported model architectures.

@callback input_template(t()) :: map()

Builds a template input for the model.

The template is used to compile the model when initializing parameters.

@callback model(t()) :: Axon.t()

Builds an Axon model according to the given configuration.