Functional top-level exports¶

class torch_lattice.nn.functional.ConvMode(*values)[source]¶

Bases: Enum

mode0 = 0¶

mode1 = 1¶

mode2 = 2¶

class torch_lattice.nn.functional.Dataflow(*values)[source]¶

Bases: Enum

ImplicitGEMM = 0¶

GatherScatter = 1¶

FetchOnDemand = 2¶

torch_lattice.nn.functional.avg_pool3d(inputs, **kwargs)[source]¶

Return type:: SparseTensor
Parameters:: inputs (SparseTensor)

torch_lattice.nn.functional.build_kernel_map(_coords, input_node_num, kernel_size=2, stride=2, padding=0, hashmap_keys=None, hashmap_vals=None, spatial_range=None, mode='hashmap', dataflow=Dataflow.ImplicitGEMM, downsample_mode='spconv', training=False, ifsort=False, generative=False, subm=False, split_mask_num=1, split_mask_num_bwd=1, FOD_fusion=True, IGEMM_center_only=False, inference=False)[source]¶

Return type:

Parameters:

_coords (Tensor)
input_node_num (int)
kernel_size (int | Tuple[int, ...])
stride (int | Tuple[int, ...])
padding (int | Tuple[int, ...])
hashmap_keys (Tensor)
hashmap_vals (Tensor)
spatial_range (tuple[int, ...] | None)
training (bool)
ifsort (bool)
generative (bool)
subm (bool)
split_mask_num (int)
split_mask_num_bwd (int)
FOD_fusion (bool)
IGEMM_center_only (bool)
inference (bool)

torch_lattice.nn.functional.build_pool_output_coords(coords, *, kernel_size, stride=1, padding=0, dilation=1, spatial_range=None)[source]¶

Generate convolution-style output support entirely on the input device.

Return type:: Tensor
Parameters:: coords (Tensor)

torch_lattice.nn.functional.build_target_out_in_map(input_coords, target_coords, *, kernel_size, stride=1, padding=0, dilation=1)[source]¶

Build (N_target, kernel_volume) target-to-input row indices.

Return type:

Parameters:

input_coords (Tensor)
target_coords (Tensor)

torch_lattice.nn.functional.build_target_transposed_out_in_map(input_coords, target_coords, *, kernel_size, stride=1, padding=0, dilation=1)[source]¶

Build target-to-input rows for transposed convolution geometry.

Return type:

Parameters:

input_coords (Tensor)
target_coords (Tensor)

torch_lattice.nn.functional.build_transposed_output_coords(coords, *, kernel_size, stride=1, padding=0, dilation=1)[source]¶

Generate unique support for transposed convolution geometry.

Return type:: Tensor
Parameters:: coords (Tensor)

torch_lattice.nn.functional.calc_ti_weights(coords, idx_query, scale=1)[source]¶

Return type:

Parameters:

coords (Tensor)
idx_query (Tensor)
scale (float)

torch_lattice.nn.functional.conv3d(input, weight, kernel_size, bias=None, stride=1, padding=0, dilation=1, config=None, subm=False, transposed=False, generative=False, training=False, coordinates=None)[source]¶

Apply sparse convolution with generated or explicit target support.

Return type:

Parameters:

input (SparseTensor)
weight (Tensor)
kernel_size (int | Sequence[int])
bias (Tensor | None)
stride (int | Sequence[int])
padding (int | Sequence[int])
dilation (int | Sequence[int])
subm (bool)
transposed (bool)
generative (bool)
training (bool)
coordinates (SparseTensor | None)

torch_lattice.nn.functional.convert_transposed_out_in_map(out_in_map, size)[source]¶

Invert an int32 output-to-input relation for transposed convolution.

Return type:

Parameters:

out_in_map (Tensor)
size (int)

torch_lattice.nn.functional.gather_scatter_kmap_from_out_in_map(out_in_map, *, input_size)[source]¶

Convert a target relation to the gather/scatter kernel-map ABI.

Return type:

Parameters:

out_in_map (Tensor)
input_size (int)

torch_lattice.nn.functional.get_conv_mode()[source]¶

Return type:: ConvMode

torch_lattice.nn.functional.global_avg_pool(inputs, *, batch_size=None)[source]¶

Return type:

Parameters:

inputs (SparseTensor)
batch_size (int | None)

torch_lattice.nn.functional.global_max_pool(inputs, *, batch_size=None)[source]¶

Return type:

Parameters:

inputs (SparseTensor)
batch_size (int | None)

torch_lattice.nn.functional.global_pool(inputs, *, mode='sum', batch_size=None)[source]¶

Reduce sparse features independently for every declared batch.

Return type:

Parameters:

inputs (SparseTensor)
mode (Literal['sum', 'avg', 'max'])
batch_size (int | None)

torch_lattice.nn.functional.global_sum_pool(inputs, *, batch_size=None)[source]¶

Return type:

Parameters:

inputs (SparseTensor)
batch_size (int | None)

torch_lattice.nn.functional.leaky_relu(input, negative_slope=0.1, inplace=True)[source]¶

