kd.metrics.TreeMap#

class kauldron.metrics.TreeMap(*, metric: kauldron.metrics.base.Metric)[source]

Bases: kauldron.metrics.base._TreeMetric

Maps an inner metric to a pytree and returns a pytree of results.

Example

train_metrics[“param_norm”] = TreeMap(metric=Norm(tensor=”params”)) would separately track and return the norm of all parameters of the model.

metric

Any metric to apply to the leaves of the tree. Also uses the keys of that metric to resolve the kwargs.

class State( tree: "Mapping['str', PyTree[base_state.State]]" = FrozenDict({}), *, parent: '_MetricT' = <_EMPTY_TYPE.EMPTY: 1>, )[source]

Bases: kauldron.metrics.base.TreeState

merge( other: kauldron.metrics.base.TreeState, ) → kauldron.metrics.base.TreeState[source]: Merge two trees of metric states.

replace(**updates): Returns a new object replacing the specified fields with new values.