# How the SparseTensor Class Represents 3D Geometry in TRELLIS 2's Latent Space

> Learn how TRELLIS.2's SparseTensor class efficiently represents 3D geometry in latent space using coordinate-feature pairs for dual-backend support. Optimize your sparse data operations.

- Repository: [Microsoft/TRELLIS.2](https://github.com/microsoft/TRELLIS.2)
- Tags: internals
- Published: 2026-08-04

---

**The `SparseTensor` class encodes sparse 3D geometry as coordinate-feature pairs with dual-backend support, enabling efficient latent space operations on only occupied voxels while providing a unified interface for both torchsparse and spconv libraries.**

TRELLIS 2 relies on `SparseTensor` as its core data structure for modeling 3D assets in a compressed latent representation. Building on the generic `VarLenTensor` base class, it stores only active spatial locations alongside their learned feature vectors to power the voxel-based generative pipeline. Understanding how the `SparseTensor` class works for representing 3D geometry in the latent space is essential for extending the architecture or debugging the sparse convolution pipeline.

## Core Data Structure: Features and Coordinates

In [`trellis2/modules/sparse/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/basic.py) (lines 43-51), the `SparseTensor` class maintains two parallel tensors that define the latent spatial representation:

- **`feats`**: The per-voxel latent feature vectors of shape `(N, C)`, where `N` is the number of active voxels and `C` is the feature dimension
- **`coords`**: Integer voxel coordinates of shape `(N, 4)`, where the first column represents the batch index and the remaining three columns store X, Y, and Z spatial positions

Both tensors share the same first dimension, ensuring each active voxel has exactly one coordinate and one feature vector. This design eliminates storage overhead from empty space, which is critical for processing high-resolution 3D volumes efficiently.

## Dual-Backend Abstraction

The class implements lazy backend selection to support both **torchsparse** and **spconv** without requiring code changes. In [`trellis2/modules/sparse/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/basic.py) (lines 66-74), the constructor dynamically imports the chosen library based on the global configuration flag `config.CONV`. This abstraction allows the same model code to run on either backend depending on which library is installed in the environment.

The underlying backend tensor is wrapped while preserving access to library-specific optimizations, ensuring that operations like sparse convolutions execute with native performance regardless of the selected engine.

## Construction Modes

The `__init__` method in [`trellis2/modules/sparse/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/basic.py) (lines 76-88) automatically distinguishes between two construction patterns.

**Tensor-first construction** builds the backend from raw PyTorch tensors:

```python
feats = torch.randn(num_active, latent_dim)
coords = torch.randint(0, 64, (num_active, 4))
sparse = SparseTensor(feats, coords, shape=torch.Size([batch_sz, 64, 64, 64]))

```

**Data-first construction** wraps an existing backend tensor:

```python
sparse = SparseTensor(backend_data, shape)

```

When no explicit shape is provided, the constructor infers the full 3D grid dimensions from the maximum batch index in coordinates and the feature dimensions (lines 61-65).

## Batch Layout and Fast Indexing

Efficient batch processing requires mapping each batch element to its contiguous data block. The `SparseTensor` class computes a `layout` attribute (lines 66-71) containing Python `slice` objects derived from the coordinate batch column using a bincount-plus-cumsum strategy.

This layout enables O(1) batch access:

```python
batch_idx = 0
batch_feats = sparse.feats[sparse.layout[batch_idx]]
batch_coords = sparse.coords[sparse.layout[batch_idx]]

```

The spatial shape—the active X-Y-Z extents—is computed on-demand from `coords[:, 1:]` and cached for reuse (lines 73-78).

## Conversion Utilities for Batch Processing

The class provides static and instance methods for converting between batched list formats and unified sparse tensors. In [`trellis2/modules/sparse/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/basic.py), the `from_tensor_list` method (lines 34-41) concatenates per-batch tensors:

```python
feats_list = [torch.randn(N_i, latent_dim) for N_i in num_vox_per_batch]
coords_list = [torch.randint(0, 64, (N_i, 4)) for N_i in num_vox_per_batch]
sparse = SparseTensor.from_tensor_list(feats_list, coords_list)

```

Conversely, `to_tensor_list` (lines 46-55) splits a `SparseTensor` back into per-batch lists using the precomputed layout slices:

```python
feats_list, coords_list = sparse.to_tensor_list()

```

These utilities are heavily utilized in the image-to-3D and texturing pipelines where individual batch elements require separate processing.

## Device and Dtype Management

The `SparseTensor` class forwards device transfer methods to the underlying backend while preserving the wrapper structure. In [`trellis2/modules/sparse/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/basic.py) (lines 57-63), methods like `.to()`, `.cpu()`, and `.cuda()` ensure the entire sparse representation moves between devices with a single call:

```python
sparse = sparse.cuda()

```

This preserves the lazy-imported backend type and maintains the coordinate-feature synchronization across GPU transfers.

## Debug Validation

When `config.DEBUG` is enabled, the constructor performs validation checks (lines 19-27) that assert:

- The number of features matches the number of coordinates
- Inferred shapes and layouts remain consistent
- Each batch's data occupies a contiguous memory block

These checks help catch malformed tensors early in the pipeline, particularly when integrating custom data loaders or preprocessing steps.

## Summary

- The `SparseTensor` class stores 3D geometry as paired **feature** and **coordinate** tensors in [`trellis2/modules/sparse/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/basic.py), eliminating empty voxel overhead.
- **Dual-backend support** for torchsparse and spconv is implemented via lazy import based on `config.CONV`, enabling seamless library switching.
- **Layout slicing** provides O(1) batch access through precomputed slice objects derived from coordinate batch indices.
- **Conversion utilities** `from_tensor_list` and `to_tensor_list` handle the common pattern of splitting and merging batched sparse data.
- **Device management** methods ensure the entire sparse representation moves between CPU and GPU atomically.

## Frequently Asked Questions

### How does SparseTensor handle different batch sizes within the same tensor?

The class uses a `layout` attribute containing Python `slice` objects that map each batch index to its specific range of rows in the feature and coordinate tensors. This is computed using a bincount and cumsum operation on the batch column of the coordinates (lines 66-71 in [`basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/basic.py)), allowing variable-length batches to coexist in a single sparse tensor while maintaining fast indexed access.

### Can I use SparseTensor with both torchsparse and spconv in the same project?

Yes, but only one backend can be active per `SparseTensor` instance. The class lazily imports the backend specified by `config.CONV` during construction (lines 66-74). While you cannot mix backends within a single tensor, you can instantiate different tensors with different backends by modifying the configuration flag before construction, though this is rarely necessary since both backends implement the same convolutional operations.

### What happens if I don't specify the shape parameter when creating a SparseTensor?

If the shape parameter is omitted, the constructor automatically infers the full tensor shape from the maximum batch index found in the coordinates and the feature dimensions (lines 61-65). The resulting shape is a `torch.Size` representing the complete 3D grid dimensions including the batch dimension, calculated as `[max_batch_idx + 1, *feats.shape[1:]]`.

### How does SparseTensor ensure data consistency when moving to GPU?

The class overrides device transfer methods (`.to()`, `.cpu()`, `.cuda()`) to forward these operations to the underlying backend tensor while maintaining the wrapper structure (lines 57-63). This ensures that both the feature and coordinate tensors move to the target device simultaneously, preserving the synchronization required for sparse convolution operations.