jax_sn.operators.multi_group_operator.composed_multi_group_strategies module¶

Multi-group strategies

Just composed for now; fused kernel is supported but not implemented.

class MultiGroupChainGatherStrategy(single_group_chain_strategy: SingleGroupChainStrategy, single_group_gather_scatter_strategy: SingleGroupGatherScatterStrategy)[source]¶

Bases: MultiGroupStrategy

Uses a chain and a gather-scatter that fully define the action of a single group, mapping in some way over all the groups.

single_group_chain_strategy: SingleGroupChainStrategy¶

single_group_gather_scatter_strategy: SingleGroupGatherScatterStrategy¶

forward(operator_data: OperatorData, x)[source]¶

transpose(operator_data: OperatorData, x)[source]¶

accepts(non_differentiable_data: NonDifferentiableData)[source]¶

Whether it can accept this non-differentiable data.

Return type:: bool

build_default_non_differentiable_data(geometric_domain, angles)[source]¶

Require each strategy to provide a way to construct this non-differentiable data from the geometric domain and angles.

Return type:: NonDifferentiableData

class MultiGroupChainGatherScanStrategy(single_group_chain_strategy: SingleGroupChainStrategy, single_group_gather_scatter_strategy: SingleGroupGatherScatterStrategy, checkpoint: bool = False)[source]¶

Bases: _ScanLift, MultiGroupChainGatherStrategy

Lift the per-group body by a serialising scan. Optionally checkpoint for reduced memory usage

class MultiGroupChainGatherVmapStrategy(single_group_chain_strategy: SingleGroupChainStrategy, single_group_gather_scatter_strategy: SingleGroupGatherScatterStrategy)[source]¶

Bases: _VmapLift, MultiGroupChainGatherStrategy

Lift by batching the entire vmap. Naturally most parallelizable but significantly increases memory (use on small problems only)

class ComposedMultiGroupChainStrategy(single_group_chain_strategy: SingleGroupChainStrategy)[source]¶

Bases: MultiGroupChainStrategy

Applies a SingleGroupChainStrategy independently over every group. Used to precondition a per-group right-hand side with the same group-lift (scan/vmap, checkpoint) as the operator it accompanies; see from_composed().

single_group_chain_strategy: SingleGroupChainStrategy¶

forward(operator_data: OperatorData, x)[source]¶

transpose(operator_data: OperatorData, x)[source]¶

classmethod from_composed(composed: MultiGroupChainGatherStrategy)[source]¶

The chain alone, lifted the same way (scan/vmap, checkpoint) as composed.

Return type:: ComposedMultiGroupChainStrategy

accepts(non_differentiable_data: NonDifferentiableData)[source]¶

Whether it can accept this non-differentiable data.

Return type:: bool

build_default_non_differentiable_data(geometric_domain, angles)[source]¶

Require each strategy to provide a way to construct this non-differentiable data from the geometric domain and angles.

Return type:: NonDifferentiableData

class MultiGroupChainScanStrategy(single_group_chain_strategy: SingleGroupChainStrategy, checkpoint: bool = False)[source]¶

Bases: _ScanLift, ComposedMultiGroupChainStrategy

Lift the chain by a serialising scan. Optionally checkpoint for reduced memory usage

class MultiGroupChainVmapStrategy(single_group_chain_strategy: SingleGroupChainStrategy)[source]¶

Bases: _VmapLift, ComposedMultiGroupChainStrategy

Lift the chain by batching the entire vmap.