nvalchemiops.torch.segment_ops: Segment Operations

The segment-ops module provides differentiable PyTorch bindings for GPU-accelerated segmented reductions and per-segment algebra.

Tip

For usage guidance and CUDA graph capture patterns, see Segment Operations.

PyTorch autograd bindings for segment operations.

Each public function accepts PyTorch tensors and returns a PyTorch tensor with full first-order and second-order backward support. Integer metadata (idx, num_segments) always receives None gradient.

Every op is wired through register_warp_op_chain(), so the Warp launches are opaque torch.library custom ops: the bindings are torch.compile-clean (single-graph capturable, no graph breaks) and differentiable to second order.

Tensor layout conventions

• Scalar arrays : shape (N,) or (num_segments,)

• Vec3 arrays : shape (N, 3) or (num_segments, 3)

• Mat33 arrays : shape (num_segments, 3, 3)

The dtype (float32 / float64) is inferred from the input tensor.

Public Operations

nvalchemiops.torch.segment_ops.segmented_sum(x, idx, num_segments)

Differentiable segmented sum.

Parameters:

• x (torch.Tensor) – Shape (N,) or (N, 3). dtype float32 or float64.

• idx (torch.Tensor) – Shape (N,), dtype int32. Sorted segment indices in [0, num_segments).

• num_segments (int) – Number of segments.

Returns:

Shape (num_segments,) or (num_segments, 3).

Return type:

torch.Tensor

nvalchemiops.torch.segment_ops.segmented_dot(x, y, idx, num_segments)

Differentiable per-segment dot product.

out[s] = sum_i dot(x[i], y[i])

Parameters:

• x (torch.Tensor) – Shape (N,) or (N, 3). Same dtype and device.

• y (torch.Tensor) – Shape (N,) or (N, 3). Same dtype and device.

• idx (torch.Tensor) – Shape (N,), dtype int32.

• num_segments (int) – Number of segments.

Returns:

Shape (num_segments,) — scalar per segment.

Return type:

torch.Tensor

nvalchemiops.torch.segment_ops.segmented_mul(x, y, idx, num_segments)

Differentiable per-element scale by a per-segment scalar.

out[i] = x[i] * y[idx[i]]

Parameters:

• x (torch.Tensor) – Shape (N,) or (N, 3).

• y (torch.Tensor) – Shape (num_segments,) — one scalar per segment.

• idx (torch.Tensor) – Shape (N,), dtype int32.

• num_segments (int) – Number of segments. Must equal y.shape[0].

Returns:

Same shape as x.

Return type:

torch.Tensor

Raises:

ValueError – If num_segments != y.shape[0]. Without this guard, forward would still succeed (y is only indexed by idx) but backward would allocate grad_y with shape (num_segments,) — a tensor whose shape disagrees with the leaf y, breaking the autograd contract.

nvalchemiops.torch.segment_ops.segmented_mean(x, idx, num_segments)

Differentiable per-segment mean.

out[s] = mean(x[i] for i in segment s)

Parameters:

• x (torch.Tensor) – Shape (N,) or (N, 3).

• idx (torch.Tensor) – Shape (N,), dtype int32. Sorted.

• num_segments (int) – Number of segments.

Returns:

Shape (num_segments,) or (num_segments, 3).

Return type:

torch.Tensor

nvalchemiops.torch.segment_ops.segmented_rms_norm(x, idx, num_segments)

Differentiable per-segment RMS vector norm.

out[s] = sqrt(mean(||x[i]||^2 for i in segment s))

Parameters:

• x (torch.Tensor) – Shape (N, 3). dtype float32 or float64.

• idx (torch.Tensor) – Shape (N,), dtype int32. Sorted.

• num_segments (int) – Number of segments.

Returns:

Shape (num_segments,) — scalar RMS norm per segment.

Return type:

torch.Tensor

nvalchemiops.torch.segment_ops.segmented_matvec(v, m, idx, num_segments)

Differentiable per-segment matrix-vector multiply.

out[i] = m[idx[i]]^T @ v[i]

Parameters:

• v (torch.Tensor) – Shape (N, 3).

• m (torch.Tensor) – Shape (num_segments, 3, 3) — one matrix per segment.

• idx (torch.Tensor) – Shape (N,), dtype int32.

• num_segments (int) – Number of segments. Must equal m.shape[0].

Returns:

Shape (N, 3).

Return type:

torch.Tensor

Raises:

ValueError – If num_segments != m.shape[0]. Without this guard, forward would succeed (m is only indexed by idx) but backward would allocate grad_m with the wrong leading dimension, breaking the autograd contract.