Add mixed precision (elixir-nx#74)

zblanco · Jun 20, 2021 · 103a1f3 · 103a1f3
1 parent 6e94494
commit 103a1f3
Show file tree

Hide file tree

Showing 8 changed files with 511 additions and 126 deletions.
diff --git a/lib/axon.ex b/lib/axon.ex
@@ -78,7 +78,7 @@ defmodule Axon do
   @type t :: %__MODULE__{}
 
   @doc false
-  defstruct [:id, :name, :output_shape, :parent, :op, :params, :opts]
+  defstruct [:id, :name, :output_shape, :parent, :op, :params, :policy, :opts]
 
   @doc """
   Custom Axon layer with given parent.
@@ -125,6 +125,7 @@ defmodule Axon do
       parent: parent,
       op: op,
       params: parameters,
+      policy: Axon.MixedPrecision.create_policy(),
       opts: opts
     }
   end