jax_sn.operators.scatter_operator.gather_scatter_operator_strategies module

Concrete gather-scatter strategies, mirroring the transport pattern (per-angle kernel + concrete vmap lift):

• BroadcastedSingleGroupGather – the per-group kernel (a SingleGroupGatherScatterStrategy), one operation for both directions; sum_over_outgoing only selects which group axis of \(\sigma_{g_{out}g_{in}}\) is contracted.

• VmapGatherScatterStrategy – a concrete ComposedGatherScatterStrategy that lifts the single-group block over the group axis by vmap (the group-axis analogue of VmapAnglesForward).

class BroadcastedSingleGroupGather

Bases: SingleGroupGatherScatterStrategy

Broadcasted (vmap-over-elements) single-group gather kernel.

gather_group(data: GatherScatterData, x: Array, g: int, sum_over_outgoing: bool)

Return type:

Array

class VmapGatherScatterStrategy(single_group_strategy: SingleGroupGatherScatterStrategy)

Bases: ComposedGatherScatterStrategy

Full \(\mathbf{U}\) / \(\mathbf{U}^T\) by vmapping the single-group block over the group axis.

forward(data: GatherScatterData, x: Array)

Return type:

Array

transpose(data: GatherScatterData, x: Array)

Return type:

Array