Merge metadata in compiler, resolves #567

elixir-nx · May 14, 2024 · 0473cf9 · 0473cf9
1 parent 5f9e7bc
commit 0473cf9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/lib/axon/compiler.ex b/lib/axon/compiler.ex
@@ -1128,6 +1128,9 @@ defmodule Axon.Compiler do
           out = Nx.Defn.Expr.metadata(Nx.Defn.Expr.tensor(out), %{axon_layer: op_name})
           %{stateful | output: out}
 
+        %Nx.Tensor{data: %{op: :metadata, args: [arg, metadata]} = expr} = out ->
+          %{out | data: %{expr | args: [arg, Map.put(metadata, :axon_layer, op_name)]}}
+
         %Nx.Tensor{} = out ->
           Nx.Defn.Expr.metadata(Nx.Defn.Expr.tensor(out), %{axon_layer: op_name})
 

diff --git a/test/axon/integration_test.exs b/test/axon/integration_test.exs
@@ -32,7 +32,7 @@ defmodule Axon.IntegrationTest do
       model_state =
         model
         |> Axon.Loop.trainer(:binary_cross_entropy, :sgd)
-        |> Axon.Loop.run(data, Axon.ModelState.empty(), iterations: 100, epochs: 10)
+        |> Axon.Loop.run(data, Axon.ModelState.empty(), iterations: 100, epochs: 20)
 
       eval_results =
         model