flashcart.o3

Expand directions into irreducible Cartesian tensors, couple tensor features, and mix their channels with equivariant linear maps. A tensor of rank l is stored as 2 * l + 1 independent components. Blocks of components are concatenated in increasing tensor rank, with feature channels contiguous within each component. Tensor-product outputs can also retain a separate coupling-path axis. The layout helpers below describe and convert these representations.

Irreps expands edge directions, TensorProduct combines tensor features, and LinearLayer mixes paths and channels without mixing tensor ranks. Products retain natural-parity couplings: l1 + l2 + l_out must be even in addition to satisfying the triangle rule.

The operations provide PyTorch implementations and optional generated Triton kernels. Select Triton with use_triton=True for CUDA inputs, and configure the generated maximum tensor rank before importing the tensor modules. See Environment variables for the rank setting, Kernel generation and compilation for compilation, and Run the profiler for profiling.

Symmetry and tracelessness determine the full tensor from its independent components. The stored components use the following order:

Tensor rank	Component order
0	Scalar
1	x, y, z
2	xx, yy, xy, xz, yz
3	xxx, xxy, xxz, xyy, xyz, yyz, yyy

For ranks above three, a component is represented by the counts (a, b, c) of its Cartesian indices, with a + b + c == l. Components with c == 0 come first, ordered by decreasing a. Those with c == 1 follow in the same order. Symmetry and tracelessness supply the remaining components.

flashcart.o3.irreps

Expand direction vectors into irreducible Cartesian tensors.

class flashcart.o3.irreps.Irreps(l_max, use_triton=False)

Bases: Module

Expand direction vectors into irreducible Cartesian tensors.

For each unit vector, the expansion contains the independent components of symmetric traceless tensors of ranks zero through l_max. The rank-l tensor contributes 2 * l + 1 components, giving (l_max + 1) ** 2 components in total.

Parameters:

• l_max (int) – Maximum tensor rank in the expansion.

• use_triton (bool, optional) – Use Triton kernels for CUDA inputs when available. Enabling this option requires l_max to fit within the generated rank limit, even when the input is on CPU. Default: False.

Raises:

RuntimeError – use_triton=True and l_max exceeds the generated rank limit configured before import.

forward(x)

Normalize direction vectors and evaluate their Cartesian expansion.

Parameters:

x (torch.Tensor) – Nonzero direction vectors of shape (n, 3). Vectors that remain nonunit after normalization are rejected.

Returns:

Independent tensor components of shape

(n, (l_max + 1) ** 2), concatenated in increasing tensor rank.

Return type:

torch.Tensor

forward_normalized(x)

Evaluate the Cartesian expansion of unit direction vectors.

Parameters:

x (torch.Tensor) – Unit vectors of shape (n, 3). The squared norm of every vector must differ from one by at most 1e-5.

Returns:

Independent tensor components of shape

(n, (l_max + 1) ** 2), concatenated in increasing tensor rank.

Return type:

torch.Tensor

flashcart.o3.tensor_product

Couple Cartesian tensor features with weighted equivariant products.

class flashcart.o3.tensor_product.TensorProduct(in1_l_max, in2_l_max, out_l_max, in1_features, in2_features, symmetric_product=False, use_triton=False, shared_weights=False, reduce_paths=False)

Bases: Module

Evaluate weighted products of irreducible Cartesian tensor features.

Each coupling (l1, l2) -> l_out satisfying the triangle rule and an even value of l1 + l2 + l_out defines one path. These couplings preserve the natural parity (-1) ** l of rank-l tensors. Each path has a separate weight for every output feature channel.

With reduce_paths=False, the output retains the individual paths. Otherwise, their contributions are summed at each output tensor rank. Providing both idx_i and idx_j applies the product as an edge convolution: first-input features are gathered from sender nodes, combined with second-input edge features, and summed at receiver nodes.

Parameters:

• in1_l_max (int) – Maximum tensor rank of the first input.

