jax_sn.operator_creation.default_strategies module¶

class DefaultMultiGroupStrategy[source]¶

Bases: object

classmethod default_multi_group_strategy(batch_size: int | None = None, checkpoint=True)[source]¶

classmethod default_multi_group_pallas_strategy(batch_size: int | None = None, checkpoint: bool = True)[source]¶

classmethod default_multi_group_warp_strategy(batch_size: int | None = None, checkpoint: bool = True, maximum_number_of_levels: int | None = None)[source]¶: maximum_number_of_levels pins the width the leveled sweep pads its level arrays to; leave it None unless the sweep is rebuilt under a trace (see maximum_number_of_levels).