Skip to content

Add Axon.capture - #643

Merged
polvalente merged 4 commits into
elixir-nx:mainfrom
Chapaman:sm-capture-continue
Oct 8, 2026
Merged

polvalente merged 4 commits into
elixir-nx:mainfrom
Chapaman:sm-capture-continue

Conversation

@Chapaman

Copy link
Copy Markdown
Contributor

Continues #626

@Chapaman
Chapaman marked this pull request as draft August 14, 2026 23:44
@Chapaman
Chapaman marked this pull request as ready for review August 17, 2026 21:06
@polvalente
polvalente requested review from polvalente and a balanced review from Copilot October 7, 2026 22:14

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🟡 Changes recommended

Graph truncation, output rewiring, parameter naming, input validation, and test coverage have unresolved issues.

6 open findings
What changed in this PR

Adds Axon.capture/2 for reusing and composing existing model graphs, optionally truncated at a named layer.

Changes:

  • Adds the documented capture API.
  • Resolves layer and input names.
  • Rewires captured graphs to replacement inputs.
File Description
lib/​axon.ex Implements model capture, truncation, validation, and graph rewiring.

🧠 Review effort: Balanced


💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread lib/axon.ex
Comment on lines +4569 to +4572
merged_nodes =
Enum.reduce(new_inputs_map, nodes, fn {_k, %Axon{nodes: new_nodes}}, acc_nodes ->
Map.merge(new_nodes, acc_nodes)
end)
Comment thread lib/axon.ex
{node_id, %{node | parent: updated_parents}}
end)

%Axon{output: output_id, nodes: updated_nodes}
Comment thread lib/axon.ex

"""
@doc type: :graph
def capture(%Axon{} = axon, opts \\ []) when is_list(opts) do
Comment thread lib/axon.ex
"Available layers: #{inspect(available)}"
end

%Axon{axon | output: target_id}
Comment thread lib/axon.ex
Comment on lines +4540 to +4541
%{} = inputs_map ->
inputs_map
Comment thread lib/axon.ex
Comment on lines +4486 to +4488
fn new_inputs ->
rewire_inputs(truncated, input_name_to_id, new_inputs)
end
Each pairing of one or many inputs with a tensor or a container output should keep the same values after the inputs are replaced.
@polvalente
polvalente merged commit b182288 into elixir-nx:main Oct 8, 2026
7 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.

4 participants