Scope Specification

Scope-level sparsity specification over block grid.

class sparsekit.scope.ScopeSpec(block, shape, name=None)[source]

Bases: SparseScope

Organizes blocks from a SparseNode into scopes (decision units).

Divides the block grid into scopes of shape blocks. Pruning decisions (hard/soft threshold, mask selection) operate at the scope level: within each scope, blocks compete to survive.

Parameters:

  • block (SparseBlock) – SparseBlock or BlockCoupling that defines the block grid. Must have grid_shape divisible by shape.

  • shape (Tuple[int, ...]) – Number of blocks of scope in each dimension. Use -1 to span the entire dimension.

  • name (str | None) – Optional name for identification.

block: SparseBlock
shape: Tuple[int, ...]
name: str | None = None
property grid_shape: Tuple[int, ...]

Full grid shape including singleton dimensions.

numel_per_scope()[source]

Number of blocks per scope.

Return type:

int

numblk_per_scope()[source]

Number of blocks per scope.

Return type:

int

blocks()[source]
Return type:

Iterable[SparseBlock]

nnz(eps=1e-08)[source]
Return type:

int

block_to_scope(b, reorder=True, merge=False)[source]

Reshape a block-grid tensor into scope layout.

Parameters:

  • b (Tensor) – Tensor with shape (B1, B2, ..., Bm, ...).

  • reorder (bool) – If True, permute so scope dims precede block dims.

  • merge (bool) – If True, collapse block dims into a single trailing dim.

Returns:

Tensor with shape (G1, G2, ..., g1, g2, ..., ...) (or merged).

Return type:

Tensor

scope_to_block(block_values)[source]

Broadcast scope-level values back to the block grid.

Parameters:

block_values (Tensor) – Tensor with shape (G1, ..., Gm).

Returns:

Tensor with shape (B1, ..., Bm) where Bi = Gi * gi.

Return type:

Tensor

block_norms(values)[source]

Compute block L2 norms arranged in scope layout.

Parameters:

values (Tensor | Mapping[Any, Tensor] | None) – Element values to compute norms from (None uses param data).

Returns:

Tensor with shape (*grid_shape, block_numel).

Return type:

Tensor

kth_largest(element_values, nnz)[source]

Calculates the k-th largest score across all blocks from all specs. This is used to determine the threshold for pruning.

Parameters:
Return type:

Tensor

kth_mid(element_values, nnz, k_weight=1.0)[source]

Calculates the k-th largest score across all blocks from all specs. This is used to determine the threshold for pruning.

Parameters:
Return type:

Tensor

hard_threshold(thresholds=None, nnz=None, values=None)[source]

Zero out blocks in-place based on thresholds.

Exactly one of thresholds or nnz must be given.

Parameters:

  • thresholds (Tensor | None) – Per-block thresholds, shape grid_shape.

  • nnz (int | None) – Number of blocks to keep per scope.

  • values (Tensor | Mapping[Any, Tensor] | None) – Element values for computing norms.

get_masks(nnz, values=None, block_scores=None, block_mask=None, **kwargs)[source]

Compute element-level boolean masks from scope-level scores.

Parameters:

  • nnz (int) – Number of blocks to keep per scope.

  • block_scores (Tensor | None) – Pre-computed scores with shape (*grid_shape, block_numel).

  • values (Tensor | Mapping[Any, Tensor] | None) – Element values for computing block norms (if scores not given).

  • block_mask (Tensor | None) – Pre-computed boolean mask to use directly.

Returns:

Dict mapping each SparseBlock to its element-level boolean mask.

Return type:

Mapping[SparseBlock, Tensor]

soft_threshold(thresholds, conditioners=None, scale=False, max_iter=20, atol=1e-08, eps=None)[source]

Apply soft thresholding (L1 proximal operator) to scopes in-place.

Parameters:

  • thresholds (Tensor) – Per-scope threshold values with shape grid_shape.

  • conditioners (Tensor | Mapping[Any, Tensor] | None) – Diagonal preconditioner per SparseBlock.

  • scale (bool) – If True, scale thresholds by sqrt(block_numel).

  • max_iter (int) – Maximum bisection iterations for Adam variant.

  • eps (float | None) – Small constant for numerical stability.

  • atol (float) – Absolute tolerance for convergence.

Return type:

None

class sparsekit.scope.ScopeCoupling(scopes, orders, name=None)[source]

Bases: SparseScope

Couples multiple ScopeSpec instances for joint pruning.

Aligns scope grids from different parameters via dimension permutations (orders) so they share a common grid shape. Within each aligned scope, blocks from all specs compete to survive pruning.

Parameters:

  • scopes (List[ScopeSpec]) – List of ScopeSpec instances to couple.

  • orders (List[Tuple[int, ...]]) – Dimension permutations to align each scope’s grid. Identity permutation if omitted.

  • name (str | None) – Optional name for identification.

scopes: List[ScopeSpec]
orders: List[Tuple[int, ...]]
name: str | None = None
numblk()[source]

Total number of blocks across all coupled scopes.

Return type:

int

property params: Set[View]

Expose underlying views for optimizer integration.

blocks()[source]

All SparseBlock instances across all child scopes.

Return type:

Iterable[SparseBlock]

property grid_shape: Tuple[int, ...]

Reference scope grid shape (after order permutation).

numblk_per_scope()[source]

Total blocks per scope across all coupled scopes.

Return type:

int

numel_per_scope()[source]

Number of elements per scope.

Return type:

int

nnz(eps=1e-08)[source]
Return type:

int

specs()[source]
Return type:

Iterable[SparseBlock]

block_norms(values)[source]
Parameters:

values (Tensor | Mapping[Any, Tensor] | None)

Return type:

Tensor

kth_largest(k, values)[source]

Calculates the k-th largest score across all blocks from all specs. This is used to determine the threshold for pruning.

Parameters:
Return type:

Tensor

hard_threshold(thresholds=None, nnz=None, values=None)[source]

Compute kappa-largest block norm among coupled scopes from all specs then sends the threshold to specs to hard-threshold in-place. Note that the threshold is across coupled scopes, so some parameters might be pruned more than others (it’s expected).

Parameters:
soft_threshold(thresholds, conditioners=None, scale=False, max_iter=20, atol=1e-08, eps=None)[source]

Performs soft thresholding on all coupled parameters.

Parameters:
Return type:

None

get_masks(nnz, block_scores=None, values=None, **kwargs)[source]
Parameters:
Return type:

Mapping[SparseBlock, Tensor]