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.

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.