Add option to not raise if output is none, resolves #538

lib/axon.ex

@@ -3405,7 +3405,7 @@ defmodule Axon do
  """
  @doc type: :graph
  def get_output_shape(%Axon{} = axon, inputs, opts \\ []) do
- {init_fn, forward_fn} = build(axon, opts)
+ {init_fn, forward_fn} = build(axon, opts ++ [raise_on_none: false])
 
  out =
  Nx.Defn.jit(

lib/axon/compiler.ex

@@ -48,6 +48,7 @@ defmodule Axon.Compiler do
  @doc false
  def build(%Axon{output: id, nodes: nodes}, opts) do
  debug? = Keyword.get(opts, :debug, false)
+ raise_on_none? = Keyword.get(opts, :raise_on_none, true)
  mode = Keyword.get(opts, :mode, :inference)
  seed = Keyword.get_lazy(opts, :seed, fn -> :erlang.system_time() end)
  global_layer_options = Keyword.get(opts, :global_layer_options, [])
@@ -105,10 +106,12 @@ defmodule Axon.Compiler do
  end
 
  with %Axon.None{} <- result do
- raise ArgumentError,
- "the compiled model will always result in %Axon.None{}." <>
- " This most likely means you specified optional output and " <>
- " did not handle the case when it is missing"
+ if raise_on_none? do
+ raise ArgumentError,
+ "the compiled model will always result in %Axon.None{}." <>
+ " This most likely means you specified optional output and " <>
+ " did not handle the case when it is missing"
+ end
  end
 
  result

test/axon_test.exs

@@ -1076,5 +1076,19 @@ defmodule AxonTest do
  assert shape = Axon.get_output_shape(model, Nx.template({1, 1}, :f32))
  assert shape == {{1, 2}, {1, 2}}
  end
+
+ test "doesn't raise on none output" do
+ values = Axon.input("values")
+ mask = Axon.input("mask", optional: true)
+
+ model =
+ values
+ |> Axon.dense(10)
+ |> Axon.multiply(mask)
+ |> Axon.dense(1)
+ |> Axon.sigmoid()
+
+ assert %Axon.None{} = Axon.get_output_shape(model, %{"values" => Nx.template({1, 1}, :f32)})
+ end
  end
 end