constraints.make_bounds_projection
constraints.make_bounds_projection(metadata)Return an optax transform that clamps params to their declared bounds.
Reads bounds from each Parameter in metadata and enforces them after every gradient step. Chain after the main optimizer:
optimizer = optax.chain(
optax.adamw(lr),
make_bounds_projection(metadata),
)Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| metadata | dict[str, Parameter] | Name-keyed Parameter objects carrying bounds info. |
required |
Returns
| Name | Type | Description |
|---|---|---|
| optax.GradientTransformationExtraArgs | An optax.GradientTransformationExtraArgs that clamps updates. | |
| optax.GradientTransformationExtraArgs | Requires params to be passed to update. |