Returns an empty model state.

Freezes parameters and state in the given model state using the given mask.

The mask is an arity 1 function which takes the access path to the leaf parameter and returns true if the parameter should be frozen or false otherwise. With this, you can construct flexible masking policies:

fn
  ["dense_" <> n, "kernel"] -> String.to_integer(n) < 3
  _ -> false
end

The default mask returns true for all paths, and is equivalent to freezing the entire model.

Returns the frozen parameters in the given model state.

Returns the frozen state in the given model state.

Merges 2 states with function.

Returns a new model state struct from the given parameter map.

Ties a parameter to another parameter, enabling weight sharing.

The destination parameter will reference the source parameter's tensor, optionally applying a transformation. Both destination and source are access paths (lists of strings) into the model state data.

Options

• :transform - a function to transform the source tensor before use at the destination. For example, &Nx.transpose/1 for tying an embedding layer to an output projection.

Examples

# Tie output projection to embedding weights (transposed)
model_state = Axon.ModelState.tie(
  model_state,
  ["output", "kernel"],
  ["embed", "kernel"],
  transform: &Nx.transpose/1
)

Returns the trainable parameters in the given model state.

Returns the trainable state in the given model state.

Unfreezes parameters and state in the given model state using the given mask.

The mask is an arity 1 function which takes the access path to the leaf parameter and returns true if the parameter should be unfrozen or false otherwise. With this, you can construct flexible masking policies:

fn
  ["dense_" <> n, "kernel"] -> n < 3
  _ -> false
end

The default mask returns true for all paths, and is equivalent to unfreezing the entire model.

Updates the given model state.