Skip to content

eproperty

eproperty

A descriptor for a value the interleaver serves during a trace.

An eproperty turns a plain attribute — model.h[0].output, tracer.result, a source op's .input — into a handoff point in the run: reading it parks the worker until the model reaches that location and returns the value there; writing it swaps the worker's value in. The location is "{obj.path}.{key}" (or just key when the host has no path, as for tracer.result).

The decorated stub is the preprocess: it takes the raw value the interleaver served and returns what the user reads, so an identity view is just def output(self, value): return value. Two more callbacks refine it:

  • postprocess — runs on a written value before it's swapped in (e.g. repack a lone input back into the (args, kwargs) the model wants).
  • transform — the write-back half of preprocess. When preprocess hands back a reshaped/sliced view, the user's in-place edits to it are invisible to the model (which still holds the original). A transform maps the edited view back to the model's layout, and the result is swapped in when the block next asks the model for anything, or at the end of the block. It receives the edited view and the raw served value: transform(self, view, raw). Until it fires, a second read of the value is the same view, so x.heads[:, 5] += f(x.heads), which reads twice, edits one tensor. It fires whether or not the view was edited, so it must be an identity on an unedited view.

    class Heads(Envoy): @eproperty(key="output") def heads(self, value): # preprocess: [B,S,H] -> heads b, s, h = value.shape return value.view(b, s, self.n_heads, h // self.n_heads).transpose(1, 2)

    @heads.transform
    def heads(self, value, raw):                # write the edited heads back
        b, nh, s, hd = value.shape
        return value.transpose(1, 2).reshape(b, s, nh * hd)
    

    with model.trace(prompt): model.attn.heads[:, 5] = 0 # zero head 5; transform swaps it back

Both callbacks work in the shape of the location, which is not always a bare tensor. key="output" on a tuple-returning module serves — and must be handed back — the whole tuple. key="input" always serves the raw (args, kwargs) pair, the same value .inputs gives you, so the example above written against c_proj's input destructures on the way in and repacks on the way out::

@eproperty(key="input")
def heads(self, value):
    (x,), _ = value
    ...

@heads.transform
def heads(self, value, raw):
    ...
    return ((flat,), {})

raw is the value as served, before preprocess. A transform that hands back a copy of one element of a container — a block that returns (hidden_states, present), read as a clone of hidden_states so a saved read cannot be mutated later — needs it to rebuild the container around the edited copy: the other elements are not in the view, and a transform cannot read the location itself, since it fires on the model side, after the read, where nothing is parked there::

@eproperty(key="output")
def hidden(self, value):
    return value[0].clone()

@hidden.transform
def hidden(self, edited, raw):
    return (edited.clone(), *raw[1:])

One failure mode is worth knowing before you meet it: an eproperty is a property, and a property getter that raises AttributeError falls through to __getattr__. A preprocess that raises one — a typo, a wrong unpack — is swallowed and resurfaces as 'Heads' object (nor its module) has attribute 'heads', blaming the attribute rather than the line inside it.

WriteBack

Bases: NamedTuple

An eproperty view the worker is holding, to be swapped back into the model.

Bound by a read whose preprocess handed the worker a view of the served value (see transform). It is a swap the worker has not issued yet: Mediator.flush issues it the moment the worker moves on, on its next request or at the end of its block, tagged with the occurrence the view was read at so it lands on that visit whatever step the worker has moved to since. Until then a repeated read of the same value is answered with view, so an edit to either is the edit that goes back.

ATTRIBUTE DESCRIPTION
fn

Maps the edited view back into the model's layout; fn().

TYPE: Callable[[], Any]

view

What the worker is holding.

TYPE: Any

eproperty

The descriptor that handed it out; another view of the same location is not this one.

TYPE: Any

location

Where it was read.

TYPE: str

occurrence

Which visit of location it was read at.

TYPE: Optional[int]

fn instance-attribute

fn: Callable[[], Any]

view instance-attribute

view: Any

eproperty instance-attribute

eproperty: Any

location instance-attribute

location: str

occurrence instance-attribute

occurrence: Optional[int]

key property

key: tuple[Any, str, Optional[int]]

What a read has to match to be answered with view.

IEnvoy

Bases: Protocol

Interface for objects that host eproperty descriptors.

An eproperty reads and writes its value through the interleaver at a location derived from the host, so a host must provide:

ATTRIBUTE DESCRIPTION
interleaver

The Interleaver managing execution flow (used by eproperty.provide to serve a value from the model side).

TYPE: Interleaver

path

Optional location prefix used to build the eproperty's location ("{path}.{key}"). May be None / empty — eproperty._location then falls back to the key alone. This is how tracer-level eproperties such as InterleavingTracer.result work without a path prefix.

TYPE: Optional[str]

Notes

Hosts with no meaningful path (e.g. tracers) need not declare path at all — eproperty._location uses getattr(obj, "path", ""), so a missing attribute is treated the same as None / "". It is declared Optional[str] here only for type clarity.

interleaver instance-attribute

interleaver: Interleaver

path instance-attribute

path: Optional[str]

eproperty

eproperty(key: Optional[str] = None, description: Optional[str] = None)

Bases: property

A served value on an interleaving host (Envoy, SourceEnvoy, tracer, ...).

Define one by decorating a stub with @eproperty (or @eproperty(key=...)); the stub is the preprocess. The host only needs a path attribute (and an interleaver for provide); the value is read/written through Mediator, which raises outside a trace.

PARAMETER DESCRIPTION
key

The location suffix appended to the host's path ("{path}.{key}"). Defaults to the stub's name. Several eproperties may share a key to give different views of the same location (input and inputs).

TYPE: Optional[str] DEFAULT: None

description

A short label; only used to surface the attribute in a repr.

TYPE: Optional[str] DEFAULT: None

name instance-attribute

name: Optional[str] = None

key instance-attribute

key: Optional[str] = None

description instance-attribute

description = description

__call__

__call__(preprocess: Callable) -> 'eproperty'

Register the decorated stub as the preprocess and adopt its name/doc.

postprocess

postprocess(func: Callable) -> 'eproperty'

Register the write-side callback, run on a value before it's swapped in.

transform

transform(func: Callable) -> 'eproperty'

Register the write-back for an edited preprocess view (see class doc).

func(self, view, raw): the edited view, and the value as served before preprocess, for a view that is one element of a container the write-back has to rebuild. The result is swapped in when the block next asks the model for anything, or at the end of the block, edited or not; a repeated read before then is answered with the same view.

__str__

__str__() -> str

The line an envoy's repr shows for this value: (name) -> ReturnType: description.

The return type is the preprocess stub's return annotation when it has one, as written (a from __future__ import annotations module gives the alias name, Residual; a class gives its __name__); an unannotated or Any stub shows no arrow.

__get__

__get__(obj: Optional[IEnvoy], owner: Any = None) -> Any

__set__

__set__(obj: IEnvoy, value: Any) -> None

provide

provide(obj: IEnvoy, value: Any) -> Any

Serve this eproperty's value into the run from the model side.

The counterpart to a worker reading it: hands value to the interleaver at this location so a worker parked there is resumed with it. Used for values the model produces outside any module — e.g. a driver feeding tracer.result.