IsochoricParameter

class hypercoil.init.mapparam.IsochoricParameter(model: PyTree, *, where: Callable = <function where_weight>, volume: float = 1.0, max_condition: Optional[float] = None, softmax_temp: Optional[float] = None, spd_threshold: float = 0.001)

Parameter representing an ellipsoid of fixed volume.

Note

This parameter mapping requires the original parameter to be square and symmetric along its last 2 axes. The parameter is further forced to be positive definite.

Parameters:

modelPyTree

The model to which the parameter belongs.

whereCallable

As in equinox.tree_at: a function that takes a model (or generally a PyTree) and returns the parameter tensor to be mapped. For example: where = lambda mlp: mlp.layers[-1].linear.weight. By default, the weight attribute of the model is retrieved.

volumenonnegative float (default 1)

Parameter volume. The determinant of the parameter is set to this value.

max_conditionfloat \(\in [1, \infty)\) or None (default None)

Maximum permissible condition number of the parameter. This can be used to constrain the eccentricity of isochoric ellipsoids. To enforce this maximum, the eigenvalues of the original parameter are replaced with a convex combination of the original eigenvalues and a vector of ones such that the largest eigenvalue is no more than max_condition times the smallest eigenvalue. Note that a max_condition of 1 will always return (a potentially isotropically scaled) identity.

softmax_tempfloat or None (default None)

If this is provided, then the eigenvalues of the original parameter are passed through a softmax with the specified temperature before any other processing.

Attributes:

max_condition

softmax_temp

Methods

image_map_impl
preimage_map_impl