• in2_l_max (int) – Maximum tensor rank of the second input.

• out_l_max (int) – Maximum tensor rank of the output.

• in1_features (int) – Feature channels of the first input.

• in2_features (int) – Feature channels of the second input. Must equal in1_features unless one of the two is 1 (broadcast).

• symmetric_product (bool, optional) – Require identical inputs and retain only couplings with l1 >= l2. Default: False.

• use_triton (bool, optional) – Use Triton kernels for CUDA inputs when available. All requested tensor ranks must then fit within the configured kernel limit. Default: False.

• shared_weights (bool, optional) – Use one weight vector for every row or edge. If False, supply a separate vector for each row or edge. Default: False.

• reduce_paths (bool, optional) – Sum path outputs per tensor rank instead of stacking them. Default: False.

forward(in1, in2, weights, idx_i=None, idx_j=None)

Evaluate weighted tensor products, optionally summing over edges.

Let features = max(in1_features, in2_features). Input channels are paired elementwise, with broadcasting when one input has a single channel. Weights contain one value per path and feature. The input widths are in1_dim = (in1_l_max + 1)**2 * in1_features and in2_dim = (in2_l_max + 1)**2 * in2_features. Path weights follow the ordering described by flashcart.o3.utils.count_output_paths().

Parameters:

• in1 (torch.Tensor) – First input of shape (n, in1_dim). In convolution mode, the rows correspond to nodes.

• in2 (torch.Tensor) – Second input of shape (n, in2_dim). In convolution mode, its shape is (n_edges, in2_dim).

• weights (torch.Tensor) – Weights of shape (n_total_paths * features,) when shared_weights=True. Otherwise, the shape is (n, n_total_paths * features), or (n_edges, n_total_paths * features) in convolution mode.

• idx_i (torch.Tensor, optional) – Integer receiver-node indices of shape (n_edges,). Provide together with idx_j to enable convolution mode. Default: None.

• idx_j (torch.Tensor, optional) – Integer sender-node indices of shape (n_edges,). Default: None.

Returns:

Features with one row per first-input row, including in

convolution mode. Blocks are concatenated in increasing output tensor rank. Without path reduction, the rank-l block has shape (2 * l + 1, out_paths[l], features) before flattening. With path reduction, it has shape (2 * l + 1, features). Ranks with no allowed paths contribute no output components.

Return type:

torch.Tensor

flashcart.o3.linear

Mix Cartesian tensor channels and coupling paths with equivariant linear maps.

class flashcart.o3.linear.LinearLayer(in_l_max, out_l_max, in_features, out_features, in_paths=None, weight_scale=None, use_triton=False)

Bases: Module

Mix feature channels with an equivariant linear map.

A separate weight matrix at each tensor rank mixes the input channels and coupling paths. The same matrix acts on every tensor component at that rank. Weights are initialized from a standard normal distribution and multiplied by a fixed scale during evaluation.

The map retains ranks up to min(in_l_max, out_l_max). Higher input ranks are discarded, and higher output ranks are filled with zeros.

Parameters:

• in_l_max (int) – Maximum tensor rank of the input.

• out_l_max (int) – Maximum tensor rank of the output.

• in_features (int) – Input feature channels per tensor rank.

• out_features (int) – Output feature channels per tensor rank.

• in_paths (Sequence[int], optional) – Paths per input tensor rank for path-resolved inputs. Default: one per tensor rank.

• weight_scale (float, optional) – Scale applied to every weight matrix. Default: None, which uses 1 / sqrt(in_features * in_paths[l]) at rank l.

• use_triton (bool, optional) – Use Triton kernels for CUDA inputs when available and when the number of active tensor ranks fits within the configured kernel limit. Otherwise, use PyTorch. Default: False.

forward(x)

Mix input channels and paths independently at each tensor rank.

The input width is sum((2*l + 1) * in_features * in_paths[l]) over ranks zero to in_l_max. The output width is (out_l_max + 1)**2 * out_features.