Return type:

Parameters:

input (SparseTensor)
negative_slope (float)
inplace (bool)

torch_lattice.nn.functional.max_pool3d(inputs, **kwargs)[source]¶

Return type:: SparseTensor
Parameters:: inputs (SparseTensor)

torch_lattice.nn.functional.normalized_conv3d(input, weight, kernel_size, bias=None, stride=1, padding=0, dilation=1, config=None, subm=False, transposed=False, generative=False, training=False, coordinates=None, eps=1e-08)[source]¶

Apply weight-normalized sparse convolution.

Non-pointwise kernels compute conv(input, weight) and divide by sqrt(conv(ones, weight.square()) + eps) before applying bias. Both passes use the same coordinate manager and therefore reuse cached kernel relations. Pointwise kernels intentionally use ordinary matrix multiplication, matching the source normalized-convolution contract.

Return type:

Parameters:

input (SparseTensor)
weight (Tensor)
kernel_size (int | Sequence[int])
bias (Tensor | None)
stride (int | Sequence[int])
padding (int | Sequence[int])
dilation (int | Sequence[int])
subm (bool)
transposed (bool)
generative (bool)
training (bool)
coordinates (SparseTensor | None)
eps (float)

torch_lattice.nn.functional.pool3d(inputs, *, mode, kernel_size=2, stride=2, padding=0, dilation=1)[source]¶

Local sparse 3D pooling over convolution-style neighborhoods.

Return type:

Parameters:

inputs (SparseTensor)
mode (Literal['sum', 'max', 'avg'])

torch_lattice.nn.functional.pool_transpose3d(inputs, target=None, *, kernel_size=2, stride=2, padding=0, dilation=1)[source]¶

Average coarse rows onto generated or explicit fine support.

Return type:

Parameters:

inputs (SparseTensor)
target (SparseTensor | None)

torch_lattice.nn.functional.relu(input, inplace=True)[source]¶

Return type:

Parameters:

input (SparseTensor)
inplace (bool)

torch_lattice.nn.functional.set_conv_mode(conv_mode)[source]¶

Return type:: None
Parameters:: conv_mode (int | ConvMode)

torch_lattice.nn.functional.silu(input, inplace=True)[source]¶

Return type:

Parameters:

input (SparseTensor)
inplace (bool)

torch_lattice.nn.functional.spcount(coords, num)[source]¶

Return type:

Parameters:

coords (Tensor)
num (Tensor)

torch_lattice.nn.functional.spcrop(input, coords_min=None, coords_max=None)[source]¶

Return type:

Parameters:

input (SparseTensor)
coords_min (Tuple[int, ...] | None)
coords_max (Tuple[int, ...] | None)

torch_lattice.nn.functional.spdevoxelize(feats, coords, weights)[source]¶

Return type:

Parameters:

feats (Tensor)
coords (Tensor)
weights (Tensor)

torch_lattice.nn.functional.spdownsample(_coords, stride=2, kernel_size=2, padding=0, spatial_range=None, downsample_mode='spconv')[source]¶

Return type:

Parameters:

_coords (Tensor)
stride (int | Tuple[int, ...])
kernel_size (int | Tuple[int, ...])
padding (int | Tuple[int, ...])
spatial_range (Tuple[int, ...] | None)
downsample_mode (str)

torch_lattice.nn.functional.sphash(coords, offsets=None)[source]¶

Hash (batch, x, y, z) int32 coordinate rows.

Return type:

Parameters:

coords (Tensor)
offsets (Tensor | None)

torch_lattice.nn.functional.sphashquery(queries, references)[source]¶

Return the row index of every query hash, or -1 when absent.

Return type:

Parameters:

queries (Tensor)
references (Tensor)

torch_lattice.nn.functional.spupsample_generative(_coords, stride=2, kernel_size=2, padding=0, spatial_range=None)[source]¶

Return type:

Parameters:

_coords (Tensor)
stride (int | Tuple[int, ...])
kernel_size (int | Tuple[int, ...])
padding (int | Tuple[int, ...])
spatial_range (Tuple[int, ...] | None)

torch_lattice.nn.functional.spvoxelize(feats, coords, counts)[source]¶

Return type:

Parameters:

feats (Tensor)
coords (Tensor)
counts (Tensor)

torch_lattice.nn.functional.sum_pool3d(inputs, **kwargs)[source]¶

Return type:: SparseTensor
Parameters:: inputs (SparseTensor)

torch_lattice.nn.functional.transpose_kernel_map(kmap, ifsort=False, training=False, split_mask_num=1, split_mask_num_bwd=1)[source]¶

Return type:

Parameters:

kmap (Dict)
ifsort (bool)
training (bool)
split_mask_num (int)
split_mask_num_bwd (int)

torch_lattice.nn.functional.trilinear_upsample3d(inputs, target=None, *, stride=2)[source]¶

Upsample sparse features with normalized trilinear interpolation.

Return type:

Parameters:

inputs (SparseTensor)
target (SparseTensor | None)