jax_sn.operators.moment_operator.fluxmoment_strategies module¶

class EinsumFluxMoment[source]¶

Bases: FluxMomentStrategy

Batched-matrix (einsum) implementation. Batches over the trailing dimensions.

contract(data: DiscreteMomentData, x, axis: int)[source]¶: Contract x over axis of data.matrix (0 = angles, 1 = moments).

class StreamingFluxMoment[source]¶

Bases: FluxMomentStrategy

Streaming (scan-over-the-contracted-axis) implementation. Accumulates the contributions one slice at a time, trading compute for a smaller peak memory footprint than the batched einsum.

contract(data: DiscreteMomentData, x, axis: int)[source]¶: Contract x over axis of data.matrix (0 = angles, 1 = moments).