Top-level package

torch_lattice re-exports the sparse tensor container, top-level sparse tensor operators, voxelization helpers, and the tuning entry point. Prefer importing module families explicitly in larger codebases, for example torch_lattice.nn as spnn and torch_lattice.nn.functional as F.

class torch_lattice.SparseTensor(feats, coords, stride=1, spatial_range=None, *, batch_counts=None, coord_manager=None, coord_key=None)[source]

Bases: object

Sparse feature value with explicit coordinate-support identity.

Coordinates have shape (N, 4) and use (batch, x, y, z) order. Features have shape (N, C) and share their row order with coordinates. Coordinate managers own support identity and cached sparse relations; feature-only transformations preserve that identity, while row-changing operations create a new coordinate key.

Parameters:
replace(*, feats)[source]

Return a feature replacement on the same coordinate support.

Return type:

SparseTensor

Parameters:

feats (Tensor)

with_coordinates(*, feats, coords, stride=None, spatial_range=<object object>, batch_counts=None)[source]

Return a value on newly created coordinate support.

Return type:

SparseTensor

Parameters:
cpu()[source]
Return type:

SparseTensor

cuda(device=None)[source]
Return type:

SparseTensor

Parameters:

device (device | int | None)

half()[source]
Return type:

SparseTensor

detach()[source]
Return type:

SparseTensor

to(device, *, non_blocking=False)[source]
Return type:

SparseTensor

Parameters:
dense()[source]
Return type:

Tensor

property batch_indices: Tensor

Batch column from coords.

property batch_rows: tuple[Tensor, ...]

Row indices grouped by coordinate batch value.

property decomposed_coordinates: tuple[Tensor, ...]

Spatial coordinates split by batch.

property decomposed_features: tuple[Tensor, ...]

Feature rows split by batch.

property decomposed_coordinates_and_features: tuple[tuple[Tensor, ...], tuple[Tensor, ...]]

Spatial coordinates and features split by batch.

torch_lattice.cat(inputs, *, join='inner')[source]
Return type:

SparseTensor

Parameters:
torch_lattice.devoxelize(points, voxels, *, batch_indices=None, point_active_rows=None, voxel_size=1.0, origin=0.0, interpolation='nearest')[source]

Sample sparse voxel features at dense point rows.

Return type:

Tensor

Parameters:
torch_lattice.generative_add(a, b)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.prune(input, rows)[source]

Keep sparse rows in caller-supplied order.

Return type:

SparseTensor

Parameters:
torch_lattice.prune_mask(input, mask)[source]

Keep sparse rows selected by a boolean mask.

Return type:

SparseTensor

Parameters:
torch_lattice.reindex_sparse(input, target, *, fill=0.0)[source]

Gather input features onto the exact row order of target.

Return type:

SparseTensor

Parameters:
torch_lattice.sparse_add(lhs, rhs, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.sparse_binary(lhs, rhs, op, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.sparse_cat(inputs, *, join='inner')[source]
Return type:

SparseTensor

Parameters:
torch_lattice.sparse_from_coordinates(coords, feats, *, stride=1, spatial_range=None, batch_counts=None, duplicate_reduction='none')[source]

Construct a sparse tensor with explicit duplicate-row semantics.

Return type:

SparseTensor

Parameters:
torch_lattice.sparse_maximum(lhs, rhs, *, join='inner', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.sparse_minimum(lhs, rhs, *, join='inner', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.sparse_mul(lhs, rhs, *, join='inner', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.sparse_sub(lhs, rhs, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

SparseTensor

Parameters:
torch_lattice.tune(model, data_loader, n_samples=100, collect_fn=<function <lambda>>, enable_fp16=False, save_dir='.torch-lattice-tune', tune_tag='temp', force_retune=False, dataflow_range=None, dataflow_prune=False, tune_with_bwd=False, verbose=True, skip_warning=False)[source]

Tune sparse convolution backend configuration for a model.

Parameters:

  • model (Module) – Module to profile for convolution backend configuration.

  • data_loader (Iterable) – Iterable that yields representative training samples.

  • n_samples (int) – Number of samples used while profiling candidate configs.

  • collect_fn (Callable) – Function that converts one data-loader item into model input. The tuned call is equivalent to model(collect_fn(data)) unless the callable returns a structure consumed by the model itself.

  • enable_fp16 (bool) – Profile with half precision and CUDA autocast enabled.

  • save_dir (str) – Directory used to cache tuned configuration files.

  • tune_tag (str) – Cache file name under save_dir.

  • force_retune (bool) – Ignore an existing cache file and profile again.

  • dataflow_range (List) – Candidate convolution dataflows. When omitted, forward-only tuning checks implicit GEMM and Fetch-on-Demand; backward tuning uses implicit GEMM.

  • dataflow_prune (bool) – Select the best dataflow before tuning lower-level config thresholds.

  • tune_with_bwd (bool) – Include backward timing in the tuning objective.

  • verbose (bool) – Print tuning progress and cache information.

  • skip_warning (bool) – Suppress iterator and backend-mode warnings.

torch_lattice.voxelize(points, features, *, batch_indices=None, active_rows=None, voxel_size=1.0, origin=0.0, reduction='mean', stride=1)[source]

Quantize point rows into a sparse voxel tensor.

Return type:

SparseTensor

Parameters: