hawk.kernel

Contents

hawk.kernel#

hawk.kernel(fn=None, *, kind=None, steps=None, time=None, time_step=None, jump=None, group=1)[source]#

Trace fn into a Kernel. Refusals fire here.

kind supplies the kernel’s hawk.ext.Kind (vocabulary, output set, seams); omitted, it gets hawk.ext.DEFAULT_KIND.

steps is how many steps one launch advances a kernel that finishes its own sample (terminated = cond):

@hawk.kernel(steps=16)
def oscillator_step(...): ...

is hawk.steps(<kernel>, 16), named oscillator_step_x16 with state kept in registers across the steps. steps="auto" makes ONE kernel whose launch reads the step count from the reserved fused_steps word (at most hawk.trace.steps.AUTO_K_MAX); an invalid steps refuses at decoration.

Without steps=, a kernel hawk.steps() admits for "auto" is built as the automatic kernel instead, paced by eagle.until_done(). steps=1 is the explicit opt-out: the plain kernel, always.

time_step names the step-size parameter (a declared scalar input) and time the time plane (default "t" once time_step is given). The plain kernel is unchanged; hawk.diff.tape() reads them to promote the step size to a per-sample plane and to write the time plane in index form. A time plane the body updates must update as t = t + dt (checked here, at decoration); one the body does not declare is minted by the twin.

jump makes hawk.diff.vjp()/jvp() differentiate the kernel’s parameter-dependent if/where switches as the smoothed loss: a comparison l < r passes on the jump between its branches, spread over |l - r| < sqrt(5) * jump (Epanechnikov kernel of standard deviation jump). jump= is the width in the units of the compared values (about 2% of their span), or the name of a Param[Scalar] holding it; "auto" picks it per batch from the spread of the compared values. The primal kernel is unchanged.

group is the compaction unit: a lane count (group=32) or the label of a Grid axis (group="sample": that axis and every axis after it). A run that drops finished lanes keeps whole groups together. The default, 1, compacts lane by lane.

Return type:

Any

Parameters:

  • fn (Callable[[...], None] | None)

  • kind (Any)

  • steps (Any)

  • time (str | None)

  • time_step (str | None)

  • jump (Any)

  • group (int | str)

Contents