Parameters:

x (torch.Tensor) – Input features of shape (batch, in_dim), with blocks concatenated in increasing tensor rank.

Returns:

Output features of shape (batch, out_dim), with one path

and out_features channels per tensor rank.

Return type:

torch.Tensor

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

flashcart.o3.utils

Describe tensor layouts, convert Cartesian components, and select kernel precision.

flashcart.o3.utils.get_triton_dot_input_precision()

Select Triton dot precision from the PyTorch matrix-product settings.

The result is treated as constant during torch.compile.

Returns:

"tf32" when CUDA is available, TF32 matrix products are enabled, and

the float32 matrix-product precision is "high". Otherwise, return "ieee".

Return type:

str

flashcart.o3.utils.get_linear_triton_kernel_options(dtype)

Select the accumulator mode and dot precision for the linear kernels.

Parameters:

dtype (torch.dtype) – Input dtype.

Returns:

Whether to use float64 accumulation and the Triton dot-input

precision. Float64 inputs return (True, "ieee").

Return type:

tuple[bool, str]

flashcart.o3.utils.apply_path_weight(x, weights, path_idx, out_features, shared)

Scale one path’s output block by its per-feature weights.

Parameters:

• x (torch.Tensor) – Path output of shape (batch, ..., out_features).

• weights (torch.Tensor) – Flat weights over (path, feature). One-dimensional when shared, and of shape (batch, n) otherwise.

• path_idx (int) – Which path’s weight block to apply.

• out_features (int) – Features per path.

• shared (bool) – Weights are batch-shared (1-D).

Returns:

Weighted path features with the same shape as x.

Return type:

torch.Tensor

flashcart.o3.utils.count_output_paths(in1_l_max, in2_l_max, out_l_max, symmetric_product=False)

Count the allowed coupling paths at each output tensor rank.

A path satisfies abs(l1 - l2) <= l_out <= l1 + l2 and an even value of l1 + l2 + l_out. At each output rank, paths with l1 >= l2 are ordered by increasing l1 and then l2. For nonsymmetric products, exchanged unequal-rank paths follow in the same order. Tensor-product weights use this ordering.

Parameters:

• in1_l_max (int) – Maximum tensor rank of the first input.

• in2_l_max (int) – Maximum tensor rank of the second input.

• out_l_max (int) – Maximum tensor rank of the output.

• symmetric_product (bool, optional) – Count only unordered pairs (identical inputs). Default: False.

Returns:

Path counts for output ranks zero through out_l_max,

followed by their sum.

Return type:

tuple[list[int], int]

flashcart.o3.utils.get_irreps_slices(l_max, n_features, n_paths=None)

Locate tensor-rank blocks in the flattened stored-component layout.

Parameters:

• l_max (int) – Maximum tensor rank.

• n_features (int) – Feature channels per tensor rank.

• n_paths (list[int], optional) – Paths per tensor rank. Default: one each.

Returns:

Start and exclusive stop offsets for each tensor rank, in

increasing order.

Return type:

list[tuple[int, int]]

flashcart.o3.utils.get_irreps_shapes(l_max, n_features, n_paths=None, mode='flattened')

Return shape tuples for blocks at each tensor rank.

With mode="factorized", the axes are component, feature, and path. Tensor-product outputs instead store component, path, and feature. The factorized shape is therefore not a direct reshape specification for a tensor-product output when both path and feature counts exceed one.

Parameters:

• l_max (int) – Maximum tensor rank.

• n_features (int) – Feature channels per tensor rank.

• n_paths (list[int], optional) – Paths per tensor rank. Default: one each.

• mode (str, optional) – "flattened" gives (2*l + 1, n_features * n_paths[l]). "factorized" gives (2*l + 1, n_features, n_paths[l]). Default: "flattened".

Returns:

Block shapes in increasing tensor rank, with axes

determined by mode and without a batch dimension.

Return type:

list[tuple[int, …]]

