Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for global layer options #563

Merged
merged 2 commits into from
Mar 6, 2024
Merged

Conversation

jonatanklosko
Copy link
Member

This allows to configure model layers at Axon-compilation time. This is useful if we want to make some compilation-time decisions when calling the model, without building a different model altogether.

API-wise the options are passed as

Axon.build(model, global_layer_options: [output_hidden_states: true])

Those options are merged into layer options, but only when declared explicitly, like so:

Axon.layer(
  fn hidden_state,  opts ->
    if opts[:output_hidden_states] do
      ...
    else
      ...
    end
  end,
  [input],
  global_options: [:output_hidden_states]
)

We could pass them nested too and access as opts[:global_options][:output_hidden_states], but this would break all existing layers that have options allowlist using keyword!. I actually think that merging them in a flat manner is better, because they can be validated with keyword! in the exact same way.

Naming suggestions are welcome, for example I was considering :global_layer_options -> :layer_options and :global_options -> :expose_options, though not sure if "expose" is clear enough.

@jonatanklosko jonatanklosko merged commit 7e0e593 into main Mar 6, 2024
3 of 5 checks passed
@jonatanklosko jonatanklosko deleted the jk-global-layer-options branch March 6, 2024 07:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants