Merge metadata in compiler, resolves #567

Author: seanmor5

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})
 

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