deepmd.pt_expt.kernels.triton.sezm.wigner_monomials

Quaternion monomial design matrices with compile-time exponent tables.

The Wigner-D construction for degrees l >= 2 evaluates, per edge, a fixed monomial basis of the unit quaternion

M[e, m] = q0^a_m * q1^b_m * q2^c_m * q3^d_m,

with a_m + b_m + c_m + d_m equal to the kernel degree, followed by one matrix multiply against a precomputed coefficient table. The reference chain (power table, gather, prod) materializes three (4, P + 1, E) intermediates per degree kernel, and its prod backward lowers to a cumprod scan pair – several milliseconds per model call at typical edge counts. Here the exponent table is a compile-time constant: the kernel builds the four scalar power ladders in registers and emits every monomial (and, in the backward, its four leave-one-out derivatives d M_m / d q_i = e_i * q_i^{e_i - 1} * prod_{j != i} q_j^{e_j}) as an unrolled register product. No intermediate ever touches DRAM.

The operator is functional (mutates_args=()) with a fake kernel and an autograd formula whose backward is itself a triton_op, so it composes with the SeZM make_fx lowering and the AOTInductor freeze exactly like the other sezm_triton operators. The exponent table is passed as a Python list[int] and must be extracted from the coefficient buffers in eager context (module construction), never at trace time: a trace-time .tolist() on a tensor creates unbacked symbols and aborts export.

Attributes

WIGNER_MONOMIALS_TRITON_AVAILABLE

Functions

wigner_monomials(→ torch.Tensor)

Evaluate a fixed quaternion monomial basis per edge.

Module Contents

deepmd.pt_expt.kernels.triton.sezm.wigner_monomials.WIGNER_MONOMIALS_TRITON_AVAILABLE = True

deepmd.pt_expt.kernels.triton.sezm.wigner_monomials.wigner_monomials(q: torch.Tensor, exponents: list[int], max_power: int) → torch.Tensor

Evaluate a fixed quaternion monomial basis per edge.

Parameters:

qTensor

Unit quaternions with shape (E, 4).

exponentslist[int]

Flattened exponent table (a0, b0, c0, d0, a1, ...) with 4 * M entries; must be a Python list of compile-time constants (extracted in eager context, never at trace time).

max_powerint

Largest exponent appearing in the table (the power-ladder depth).

Returns:

Tensor

Monomial design matrix with shape (E, M), where column m is q0^a_m * q1^b_m * q2^c_m * q3^d_m.