Skip to content

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_size controls detector tiling;
  • gpu_time_block_size, gpu_num_warps, and gpu_num_stages control GPU launches;
  • rounding_mode and periodic_axes control 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.