Refactor containers (#590)

* Refactor containers

* Fix warning

lib/axon.ex

@@ -279,6 +279,8 @@ defmodule Axon do
  alias __MODULE__, as: Axon
  alias Axon.Parameter
 
+ import Axon.Shared
+
  require Logger
 
  @type t :: %__MODULE__{}
@@ -380,15 +382,6 @@ defmodule Axon do
  }
  end
 
- defp split_inputs(:container, [inputs]) do
- {inputs, cache} =
- deep_map_reduce(inputs, %{}, fn %Axon{output: id, nodes: nodes}, cache ->
- {id, Map.merge(nodes, cache)}
- end)
-
- {[inputs], [], [:layer], cache}
- end
-
  defp split_inputs(_op, inputs) do
  Enum.reduce(inputs, {[], [], [], %{}}, fn
  %Axon{output: layer_input, nodes: nodes}, {layers, params, args, cache} ->
@@ -704,62 +697,47 @@ defmodule Axon do
  @doc type: :special
  def container(container, opts \\ []) do
  opts = Keyword.validate!(opts, [:name, :meta])
-
- layer(:container, [container], name: opts[:name], meta: opts[:meta], op_name: :container)
+ {structure_fn, nodes} = destructure(container)
+ layer(structure_fn, nodes, name: opts[:name], meta: opts[:meta], op_name: :container)
  end
 
- # TODO: This should not be duplicated
- defp deep_new(%Nx.Tensor{} = x, fun), do: fun.(x)
-
- defp deep_new(x, fun) when is_number(x), do: fun.(x)
-
- defp deep_new(map, fun) do
- {cont, :ok} = Nx.Container.traverse(map, :ok, &recur_traverse(&1, &2, fun))
- cont
+ defp destructure(container) do
+ {structure, {nodes, _}} = recur_destructure(container, {[], 0})
+ fun = restructure(length(nodes) + 1, structure)
+ {fun, Enum.reverse(nodes)}
  end
 
- defp recur_traverse(item, :ok, fun) do
- case item do
- %Axon{} = t ->
- {fun.(t), :ok}
-
- %{axon: :axon} = t ->
- {fun.(t), :ok}
+ defp recur_destructure(container, acc) do
+ Nx.Container.traverse(container, acc, fn value, {leaves, idx} ->
+ case value do
+ %Axon{} = leaf ->
+ {idx, {[leaf | leaves], idx + 1}}
 
- container ->
- {deep_new(container, fun), :ok}
- end
+ container ->
+ recur_destructure(container, {leaves, idx})
+ end
+ end)
  end
 
- defp deep_merge(left, right, fun) do
- case Nx.Container.traverse(left, leaves(right), &recur_merge(&1, &2, fun)) do
- {merged, []} ->
- merged
+ for i <- 0..128 do
+ args = Macro.generate_arguments(i, __MODULE__)
 
- {_merged, _leftover} ->
- raise ArgumentError,
- "unable to merge arguments with incompatible" <>
- " structure"
+ defp restructure(unquote(i), structure) do
+ fn unquote_splicing(args) ->
+ args_tuple = {unquote_splicing(args)}
+ {container, :ok} = recur_restructure(structure, args_tuple)
+ container
+ end
  end
  end
 
- defp leaves(container) do
- container
- |> Nx.Container.reduce([], fn x, acc -> [x | acc] end)
- |> Enum.reverse()
- end
-
- defp recur_merge(left, [right | right_leaves], fun) do
- case {left, right} do
- {%Nx.Tensor{} = left, %Nx.Tensor{} = right} ->
- {fun.(left, right), right_leaves}
-
- {%Axon{} = left, %Axon{} = right} ->
- {fun.(left, right), right_leaves}
-
- {left, right} ->
- {deep_merge(left, right, fun), right_leaves}
- end
+ defp recur_restructure(structure, args_tuple) do
+ Nx.Container.traverse(structure, :ok, fn value, :ok ->
+ case value do
+ idx when is_integer(idx) -> {elem(args_tuple, idx), :ok}
+ container -> recur_restructure(container, args_tuple)
+ end
+ end)
  end
 
  @doc """
@@ -3644,35 +3622,31 @@ defmodule Axon do
  end
 
  @doc """
- Returns a model's output shape from the given input
+ Returns a model's output template from the given input
  template.
+
+ The output template gives you access to the output shape
+ and type of the given input graph.
  """
  @doc type: :graph
  def get_output_shape(%Axon{} = axon, inputs, opts \\ []) do
  {init_fn, forward_fn} = build(axon, opts ++ [raise_on_none: false])
 
- out =
+ inputs =
+ case inputs do
+ %Nx.Tensor{} = input -> Nx.to_template(input)
+ inputs when is_map(inputs) -> Map.new(inputs, fn {k, v} -> {k, Nx.to_template(v)} end)
+ end
+
+ fun =
  Nx.Defn.jit(
  fn inputs ->
  forward_fn.(init_fn.(inputs, Axon.ModelState.empty()), inputs)
  end,
  compiler: Axon.Defn
- ).(inputs)
-
- safe_shape(out)
- end
-
- defp safe_shape(container_or_tensor) do
- case container_or_tensor do
- %Axon.None{} = none ->
- none
-
- %Nx.Tensor{} = tensor ->
- Nx.shape(tensor)
+ )
 
- container ->
- deep_new(container, &Nx.shape/1)
- end
+ deep_new(apply(fun, [inputs]), &Nx.to_template/1)
  end
 
  @doc """
@@ -3850,74 +3824,17 @@ defmodule Axon do
  if MapSet.member?(visited, id) do
  {acc, visited}
  else
- %{op: op, parent: parents} = parent = nodes[id]
+ %{parent: parents} = parent = nodes[id]
 
  {acc, visited} =
- case op do
- :container ->
- [container] = parents
-
- deep_reduce(container, {acc, visited}, fn pid, {acc, visited} ->
- traverse_nodes(pid, nodes, acc, visited)
- end)
-
- _ ->
- Enum.reduce(parents, {acc, visited}, fn pid, {acc, visited} ->
- traverse_nodes(pid, nodes, acc, visited)
- end)
- end
+ Enum.reduce(parents, {acc, visited}, fn pid, {acc, visited} ->
+ traverse_nodes(pid, nodes, acc, visited)
+ end)
 
  {[parent | acc], MapSet.put(visited, id)}
  end
  end
 
- # TODO: Do not duplicate
- defp deep_reduce(item, acc, fun) when is_integer(item) do
- fun.(item, acc)
- end
-
- defp deep_reduce(map, acc, fun) do
- Nx.Container.reduce(map, acc, &recur_deep_reduce(&1, &2, fun))
- end
-
- defp recur_deep_reduce(value, acc, fun) do
- case value do
- %Axon{} = val ->
- fun.(val, acc)
-
- %Nx.Tensor{} = val ->
- fun.(val, acc)
-
- %{axon: :axon} = val ->
- fun.(val, acc)
-
- val when is_integer(val) ->
- fun.(val, acc)
-
- val ->
- deep_reduce(val, acc, fun)
- end
- end
-
- defp deep_map_reduce(leaf, acc, fun) when is_integer(leaf), do: fun.(leaf, acc)
-
- defp deep_map_reduce(container, acc, fun) do
- Nx.Container.traverse(container, acc, &recur_deep_map_reduce(&1, &2, fun))
- end
-
- defp recur_deep_map_reduce(leaf, acc, fun) do
- case leaf do
- %Axon{} = leaf ->
- fun.(leaf, acc)
-
- %Nx.Tensor{} = leaf ->
- fun.(leaf, acc)
-
- container ->
- deep_map_reduce(container, acc, fun)
- end
- end
-
  @doc """
  Pops the top node off of the graph.
 

lib/axon/compiler.ex

@@ -40,7 +40,6 @@ defmodule Axon.Compiler do
  @moduledoc false
  require Logger
 
- import Axon.Shared
  alias Axon.StatefulOutput
 
  ## Init JIT Compilation
@@ -549,72 +548,6 @@ defmodule Axon.Compiler do
  {id, model_funs, cache, op_counts, block_cache, model_state_meta}
  end
 
- defp recur_model_funs(
- %Axon.Node{id: id, op: :container, parent: [parents]},
- nodes,
- cache_and_counts,
- config
- ) do
- {parent_ids, {cache, op_counts, block_cache, model_state_meta}} =
- deep_map_reduce(parents, cache_and_counts, &to_model_funs(&1, nodes, &2, config))
-
- op_counts = Map.update(op_counts, :container, 1, fn x -> x + 1 end)
-
- predict_fun = fn params, inputs, state, cache, result_cache, fn_stacktrace ->
- {input, {state, result_cache, none?}} =
- deep_map_reduce(
- parent_ids,
- {state, result_cache, false},
- fn parent_id, {state, result_cache, none?} ->
- {input, {state, result_cache}} =
- call_predict_cache(
- parent_id,
- params,
- inputs,
- state,
- cache,
- result_cache,
- fn_stacktrace
- )
-
- none? = none? or propagating_none?(input)
- {input, {state, result_cache, none?}}
- end
- )
-
- input = if none?, do: %Axon.None{}, else: input
-
- {input, {state, result_cache}}
- end
-
- init_fun = fn template, cache, result_cache, fn_stacktrace, keys ->
- {parent_template, {parent_params, result_cache, none?}} =
- deep_map_reduce(parent_ids, {%{}, result_cache, false}, fn
- parent_id, {params, result_cache, none?} ->
- {parent_template, {params, result_cache}} =
- call_init_cache(
- parent_id,
- template,
- params,
- cache,
- result_cache,
- fn_stacktrace,
- keys
- )
-
- none? = none? or propagating_none?(parent_template)
- {parent_template, {params, result_cache, none?}}
- end)
-
- parent_template = if none?, do: %Axon.None{}, else: parent_template
-
- {parent_template, {parent_params, result_cache}}
- end
-
- model_funs = %{predict: predict_fun, init: init_fun}
- {id, model_funs, cache, op_counts, block_cache, model_state_meta}
- end
-
  defp recur_model_funs(
  %Axon.Node{
  id: id,

lib/axon/display.ex

@@ -3,7 +3,6 @@ defmodule Axon.Display do
  Module for rendering various visual representations of Axon models.
  """
 
- import Axon.Shared
  alias Axon.Parameter
 
  @compile {:no_warn_undefined, TableRex.Table}
@@ -94,7 +93,8 @@ defmodule Axon.Display do
  defp do_axon_to_rows(
  %Axon.Node{
  id: id,
- op: :container,
+ op: structure,
+ op_name: :container,
  parent: [parents],
  name: name_fn
  },
@@ -105,7 +105,7 @@ defmodule Axon.Display do
  model_info
  ) do
  {input_names, {cache, op_counts, model_info}} =
- deep_map_reduce(parents, {cache, op_counts, model_info}, fn
+ Enum.map_reduce(parents, {cache, op_counts, model_info}, fn
  parent_id, {cache, op_counts, model_info} ->
  {_, name, _shape, cache, op_counts, model_info} =
  axon_to_rows(parent_id, nodes, templates, cache, op_counts, model_info)
@@ -119,7 +119,7 @@ defmodule Axon.Display do
  shape = Axon.get_output_shape(%Axon{output: id, nodes: nodes}, templates)
 
  row = [
- "#{name} ( #{op_string} #{inspect(input_names)} )",
+ "#{name} ( #{op_string} #{inspect(apply(structure, input_names))} )",
  "#{inspect({})}",
  "#{inspect(shape)}",
  render_options([]),
@@ -311,27 +311,6 @@ defmodule Axon.Display do
  end
  end
 
- defp recur_axon_to_edges(
- %Axon.Node{id: id, op: :container, name: name_fn, parent: [parents]},
- nodes,
- templates,
- cache_counts_edgelist
- ) do
- {node_inputs, {cache, op_counts, edgelist}} =
- deep_map_reduce(parents, cache_counts_edgelist, &axon_to_edges(&1, nodes, templates, &2))
-
- name = name_fn.(:container, op_counts)
- node_shape = Axon.get_output_shape(%Axon{output: id, nodes: nodes}, templates)
- to_node = %{axon: :axon, id: id, op: :container, name: name, shape: node_shape}
-
- new_edgelist =
- deep_reduce(node_inputs, edgelist, fn from_node, acc ->
- [{from_node, to_node} | acc]
- end)
-
- {to_node, {cache, op_counts, new_edgelist}}
- end
-
  defp recur_axon_to_edges(
  %Axon.Node{id: id, op_name: op, name: name_fn, parent: parents},
  nodes,