flashcart.o3.utils.get_cartesian_slices(l_max, n_features, n_paths=None)

Locate tensor-rank blocks in the flattened full Cartesian layout.

Parameters:

• l_max (int) – Maximum tensor rank.

• n_features (int) – Feature channels per tensor rank.

• n_paths (list[int], optional) – Paths per tensor rank. Default: one each.

Returns:

Start and exclusive stop offsets for each tensor rank, in

increasing order.

Return type:

list[tuple[int, int]]

flashcart.o3.utils.irreps_to_cartesian(irreps, l_max, n_paths=None)

Reconstruct full Cartesian tensors from their independent components.

Symmetry and tracelessness determine the remaining components at each tensor rank. The stored basis is not orthonormal for ranks two and above, so the ordinary Euclidean norm of the stored components does not generally equal the full Cartesian tensor norm. Paths and features are treated as a combined trailing axis whose ordering is preserved.

Parameters:

• irreps (torch.Tensor) – Features of shape (batch, dim) in the stored layout.

• l_max (int) – Maximum tensor rank.

• n_paths (list[int], optional) – Paths per tensor rank. Default: one each.

Returns:

Full Cartesian components of shape

(batch, sum(3**l * n_features * n_paths[l])), summed over ranks zero through l_max. The feature count is inferred from the input width.

Return type:

torch.Tensor

Raises:

ValueError – The input width is incompatible with the requested tensor ranks and path counts.

flashcart.o3.utils.cartesian_to_irreps(tensors, l_max, n_paths=None)

Select the independent components of full Cartesian tensors.

This operation inverts irreps_to_cartesian for symmetric traceless inputs. It does not project arbitrary inputs onto the symmetric traceless subspace. Paths and features are treated as a combined trailing axis whose ordering is preserved.

Parameters:

• tensors (torch.Tensor) – Symmetric traceless tensor features of shape (batch, dim) in the full Cartesian layout.

• l_max (int) – Maximum tensor rank.

• n_paths (list[int], optional) – Paths per tensor rank. Default: one each.

Returns:

Stored components of shape

(batch, sum((2*l + 1) * n_features * n_paths[l])), summed over ranks zero through l_max. The feature count is inferred from the input width.

Return type:

torch.Tensor

Raises:

ValueError – The input width is incompatible with the requested tensor ranks and path counts.

flashcart.o3.utils.rotate_irreps(irreps, rot_matrix, l_max, n_paths=None)

Apply an orthogonal transformation to stored Cartesian tensor components.

The tensors are reconstructed in full Cartesian form, each tensor index is contracted with the rotation matrix, and the independent components are selected again. For rank-one features, the convention is v_rotated = rot_matrix @ v.

Parameters:

• irreps (torch.Tensor) – Features of shape (batch, dim) in the stored layout.

• rot_matrix (torch.Tensor) – Orthogonal transformation of shape (3, 3), shared by all input rows and using the input dtype and device. Rotations and reflections are supported. Orthogonality is not checked.

• l_max (int) – Maximum tensor rank. Ranks above 13 are not supported.

• n_paths (list[int], optional) – Paths per tensor rank. Default: one each.

Returns:

Rotated features with the same shape and stored-component layout

as irreps.

Return type:

torch.Tensor

Raises:

• ValueError – The input width is incompatible with the requested layout, or rot_matrix does not have shape (3, 3).

• NotImplementedError – l_max exceeds 13.

flashcart.o3.kernel_config

Configure the maximum tensor rank included in generated kernel modules.

Each kernel backend reads FLASHCART_KERNEL_L_MAX when imported and retains the result as KERNEL_L_MAX. The value also contributes to the generated-source fingerprint. Set the environment variable before importing FlashCart. Changing it does not update imported backends.

flashcart.o3.kernel_config.kernel_l_max()

Read the configured maximum tensor rank for generated kernels.

Returns:

Value of FLASHCART_KERNEL_L_MAX, or 3 when it is unset.

Return type:

int