jax_sn.operators.mass_operator.mass_operator_warp module

NVIDIA Warp mass strategy for n_basis <= MAX_NBASIS (=4).

build_kernel(n_basis: int)

Compile the mass kernel for one basis count. Memoised, so distinct counts coexist as separate Warp modules and a repeated call recompiles nothing.

warp_mass_apply(fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, x: Array)

Apply the (symmetric) mass operator via the Warp kernel. Forward == transpose; differentiable.

Return type:

Array

class WarpMass

Bases: MassStrategy

Fused Warp kernel, supporting n_basis <= MAX_NBASIS (=4). Forward == transpose; differentiable.

forward(operator, x)