Encoders

This module contains encoder utilities for sparse operations.

Pairwise Encoders

torchsparsegradutils.encoders.pairwise_encoder.calc_pairwise_coo_indices_nd(radius: float, volume_shape: Tuple[int, ...], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', dtype: dtype = torch.int64, device: device | None = device(type='cpu')) → Dict[Tuple[int, ...], Tensor][source]

Compute per-offset COO linear index pairs for an \((C,*S)\) volume.

For a volume (C, *spatial_dims) and spatial radius r, return a dictionary mapping each offset tuple (c, *spatial_offset) to a (2, M) tensor of linear index pairs [[i...],[j...]] such that the second row is the first row shifted by the offset (within bounds). Linearization follows row-major order (torch.arange(prod(volume_shape)).reshape(volume_shape).flatten()).

Offsets come from _gen_offsets_nd() (sorted), which enumerates spatial offsets with \(\|o_{spatial}\|_2 \le r\) and augments them with channel offsets according to channel_voxel_relation.

Parameters:

  • radius (float) – Neighborhood radius (>= 1).

  • volume_shape (tuple[int, ...]) – Shape (C, *spatial_dims) with at least one spatial dimension.

  • diag (bool, optional) – Include diagonal key (0,...,0) mapping to (i,i) pairs. Default False.

  • upper (bool or None, optional) – Forwarded sign filter (see _gen_offsets_nd()). Default None.

  • channel_voxel_relation ({'indep','intra','inter'}, optional) – Channel relation mode. Default 'indep'.

  • dtype (torch.dtype, optional) – Integer dtype of output index tensors (default torch.int64).

  • device (torch.device, optional) – Target device (default CPU).

Returns:

Mapping from offset tuple to a (2, M_o) tensor of linear index pairs.

Return type:

dict[tuple[int, …], torch.Tensor]

Raises:

ValueError – If arguments are inconsistent (e.g. radius < 1).

Notes

Each non-zero offset o yields pairs by trimming the index lattice twice with _trim_nd(): once by o and once by -o. Only valid in-bounds pairs are produced (no padding). Sorting matches _gen_offsets_nd().

See also

_gen_offsets_nd

Generate (sorted) offsets.

_trim_nd

Bounds-aware slicing used for forming pairs.

Examples

2D single channel: >>> idxs = calc_pairwise_coo_indices_nd( … radius=1.0, … volume_shape=(1, 3, 3), # (C,H,W) … diag=True, … upper=None, … channel_voxel_relation=’indep’, … ) >>> sorted(list(idxs.keys()))[:3] # doctest: +ELLIPSIS [(0, -1, 0), (0, 0, -1), (0, 0, 0)] >>> z = (0, 0, 0) >>> idxs[z].shape torch.Size([2, 9])

3D, inter-channel: >>> idxs3d = calc_pairwise_coo_indices_nd( … radius=1.0, … volume_shape=(2, 3, 3, 3), … channel_voxel_relation=’inter’, … ) >>> any(o[0] == 1 and o[1:] != (0, 0, 0) for o in idxs3d.keys()) True

torchsparsegradutils.encoders.pairwise_encoder.calc_pairwise_coo_indices(radius: float, volume_shape: Tuple[int, int, int, int], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', dtype: dtype = torch.int64, device: device = device(type='cpu')) → Dict[Tuple[int, int, int, int], Tensor][source]

3D wrapper for calc_pairwise_coo_indices_nd() (deprecated).

Deprecated since version 0.x: Use calc_pairwise_coo_indices_nd().

Parameters:

  • radius (float) – Spatial radius (>=1).

  • volume_shape (tuple[int,int,int,int]) – (C,H,D,W).

  • diag (bool, optional) – Include diagonal. Default False.

  • upper (bool or None, optional) – Sign-selection (see N-D version). Default None.

  • channel_voxel_relation ({'indep','intra','inter'}, optional) – Channel relation mode. Default 'indep'.

  • dtype (torch.dtype, optional) – Index dtype (default torch.int64).

  • device (torch.device, optional) – Target device (default CPU).

Returns:

Per-offset COO index pairs.

Return type:

dict[tuple[int,int,int,int], torch.Tensor]

Raises:

ValueError – If inputs are invalid.

torchsparsegradutils.encoders.pairwise_encoder.calc_pariwise_coo_indices(radius: float, volume_shape: Tuple[int, int, int, int], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', dtype: dtype = torch.int64, device: device = device(type='cpu')) → Dict[Tuple[int, int, int, int], Tensor]

3D wrapper for calc_pairwise_coo_indices_nd() (deprecated).

Deprecated since version 0.x: Use calc_pairwise_coo_indices_nd().

Parameters:

  • radius (float) – Spatial radius (>=1).

  • volume_shape (tuple[int,int,int,int]) – (C,H,D,W).

  • diag (bool, optional) – Include diagonal. Default False.

  • upper (bool or None, optional) – Sign-selection (see N-D version). Default None.

  • channel_voxel_relation ({'indep','intra','inter'}, optional) – Channel relation mode. Default 'indep'.

  • dtype (torch.dtype, optional) – Index dtype (default torch.int64).

  • device (torch.device, optional) – Target device (default CPU).

Returns:

Per-offset COO index pairs.

Return type:

dict[tuple[int,int,int,int], torch.Tensor]

Raises:

ValueError – If inputs are invalid.

class torchsparsegradutils.encoders.pairwise_encoder.PairwiseEncoder(radius: float, volume_shape: Tuple[int, ...], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', layout=torch.sparse_coo, indices_dtype: dtype = torch.int64, device: device = device(type='cpu'))[source]

Bases: Module

Encode pairwise spatial–channel neighborhoods as sparse tensors.

Precomputes a mapping from local neighborhoods (within spatial radius r) and optional channel interactions to global sparse matrix indices over the linearized volume. Useful for graph-like layers, covariance assembly, sparse attention and any operator exploiting local geometric structure.

Parameters:

  • radius (float) – Spatial neighborhood radius r.

  • volume_shape (tuple[int, ...]) – (C,*spatial_dims) (at least one spatial dimension).

  • diag (bool, optional) – Include diagonal (self-edges). Default False.

  • upper (bool or None, optional) – Triangular selection on offset set (first non-zero criterion). None keeps all.

  • channel_voxel_relation ({'indep','intra','inter'}, optional) – Channel interaction model. Default 'indep'.

  • layout (torch.layout, optional) – torch.sparse_coo (default) or torch.sparse_csr.

  • indices_dtype (torch.dtype, optional) – Integer dtype for indices (int32 or int64). Default int64.

  • device (torch.device, optional) – Device to store cached indices (default CPU).

volume_numel

C * prod(spatial_dims).

Type:

int

spatial_dims

len(volume_shape) - 1.

Type:

int

offsets

Ordered offsets (optionally with diagonal key first if diag).

Type:

list[tuple[int,…]]

indices

(COO) (2, nnz_total) tensor of linear index pairs.

Type:

torch.Tensor

crow_indices, col_indices, csr_permutation

(CSR) components & permutation to reorder values into CSR order.

Type:

torch.Tensor

Notes

  • Input to __call__() must have shape [(B), N, C, *S] where N == len(self.offsets). Batch dimension optional.

  • Edge handling uses trimming (no wrap, no padding).

  • CSR values are internally reordered via self.csr_permutation.

  • Complexity scales with number of offsets times valid pairs (≈ \(O(r^2)\) in 2D, \(O(r^3)\) in 3D).

See also

calc_pairwise_coo_indices_nd

Build per-offset COO indices.

convert_coo_to_csr_indices_values

COO→CSR conversion + permutation.

_gen_offsets_nd

Construct ordered offset set.

_trim_nd

Bounds-aware slicing for forming value blocks.

Examples

Basic 2D: >>> from torchsparsegradutils.encoders import PairwiseEncoder >>> encoder = PairwiseEncoder( … radius=1.5, … volume_shape=(3, 8, 8), … diag=True, … channel_voxel_relation=’indep’ … ) >>> encoder.volume_numel 192 >>> len(encoder.offsets) # doctest: +SKIP 13

Create sparse tensor from values: >>> values = torch.randn(len(encoder.offsets), 3, 8, 8) >>> sp = encoder(values) >>> sp.shape torch.Size([192, 192]) >>> sp.is_sparse True

Batched: >>> values_b = torch.randn(4, len(encoder.offsets), 3, 8, 8) >>> sp_b = encoder(values_b) >>> sp_b.shape torch.Size([4, 192, 192])

3D with inter-channel relations: >>> encoder3d = PairwiseEncoder( … radius=2.0, … volume_shape=(5, 16, 16, 16), … channel_voxel_relation=’inter’, … layout=torch.sparse_csr, … )

Upper-triangular (symmetric use-case): >>> sym = PairwiseEncoder( … radius=1.0, … volume_shape=(1, 10, 10), … upper=True, … diag=True, … ) >>> v = torch.randn(len(sym.offsets), 1, 10, 10) >>> _ = sym(v)

__init__(radius: float, volume_shape: Tuple[int, ...], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', layout=torch.sparse_coo, indices_dtype: dtype = torch.int64, device: device = device(type='cpu'))[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

property device
__call__(values: Tensor) → Tensor[source]

Construct sparse tensor (COO or CSR) from per-offset value blocks.

Parameters:

values (torch.Tensor) – Shape [(B), N, C, *spatial_dims] with optional batch B and N == len(self.offsets).

Returns:

Sparse tensor of shape [(B), S, S] where S = C * prod(spatial_dims).

Return type:

torch.Tensor

Raises:

ValueError – If shape, dtype or offset count are inconsistent.

Pairwise Voxel Encoders

class torchsparsegradutils.encoders.pairwise_voxel_encoder.PairwiseVoxelEncoder(radius: float, volume_shape: Tuple[int, int, int, int], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', layout=torch.sparse_coo, indices_dtype: dtype = torch.int64, device: device = device(type='cpu'))[source]

Bases: PairwiseEncoder

3D voxel specialization of PairwiseEncoder (deprecated).

Deprecated since version 0.x: Use PairwiseEncoder which supports arbitrary N-D spatial relationships.

Encodes pairwise 3D spatial neighborhoods and channel relations into sparse COO or CSR tensors for volumes (C, H, D, W). Maintained only for backward compatibility; internally delegates to the generic encoder.

Parameters:

  • radius (float) – Spatial radius (neighborhood extent).

  • volume_shape (Tuple[int, int, int, int]) – (C, H, D, W) with all positive integers.

  • diag (bool, optional) – Include diagonal offsets. Default False.

  • upper (bool or None, optional) – Triangular selection over offsets (see generic encoder). Default None.

  • channel_voxel_relation ({"indep","intra","inter"}, optional) – Channel relation mode. Default "indep".

  • layout (torch.layout, optional) – Output layout (sparse_coo default or sparse_csr).

  • indices_dtype (torch.dtype, optional) – Index dtype (int32 or int64; default int64).

  • device (torch.device, optional) – Storage device (default CPU).

volume_numel

Total sites C * H * D * W.

Type:

int

offsets

Ordered offset tuples.

Type:

list[tuple[int,int,int,int]]

indices / crow_indices / col_indices / csr_permutation

Sparse structure (depending on layout).

Type:

torch.Tensor

Warning

DeprecationWarning

Issued upon instantiation.

Examples

Deprecated usage: >>> from torchsparsegradutils.encoders import PairwiseVoxelEncoder >>> import warnings >>> with warnings.catch_warnings(): … warnings.simplefilter(‘ignore’, DeprecationWarning) … enc = PairwiseVoxelEncoder(radius=1.5, volume_shape=(2, 8, 8, 8)) >>> enc.indices.shape[0] 2

Preferred replacement: >>> from torchsparsegradutils.encoders import PairwiseEncoder >>> new = PairwiseEncoder(radius=1.5, volume_shape=(2, 8, 8, 8), channel_voxel_relation=’indep’) >>> new.indices.shape[0] 2

See also

PairwiseEncoder

General N-D implementation.

__init__(radius: float, volume_shape: Tuple[int, int, int, int], diag: bool = False, upper: bool | None = None, channel_voxel_relation: str = 'indep', layout=torch.sparse_coo, indices_dtype: dtype = torch.int64, device: device = device(type='cpu'))[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.