Skip to content

Correct implementation of weight tying#625

Merged
seanmor5 merged 3 commits into
mainfrom
real-weight-tying
May 11, 2026
Merged

Correct implementation of weight tying#625
seanmor5 merged 3 commits into
mainfrom
real-weight-tying

Conversation

@seanmor5

Copy link
Copy Markdown
Contributor

Resolves #325

Simple weight tying implementation via shared parameters in the model state. These are resolved to the correct parameter, and allow an optional transform (e.g. Nx.transpose/1) so the parameter shape matches effectively.

Comment thread lib/axon/compiler.ex
Comment on lines +1183 to +1186
layer_params =
with %Axon.ModelState.SharedParameter{path: path} <- params[layer_name] do
get_in(params, path)
end

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
layer_params =
with %Axon.ModelState.SharedParameter{path: path} <- params[layer_name] do
get_in(params, path)
end
layer_params =
case params[layer_name] do
%Axon.ModelState.SharedParameter{path: path} ->
get_in(params, path)
nil ->
nil
end

I think this is slightly more readable

Comment thread lib/axon/compiler.ex

@polvalente polvalente left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor style comments

seanmor5 and others added 2 commits May 11, 2026 19:06
Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>
@seanmor5 seanmor5 merged commit 0562126 into main May 11, 2026
5 checks passed
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.

2 participants