Add parameter constraints

[seanmor5]

Closes #111.

Axon had no way to keep a trainable parameter inside a feasible region during training. Keras offers this as constraints: projections such as max_norm or non_neg that run on a weight right after every optimizer update. This PR adds the same mechanism to Axon, following the Keras (post-update projection) semantics rather than PyTorch-style parametrizations that reparameterize the forward pass.

Design

Constraints mirror how frozen parameters already work: they are metadata on %Axon.ModelState{} that the compiler collects from the graph and that the training step consumes through Axon.ModelState transforms.

• Axon.Constraints is a new module shaped like Axon.Initializers. max_norm/1, non_neg/0, unit_norm/1 and min_max_norm/1 return arity-1 functions that take a parameter tensor and return the projected tensor. Norm-based constraints take :axes (default [0], the incoming weight vector of each output unit for a dense kernel). Their implementations are defnps so they trace into the training step.
• Axon.param/3 and Axon.parameter/3 accept a :constraint option: one of the atoms above (resolved to the Axon.Constraints function with default options) or an arity-1 function. It is validated like :initializer and stored on %Axon.Parameter{}.
• %Axon.ModelState{} gains a constraints field (a keep field, nested %{layer => %{param => fun}} like data). The compiler collects it next to frozen_parameters, nests it for blocks, and merge_model_state! carries constraints from a user-supplied initial model state so they compose with Axon.Loop.run/4.
• Axon.ModelState.constrain/3 attaches a constraint to every parameter matched by a path mask (same mask style as freeze/2), which is how you constrain built-in layers such as Axon.dense. Axon.ModelState.apply_constraints/1 projects every constrained, trainable, non-frozen parameter.
• Axon.Loop.train_step/4 calls apply_constraints/1 right after Axon.ModelState.update/3, so every optimizer update is followed by the projection.

Usage

On a custom layer parameter:

w = Axon.param("w", {32, 64}, constraint: Axon.Constraints.max_norm(max: 2.0))
Axon.layer(fn x, w, _opts -> Nx.dot(x, w) end, [input, w])

On the parameters of built-in layers, by path:

{init_fn, _} = Axon.build(model)
model_state = init_fn.(template, Axon.ModelState.empty())

model_state =
  Axon.ModelState.constrain(
    model_state,
    &match?([_, "kernel"], &1),
    Axon.Constraints.unit_norm()
  )

model
|> Axon.Loop.trainer(:mean_squared_error, :sgd)
|> Axon.Loop.run(data, model_state)

Tradeoffs and limitations

• Constraints are projections applied after each update (Keras semantics). They are not applied at initialization, never to frozen parameters, and never to layer state such as batch norm statistics.
• Constraint functions live in model state metadata, so they are part of the jit cache key and travel with checkpoints. Axon.Loop.serialize_state/2 serializes them with term_to_binary; anonymous constraints only deserialize when the module that defined them is available (the Axon.Constraints functions are fine).
• No graph-level Axon.constrain/3 is added since the graph-level Axon.freeze/2 is deprecated in favor of the Axon.ModelState API, and no per-layer kernel_constraint/bias_constraint options are added.
• Pre-existing gap, not changed here: merge_model_state! still drops the frozen_parameters of a user-supplied initial model state at init time. This PR only carries constraints across; carrying frozen_parameters the same way is a candidate follow-up.

Tests

• test/axon/constraints_test.exs: doctests plus numeric checks for all four constraints (including custom :axes, :rate, type preservation for bf16, and use inside defn).
• test/axon_test.exs: :constraint validation and storage on Axon.param/3 and Axon.parameter/3.
• test/axon/model_state_test.exs: constrain/3 path masking and composition, apply_constraints/1 projecting only constrained parameters, skipping frozen ones, and both working inside defn.
• test/axon/compiler_test.exs: constraint metadata is collected for custom layers, nested for blocks, and preserved from an initial model state.
• test/axon/loop_test.exs: train_step projects a parameter after SGD (and leaves it alone when frozen), Axon.Loop.run/4 honors constraints attached to the initial model state, and constraints survive serialize_state/deserialize_state.

🤖 Generated with Claude Code

[polvalente]

eps should be an option IMO

[polvalente]

norm should be an option as well. Otherwise the docs should make it clear we're talking euclidean norm

[polvalente]

let's use full names

Suggested change

	def non_neg() do
	def non_negative() do

[polvalente]

Again, which norm is used should be an option

[polvalente]

Why not use Nx.LinAlg.norm?