Block Specification

Block-level sparsity specification.

class sparsekit.block.BlockSpec(view, shape, name=None)

Bases: SparseBlock

Treats the entire tensor as a grid of blocks.

Parameters:

• view (View)

• shape (Tuple[int, ...])

• name (str | None)

param

Parameter/View providing shape, stride, and write-through data access.

block_shape

Shape of each block in the grid.

name

Optional name for identification.

Type:

str | None

property ndim: int

Number of dimensions in the underlying tensor.

property data: Tensor

Raw tensor data of the underlying Parameter.

property grid_shape: Tuple[int, ...]

Full grid shape including singleton dimensions.

numel_per_block()

Number of elements per block.

Return type:

int

blocks()

Return self as the only BlockSpec.

Return type:

Iterable[SparseBlock]

parameters()

List containing the single underlying Parameter.

Return type:

List[View]

set_data(data)

Copy data into the underlying Parameter tensor.

nnz(eps=1e-08)

Number of non-zero elements (within tolerance).

Return type:

int

block_view(values, reorder=True, merge=False)

Reshape tensor to interleaved block view.

Parameters:

• values (Tensor | Mapping[Any, Tensor] | None) – Input values (None uses param.data).

• reorder (bool) – If True, permute grid dims before block dims.

• merge – If True, collapse block dims to trailing dim.

Return type:

Tensor

expand_block_tensor(block_values)

Convert grid-shaped tensor to full grid shape with singletons.

Parameters:

block_values (Tensor) – Tensor with shape block_grid_shape.

Returns:

Tensor reshaped to grid_shape.

Return type:

Tensor

broadcast_block_to_element(block_values, fake=False)

Broadcast block grid-shaped tensor to full tensor shape.

Parameters:

• block_values (Tensor) – Tensor with shape grid_shape.

• fake – If True, only unsqueeze without repeating.

Returns:

Tensor with shape self.shape (or interleaved if fake=True).

Return type:

Tensor

apply_mask(mask)

Zero out blocks where mask is True.

Parameters:

mask (Tensor)

apply_multiplier(multiplier)

Multiply each block by corresponding scalar in multiplier.

Parameters:

multiplier (Tensor)

reduce(values, reduce_fn)

Apply reduce_fn over each block and return grid-shaped result.

Parameters:

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

• reduce_fn (Callable[[Tensor], Tensor])

Return type:

Tensor

get_masks(block_masks)

Convert block-level mask to element-level mask.

Parameters:

block_masks (Tensor) – Boolean tensor with shape block_grid_shape.

Returns:

Dict mapping self to the broadcasted element mask.

Return type:

Mapping[BlockSpec, Tensor]

class sparsekit.block.BlockCoupling(specs, orders, name=None)

Bases: SparseBlock

Merges multiple BlockSpec objects into one coupled sparse node.

Parameters:

• specs (List[BlockSpec])

• orders (List[Tuple[int, ...]])

• name (str | None)

specs

List of BlockSpec objects to couple.

Type:

List[sparsekit.block.BlockSpec]

orders

Axis permutations to align block grids.

Type:

List[Tuple[int, …]]

name

Optional name for identification.

Type:

str | None

property shape: Tuple[int, ...]

Return placeholder shape (-1, -1) for coupled specs.

property data: Mapping[BlockSpec, Tensor]

Dict mapping each spec to its tensor data.

property grid_shape: Tuple[int, ...]

Grid shape for the coupling (after order permutation).

numel_per_block()

Total elements per block across all specs.

Return type:

int

parameters()

List of all Parameter objects from coupled specs.

Return type:

List[View]

nnz(eps=1e-08)

Count non-zero elements across all specs.

Return type:

int

block_view(values, reorder=True, merge=False)

Return concatenated block view across all specs.

Parameters:

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

• reorder (bool)

Return type:

Tensor

reduce(values, reduce_fn)

Apply reduce_fn over concatenated block view.

Parameters:

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

• reduce_fn (Callable[[Tensor], Tensor])

Return type:

Tensor

get_masks(block_masks)

Convert block-level mask to element-level masks for each spec.

Parameters:

block_masks (Tensor)

Return type:

Mapping[BlockSpec, Tensor]

apply_mask(mask)

Zero out blocks where mask is True across all specs.

Parameters:

mask (Tensor)

apply_multiplier(multiplier)

Multiply each block by corresponding scalar across all specs.

Parameters:

multiplier (Tensor)