beamax.gb.PallasConfig
Static launch, layout, and periodic-wrapping options for experimental Pallas kernels.
Use the defaults unless profiling identifies a reason to tune them:
sensor_block_sizecontrols detector tiling;gpu_time_block_size,gpu_num_warps, andgpu_num_stagescontrol GPU launches;rounding_modeandperiodic_axescontrol minimum-image wrapping; and- the
tpu_*fields select TPU layouts and block sizes.
beamax.gb.pallas_config.PallasConfig(sensor_block_size: int = 128, gpu_num_warps: int | None = 4, gpu_num_stages: int | None = 2, rounding_mode: RoundingMode = 'nearest_even', periodic_axes: tuple[bool, ...] | None = None, tpu_layout: TPULayout = 'beam_sensor', tpu_beam_block_size: int = 8, tpu_time_block_size: int = 8, tpu_grid_order: TPUGridOrder = 'time_sensor', gpu_time_block_size: int = 1)
dataclass
Compilation-static launch and wrapping options.