coreai_opt.pruning.config.OpMagnitudePrunerConfig¶

class coreai_opt.pruning.config.OpMagnitudePrunerConfig[source]¶

Bases: WeightOnlyOpValidationMixin, OpCompressionConfig[PruningSpec]

Operation-level pruning configuration.

Pruning is a weight-only compression technique. Only op_state_spec is used to configure which state tensors (e.g. weights) to prune.

op_state_spec¶

Mapping of parameter names to their pruning specs. Default includes "weight" and "in_proj_weight" at 50 % sparsity.

Type:: dict[str, PruningSpec | None]

Example

>>> config = OpMagnitudePrunerConfig()
>>> config = OpMagnitudePrunerConfig(
...     op_state_spec={"weight": PruningSpec(target_sparsity=0.75)}
... )

classmethod get_default_state_spec()[source]¶

Provide default state spec for pruning.

Return type:: dict[str, PruningSpec | None]