Add support for global layer options (#563)

Author: jonatanklosko

lib/axon.ex

@@ -301,9 +301,14 @@ defmodule Axon do
   to inference function except:
 
     * `:name` - layer name.
+
     * `:op_name` - layer operation for inspection and building parameter map.
+
     * `:mode` - if the layer should run only on `:inference` or `:train`. Defaults to `:both`
 
+    * `:global_options` - a list of global option names that this layer
+      supports. Global options passed to `build/2` will be forwarded to
+      the layer, as long as they are declared
 
   Note this means your layer should not use these as input options,
   as they will always be dropped during inference compilation.
@@ -332,14 +337,15 @@ defmodule Axon do
     {mode, opts} = Keyword.pop(opts, :mode, :both)
     {name, opts} = Keyword.pop(opts, :name)
     {op_name, opts} = Keyword.pop(opts, :op_name, :custom)
+    {global_options, opts} = Keyword.pop(opts, :global_options, [])
     name = name(op_name, name)
 
     id = System.unique_integer([:positive, :monotonic])
-    axon_node = make_node(id, op, name, op_name, mode, inputs, params, args, opts)
+    axon_node = make_node(id, op, name, op_name, mode, inputs, params, args, opts, global_options)
     %Axon{output: id, nodes: Map.put(updated_nodes, id, axon_node)}
   end
 
-  defp make_node(id, op, name, op_name, mode, inputs, params, args, layer_opts) do
+  defp make_node(id, op, name, op_name, mode, inputs, params, args, layer_opts, global_options) do
     {:current_stacktrace, [_process_info, _axon_layer | stacktrace]} =
       Process.info(self(), :current_stacktrace)
 
@@ -354,6 +360,7 @@ defmodule Axon do
       policy: Axon.MixedPrecision.create_policy(),
       hooks: [],
       opts: layer_opts,
+      global_options: global_options,
       op_name: op_name,
       stacktrace: stacktrace
     }
@@ -3651,6 +3658,9 @@ defmodule Axon do
       to control differences in compilation at training or inference time.
       Defaults to `:inference`
 
+    * `:global_layer_options` - a keyword list of options passed to
+      layers that accept said options
+
   All other options are forwarded to the underlying JIT compiler.
   """
   @doc type: :model

lib/axon/compiler.ex

@@ -50,7 +50,8 @@ defmodule Axon.Compiler do
     debug? = Keyword.get(opts, :debug, false)
     mode = Keyword.get(opts, :mode, :inference)
     seed = Keyword.get_lazy(opts, :seed, fn -> :erlang.system_time() end)
-    config = %{mode: mode, debug?: debug?}
+    global_layer_options = Keyword.get(opts, :global_layer_options, [])
+    config = %{mode: mode, debug?: debug?, global_layer_options: global_layer_options}
 
     {time, {root_id, {cache, _op_counts, _block_cache}}} =
       :timer.tc(fn ->
@@ -718,14 +719,15 @@ defmodule Axon.Compiler do
            parameters: layer_params,
            args: args,
            opts: opts,
+           global_options: global_options,
            policy: policy,
            hooks: hooks,
            op_name: op_name,
            stacktrace: stacktrace
          },
          nodes,
          cache_and_counts,
-         %{mode: mode, debug?: debug?} = config
+         %{mode: mode, debug?: debug?, global_layer_options: global_layer_options} = config
        )
        when (is_function(op) or is_atom(op)) and is_list(inputs) do
     # Traverse to accumulate cache and get parent_ids for
@@ -761,10 +763,12 @@ defmodule Axon.Compiler do
         name,
         args,
         opts,
+        global_options,
         policy,
         layer_params,
         hooks,
         mode,
+        global_layer_options,
         stacktrace
       )
 
@@ -841,10 +845,12 @@ defmodule Axon.Compiler do
          name,
          args,
          opts,
+         global_options,
          policy,
          layer_params,
          hooks,
          mode,
+         global_layer_options,
          layer_stacktrace
        ) do
     # Recurse graph inputs and invoke cache to get parent results,
@@ -914,7 +920,12 @@ defmodule Axon.Compiler do
 
       # Compute arguments to be forwarded and ensure `:mode` is included
       # for inference/training behavior dependent functions
-      args = Enum.reverse(tensor_inputs, [Keyword.put(opts, :mode, mode)])
+      layer_opts =
+        opts
+        |> Keyword.merge(Keyword.take(global_layer_options, global_options))
+        |> Keyword.put(:mode, mode)
+
+      args = Enum.reverse(tensor_inputs, [layer_opts])
 
       # For built-in layers we always just apply the equivalent function
       # in Axon.Layers. The implication of this is that every function which

lib/axon/node.ex

@@ -12,6 +12,7 @@ defmodule Axon.Node do
     :policy,
     :hooks,
     :opts,
+    :global_options,
     :op_name,
     :stacktrace
   ]

test/axon/compiler_test.exs

@@ -5837,4 +5837,27 @@ defmodule CompilerTest do
       assert predict_fn.(params, x) == Nx.add(x, a)
     end
   end
+
+  describe "global layer options" do
+    test "global options are forwarded to the layer when declared" do
+      input = Axon.input("input")
+
+      model =
+        Axon.layer(
+          fn input, opts ->
+            assert Keyword.has_key?(opts, :option1)
+            refute Keyword.has_key?(opts, :option2)
+            input
+          end,
+          [input],
+          global_options: [:option1]
+        )
+
+      {_, predict_fn} = Axon.build(model, global_layer_options: [option1: true, option2: true])
+
+      params = %{}
+      input = random({1, 1}, type: {:f, 32})
+      predict_fn.(params, input)
+    end
+  end
 end