Add regularization features

[fjclark]

Description

This adds optional regularization information to AttributeConfig and Trainable. Specifically, AttributeConfig gets:

    regularize: dict[str, float] = pydantic.Field(
        {},
        description="The regularization strength to apply to each parameter, e.g. "
        "'k': 0.01, 'epsilon': 0.001. Parameters not listed are not regularized.",
    )

and Trainable gets regularized_idxs (where the idxs are within the values returned by to_values) and regularization_weights properties, which allow you to apply regularization with e.g.

initial_parameters = trainable.to_values() 
# Copy and train to get trained_parameters

regularized_idxs = trainable.regularized_idxs
regularization_weights = trainable.regularization_weights

l2_reg_loss = (
        ((trained_parameters[regularized_idxs] - initial_parameters[regularized_idxs]) ** 2)
        * regularization_weights
    ).sum()

Status

• Ready to go

[jthorton]

LGTM the tests look good, it would be great if coverage was working to know if anything was missed though!

[fjclark]

Thanks Josh! Yes, don't think I have permissions to fix this though (thanks for raising the issue). I've checked coverage locally, and everything inside Trainable is covered. Uncovered lines are either in pydantic 1 blocks or weren't covered previously.

[lilyminium]

Just closing and re-opening to kick CI, since it looks like that never ran.

[fjclark]

Thanks Lily. Merging in #78 to get the CI running.

[lilyminium]

LGTM, at some point it would be good generate some tutorials/docs on how to use this, but not in scope atm. Just one question about Interchange!

[lilyminium]

Thanks Finlay!