eagle.frameworks.torch.KernelFunction

eagle.frameworks.torch.KernelFunction#

class eagle.frameworks.torch.KernelFunction(kernel, *, wrt=None, cache_dir=None, device_arch=None)[source]#

Bases: object

A hawk kernel, its derived reverse- and forward-mode kernels, and the torch.autograd.Function that applies them (see function()).

Variables:
  • inputs – the input names, in positional order.

  • outputs – the output names, in return order – the Terminated mask’s own name is last when finishes_terminated is true.

  • wrt – the inputs that receive a gradient.

  • finishes_terminated – whether the kernel’s own body finishes the Terminated mask (terminated = cond), rather than only reading it.

finishes_terminated#

a mask it only READS stays input-only (unchanged), but one it FINISHES is also an output, so the caller’s next step reuses the SAME decision instead of recomputing it with a second masking rule of its own.

Type:

The kernel’s own finish, surfaced

__call__(*args, layout=None, **kwargs)[source]#

Run the kernel on args/kwargs and return its outputs as tensors, differentiable through torch, in the caller’s layout (a mix of sample-major and component-major inputs is refused, see _call_layout()). layout ("samples_first" / "samples_last") and eagle.samples_first(x) / eagle.samples_last(x) say which axis holds the samples of an input whose shape reads both ways ((w, w)), which is otherwise refused.