Add metadata to layers

[seanmor5]

This attaches Axon specific metadata to layers in the form of the op_name. This allows backends and compilers to match specifically on these patterns if they're available and provide optimized versions of layers rather than depending on the lower-level Nx implementation.

[seanmor5]

I think there may be an Nx bug that surfaces here. All of my "constant with other operations" tests fail with a function clause error on scope_ids_each:

         # 1
         #Nx.Tensor<
           f32[1][2][3]
           [
             [
               [0.0, 1.0, 2.0],
               [3.0, 4.0, 5.0]
             ]
           ]
         >
     
         # 2
         {%{#Reference<0.4002093073.794296321.209144> => :metadata}, %{}}
     
         # 3
         nil
     
     Attempted function clauses (showing 3 out of 3):
     
         defp scope_ids_each(%Nx.Tensor{data: %Nx.Defn.Expr{op: :constant}}, {ids, cond_ids}, _scope)
         defp scope_ids_each(%Nx.Tensor{data: %Nx.Defn.Expr{id: id, op: op} = expr} = t, {ids, cond_ids}, nil)
         defp scope_ids_each(%Nx.Tensor{data: %Nx.Defn.Expr{id: id}} = t, {ids, cond_ids}, scope)

cc @josevalim

[josevalim]

@seanmor5 it should be good now :)

[josevalim]

FWIW, the failing tests should be fixed by using Nx.all_close or similar. :)