beamax.gb.pallas_kernels
Low-level, single-device Pallas accumulation kernels. Most users should select
these paths through MSGBSolver; call them directly only when working with
already prepared beam parameters.
beamax.gb.pallas_kernels
Experimental Pallas kernels for real Gaussian beam accumulation.
GPU and TPU use tiled kernels; CPU supports tests through interpret mode only.
sum_gaussian_beam_real_pallas(xt: jax.Array, pt: jax.Array, Mt: jax.Array, At: jax.Array, omega0: jax.Array, sensors: jax.Array, domain_size: jax.Array, periodic: jax.Array, *, interpret: bool | None = None, config: PallasConfig | None = None, initial_field: jax.Array | None = None, return_padded: bool = False) -> jax.Array
Evaluate and sum real Gaussian beams with an output-stationary kernel.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
xt
|
(array, shape(B, Nt, d))
|
Beam centres and momenta along the trajectories. |
required |
pt
|
(array, shape(B, Nt, d))
|
Beam centres and momenta along the trajectories. |
required |
Mt
|
complex array, shape (B, Nt, d, d)
|
Beam Hessians along the trajectories. |
required |
At
|
complex array, shape (B, Nt) or (B, Nt, 1)
|
Beam amplitudes along the trajectories. |
required |
omega0
|
(array, shape(B))
|
Beam angular frequencies. |
required |
sensors
|
(array, shape(Ns, d))
|
Evaluation positions. |
required |
domain_size
|
(array, shape(d))
|
Periodic wrapping parameters. |
required |
periodic
|
(array, shape(d))
|
Periodic wrapping parameters. |
required |
config
|
PallasConfig
|
Static kernel layout, launch, wrapping, and batch-accumulation
configuration. Defaults to :class: |
None
|
initial_field
|
(array, shape(Nt, Ns) or (Nt, padded_Ns))
|
Existing field to include. GPU lowering may alias it to the output. |
None
|
return_padded
|
bool
|
Retain detector padding for internal batched scans. |
False
|
interpret
|
bool
|
Run the Pallas HLO interpreter. Defaults to |
None
|
Returns:
| Type | Description |
|---|---|
(array, shape(Nt, Ns))
|
Real field summed over the beam axis. |
Notes
Reduction order differs from vectorized jnp.sum.
sum_gaussian_beam_real_hom_diag_3d_pallas(x0: jax.Array, p0: jax.Array, M0: jax.Array, a0: jax.Array, omega0: jax.Array, mode: jax.Array, c0: jax.Array, ts: jax.Array, sensors: jax.Array, domain_size: jax.Array, periodic: jax.Array, *, config: PallasConfig | None = None, initial_field: jax.Array | None = None, return_padded: bool = False, interpret: bool | None = None) -> jax.Array
Fuse homogeneous 3D propagation and real sensor accumulation.
The primal evaluates solve_hom_diag within each output tile instead of
materializing trajectories. Differentiation uses portable XLA.