hyperiax.SweepFn

class hyperiax.SweepFn(direction, fn, reads, reads_children, reads_parent, writes, writes_children=())

Bases: object

A direction-tagged, pure Tree -> Tree transform.

Hashable so the dispatcher can pin it as JIT static. Equality is value-based on (direction, fn identity, read/write spec tuples).

writes_children (up sweeps only) declares fields whose per-parent output has shape (k, *trailing) and gets scattered to the children’s slots rather than the parent’s slot — useful for caching per-edge quantities computed during the up-sweep (e.g. the per-step F_t/H_t from a per-edge ODE filter).

Parameters:

• direction (Literal['up', 'down'])

• fn (Callable)

• reads (tuple[str, ...] | None)

• reads_children (tuple[str, ...] | None)

• reads_parent (tuple[str, ...] | None)

• writes (tuple[str, ...])

• writes_children (tuple[str, ...])

__init__(direction, fn, reads, reads_children, reads_parent, writes, writes_children=())

Parameters:

• direction (Literal['up', 'down'])

• fn (Callable)

• reads (tuple[str, ...] | None)

• reads_children (tuple[str, ...] | None)

• reads_parent (tuple[str, ...] | None)

• writes (tuple[str, ...])

• writes_children (tuple[str, ...])

Return type:

None

Methods

__init__(direction, fn, reads, ...[, ...])

Attributes

writes_children
direction
fn
reads
reads_children
reads_parent
writes

direction: Literal['up', 'down']

fn: Callable

reads: tuple[str, ...] | None

reads_children: tuple[str, ...] | None

reads_parent: tuple[str, ...] | None

writes: tuple[str, ...]

writes_children: tuple[str, ...] = ()