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