Plugins

The plugin registry and the plugin-facing types: track and gate generators, the relax, width, elevation and validity stages of the track pipeline and the order and stations stages of the gate pipeline, selected by name, built-ins and installed packages alike (entry point group track_gen.plugins), their Params and the Stage protocol. Writing a Generator and Writing a Stage show how they fit together. track_gen.plugin is a subpackage outside track_gen.__all__: import it explicitly. Importing track_gen loads no generator module and no plugin; the registry loads on its first lookup.

Public extension surface: register track and gate generators, and relax / width / elevation / validity stages of the track pipeline, without editing track_gen.

A pip-installable package adds generators through an entry point in the "track_gen.plugins" group (ENTRY_POINT_GROUP) whose target is a register(registry) callable. In the plugin package’s pyproject.toml:

[project.entry-points."track_gen.plugins"]
acme = "acme_tracks:register"

and in acme_tracks/__init__.py (Superformula is the package’s generator, a Stage subclass; a GeneratorSpec registers the same way):

from acme_tracks.superformula import Superformula

track_gen_api = 1     # optional: the plugin API version this package targets

def register(reg):
    reg.add(Superformula)     # the class, not an instance

after which TrackGenConfig(generator="acme.superformula") selects it. Names: third-party names should be dotted ("package.name"); a non-dotted one registers with a one-time UserWarning. Code that registers at run time (scripts, tests) calls registry’s add() directly; registry.scoped() snapshots and restores the registry around a test (scoped(fresh=True) also re-runs discovery).

Loading. import track_gen loads nothing. The registry fills itself lazily on its first call (get, available, describe, add, unregister, quarantined, scoped(), and the deprecated track_gen._src registry shims), once per process and under a lock:

  1. the built-ins, registered explicitly by track_gen._src.builtin_plugins.register (no entry point needed, so a source checkout without metadata works);

  2. third-party plugins: every entry point of the group, sorted by (distribution name, entry-point name), each loaded and called with the registry.

Policy. The TRACK_GEN_PLUGINS environment variable (POLICY_ENV, read when the registry loads, case-insensitive) selects what happens to a third-party plugin that fails to load: warn (default: it is quarantined and one RuntimeWarning names it), strict (raise PluginLoadError at load, e.g. for CI) or off (skip third-party discovery entirely). Any other value emits a RuntimeWarning and falls back to warn. A plugin fails to load when importing or calling it raises, when it declares a newer plugin API than API_VERSION (track_gen_api = <int> in the module its entry point names, or on the register callable), when it registers a name that already exists without replace=True (built-ins always win unless a plugin explicitly overrides them), or when it unregisters a name it did not register.

Parameters and stages. A generator declares its own parameters as a frozen Params dataclass (fields marked static() or dynamic()) and users select them with TrackGenConfig(generator=..., generator_params=...) (an instance or a mapping of its fields). A legacy GeneratorSpec takes a static Params class through params=; a Stage subclass (kind="generator", output="centerline", registered as a class) allocates through a StageBuilder and may declare dynamic, per-env params that callers write through gen.params["generator"] without re-capturing a graph. Registry.describe() reports each spec’s Params JSON schema (SpecInfo.params_schema).

Pipeline stages. A Stage of kind "relax", "width", "elevation" or "validity" (registered as a class, like a generator Stage) replaces or extends a stage of the track pipeline: select it with track_gen.pipeline.presets.track() (relax="acme.taubin", checks=("default", "acme.max_curvature"), …), or with the flat shortcuts TrackGenConfig(relax_solver=...) / TrackGenConfig(z_profile=...). TrackGen’s own stages of these kinds (relax.xpbd, elevation.noise, validity.default, …) are registered with the built-in generators. track_gen.testing.check_stage() runs the conformance checks on one.

Quarantine. A quarantined plugin’s registrations are rolled back and it is listed in Registry.quarantined; an unknown-name error lists the quarantined plugins with their errors. Asking for a name the plugin tried to register, or a name namespaced under its entry-point name, entry-point module or distribution name (acme.x for an entry point named acme), raises PluginLoadError chained from its load error.

Registry

API_VERSION: int = 1

Version of the plugin API this TrackGen provides. A plugin declares the version it targets as track_gen_api = <int> in the module its entry point names (or as an attribute of its register callable); a plugin that declares a newer version is quarantined instead of being called.

ENTRY_POINT_GROUP: str = "track_gen.plugins"

The importlib.metadata entry-point group scanned for third-party plugins.

POLICY_ENV: str = "TRACK_GEN_PLUGINS"

The environment variable selecting the third-party plugin policy: warn (default), strict or off (Environment variables). Any other value emits a RuntimeWarning and falls back to warn.

registry: track_gen.plugin.Registry

The process-global Registry. TrackGenerator, GateGenerator and Course resolve config.generator through it, and TrackGenerator, GateGenerator and build() the pipeline stages of the open kinds (config.relax_solver, config.z_profile and GateGenConfig.gate_ordering included).

class Registry[source]

Bases: object

Specs keyed by (kind, name). Thread-safe.

Loading, the TRACK_GEN_PLUGINS policy and quarantine are described in track_gen.plugin. The process-global instance is registry. Construct a separate Registry() only for isolated experiments; the library’s facades resolve names through the global one.

__init__()[source]
Return type:

None

add(spec, *, kind=None, replace=False)[source]

Register spec under (kind, spec.name).

Adding the identical spec object again under the same kind is a no-op.

Parameters:

  • spec – The spec (a Stage class, GeneratorSpec, GateGeneratorSpec, or any object with a name under another kind).

  • kind (str | None) – Inferred for a Stage (Stage.kind), GeneratorSpec ("generator") and GateGeneratorSpec ("gate_generator"); required otherwise.

  • replace (bool) – Replace an existing spec of that name, emitting a PluginOverrideWarning that names both origins.

Raises:

  • DuplicateSpecError – The name is already registered in that kind (by a different spec object) and replace is false. Raised from a plugin’s register(reg), it quarantines that plugin (or raises under strict).

  • TypeError – The kind cannot be inferred, or is not a non-empty string, or a relax / width / elevation / validity spec is not a Stage class.

  • ValueError – The kind does not match the spec type, or the name is not a non-empty string.

Return type:

None

Third-party names should be dotted ("package.name"); a new non-dotted name added outside the built-ins emits a one-time UserWarning.

A Stage is registered as its CLASS (kind comes from Stage.kind); its declaration is checked first (see track_gen.plugin.stage). A spec with a params= class has the class checked (frozen dataclass, every field with a default, static fields only).

get(kind, name)[source]

The spec registered under (kind, name).

Raises:

  • PluginLoadError – name belongs to a quarantined plugin (chained from the plugin’s original load error).

  • ValueError – Unknown name; the message lists the available names and the quarantined plugins with their errors.

Parameters:
available(kind=None)[source]

Sorted registered names of kind; with kind=None, every spec as "kind:name".

Parameters:

kind (str | None)

Return type:

list[str]

unregister(kind, name)[source]

Remove (kind, name). Raises ValueError when it is not registered.

Inside a plugin’s register(reg), a plugin may only unregister names it registered itself (override others with add(spec, replace=True)).

Parameters:
Return type:

None

describe(kind, name)[source]

Provenance and execution flags of (kind, name) (see SpecInfo). Raises like get(), and TypeError for a Stage whose capture_class does not return a CaptureClass.

Parameters:
Return type:

SpecInfo

property quarantined: Mapping[str, str]

A read-only mapping of each quarantined plugin to the error text that quarantined it.

Keys are entry-point names; when two quarantined plugins share an entry-point name, each of their keys is qualified as "<distribution>:<entry point>".

scoped(*, fresh=False)[source]

Snapshot every kind (and the load/quarantine state) on enter, restore it on exit; for tests. Lookups from other threads while a scope is entered or exited may observe the change mid-way.

Parameters:

fresh (bool) – Start the scope from an unloaded, empty registry, so the next lookup re-runs built-in and third-party discovery (e.g. under a monkeypatched entry-point list). Without it the registry is loaded on enter.

Return type:

Iterator[Registry]

class SpecInfo[source]

Bases: object

What Registry.describe() reports about one registered spec.

kind

The registry kind ("generator", "gate_generator", …).

Type:

str

name

The registered name.

Type:

str

origin

The module that defines the spec’s callables.

Type:

str

distribution

(name, version) of the distribution whose entry point registered the spec; None for built-ins and for specs added at run time.

Type:

tuple[str, str] | None

capturable

The spec’s capturable flag (True when it has none). For a Stage: capture != HOST, computed from the DEFAULT Params; a stage whose capture_class() depends on its params may differ for other values (the generator reports its own).

Type:

bool

accepts_ctx

The spec’s accepts_ctx flag (False when it has none).

Type:

bool

builtin

True for the specs shipped with track_gen.

Type:

bool

entry_point

The name of the entry point that registered the spec, or None.

Type:

str | None

params_schema

The JSON schema (draft 2020-12) of the spec’s Params class: per field type, default, minimum / maximum (from lo / hi), enum (choices), description (doc), x-static or x-dynamic (with x-dtype) and x-ui. {} for a spec without parameters (a legacy spec without a params= class, or a Stage with NoParams). Always json.dumps-able.

Type:

Mapping[str, Any]

capture

A Stage’s CaptureClass for its default Params; None for a spec that is not a Stage.

Type:

track_gen.plugin.stage.CaptureClass | None

output

A Stage’s output ("centerline" for a generator, None for a track-pipeline stage); None for a spec that is not a Stage.

Type:

str | None

writes

A Stage’s writes (local port -> channel); {} otherwise. A centerline generator’s DECLARED writes for its default Params (Stage.declare; other params may declare more, like capture).

Type:

Mapping[str, str]

provides

A Stage’s provides; empty otherwise (a centerline generator: its declaration at the default Params).

Type:

frozenset

reads

A Stage’s reads (local port -> channel); {} otherwise.

Type:

Mapping[str, str]

orderings

A gate generator’s order whitelist: a Stage’s orderings (None: every registered order but excluded_orderings), a legacy spec’s frozenset(supported_orderings); None for other specs.

Type:

frozenset | None

excluded_orderings

A gate generator Stage’s excluded_orderings.

Type:

frozenset

preferred

A generator Stage’s preferred (a gate generator’s {"order": ..., "normalize": ...}, a centerline generator’s {"relax": ...}); {} otherwise.

Type:

Mapping[str, str]

supported_orderings

A gate generator’s accepted orders, resolved against the registry at describe time and sorted: its whitelist, or every registered order but its excluded_orderings (a legacy spec: its supported_orderings); () for other specs.

Type:

tuple[str, …]

seeded

A generator’s seeded (False: its output does not depend on the seed, e.g. an ingestion generator); True for a legacy spec and every other kind.

Type:

bool

__init__(kind, name, origin, distribution, capturable, accepts_ctx, builtin=False, entry_point=None, params_schema=<factory>, capture=None, output=None, writes=<factory>, provides=frozenset({}), reads=<factory>, orderings=None, excluded_orderings=frozenset({}), preferred=<factory>, supported_orderings=(), seeded=True)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)
class GeneratorSpec[source]

Bases: object

One registered generator.

name: the config.generator string that selects it. Third-party names

should be dotted ("package.name").

alloc_scratch: (config) -> scratch — allocate this generator’s PRIVATE working

buffers ONCE (fixed shapes from config, on config.device). The generation OUTPUT buffers (centerline, valid) are owned by the orchestrator and passed to generate; they are NOT part of this.

generate: (seeds_wp, config, out_centerline, out_valid_wp, scratch) -> None —

write the closed centerline ([E*num_points] vec2f) into out_centerline and per-env validity ([E] int32) into out_valid_wp, using scratch. Pure Warp, in-place, zero-alloc, no host sync WHEN capturable=True.

capturable: whether generate obeys the CUDA-graph-capture contract (pure Warp,

no host sync, no per-call allocation). Defaults True so existing three-arg registrations are unchanged and stay on the capture-then-replay CUDA path. A capturable=False generator may use host-side loop control and per-call allocation; TrackGenerator runs it EAGERLY on CUDA every call (the same code path as cpu), never building a wp.Graph for it.

accepts_ctx: True when generate accepts a ctx= keyword (a

track_gen.exec.LaunchContext) and emits all of its device work through it; the orchestrator then calls generate(..., ctx=ctx). Legacy plugins leave it False and the orchestrator runs generate inside ctx.scope(), so the warp_pipeline helpers they call (which default ctx=None to the ambient context) inherit the call’s stream; their own direct wp.launch calls still go to Warp’s current stream.

params: optional track_gen.plugin.Params subclass (STATIC fields

only) declaring this generator’s own parameters. Users set them with TrackGenConfig(generator=name, generator_params=...) (an instance or a mapping of its fields; None resolves params.from_legacy(config)). generate / alloc_scratch read the resolved instance as config.generator_params, and registry.describe reports its JSON schema. None (the default): the generator takes no generator_params.

A generator written against this function-based form is supported through 0.x; a Stage is the recommended form for a new generator (GeneratorSpec and the 0.1 plugin ABIs).

__init__(name, alloc_scratch, generate, capturable=True, accepts_ctx=False, params=None)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)
class GateGeneratorSpec[source]

Bases: object

One registered native gate generator (the legacy gate generator ABI).

Deprecated: write a Stage with kind="gate_generator" and output="anchors" (Writing a Gate Generator). A GateGeneratorSpec runs through an adapter stage until 1.0, and binding one emits a DeprecationWarning. Its max_gates callable is read only as the adapter’s capacity (the most gates generate writes, checked against GateGenConfig.min_gates / max_gates).

capturable: whether generate obeys the CUDA-graph-capture contract (pure Warp,

no host sync, no per-call allocation). Defaults True. A capturable=False generator may use host-side loop control and per-call allocation; GateGenerator runs it EAGERLY on CUDA every call (the same code path as cpu), never building a wp.Graph. Mirrors GeneratorSpec.capturable.

accepts_ctx: True when generate accepts a ctx= keyword (a

track_gen.exec.LaunchContext) and emits all of its device work through it; the gate Program’s legacy adapter stage (track_gen.stages.gates._LegacyGateSlot) then calls generate(seeds_wp, config, out, scratch, ctx=ctx). Legacy plugins leave it False and the adapter runs generate inside ctx.scope(), so the warp_gate / warp_pipeline helpers they call (which default ctx=None to the ambient context) inherit the call’s stream; their own direct wp.launch calls still go to Warp’s current stream. Mirrors GeneratorSpec.accepts_ctx.

params: optional track_gen.plugin.Params subclass (static fields only)

set through GateGenConfig.generator_params; read as config.generator_params by generate / alloc_scratch / max_gates. Mirrors GeneratorSpec.params.

Deprecated: binding one (GateGenerator, Course, build) emits a DeprecationWarning; it runs through an adapter until 1.0. Write a gate generator Stage instead (The anchors stage contract).

__init__(name, alloc_scratch, generate, max_gates, supported_orderings, capturable=True, accepts_ctx=False, params=None)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)
exception DuplicateSpecError[source]

Bases: ValueError

A spec was added under a (kind, name) that is already registered (and replace=True was not passed). The message names both origins.

exception PluginOverrideWarning[source]

Bases: UserWarning

add(..., replace=True) replaced an existing spec. The message names both origins.

exception PluginLoadError[source]

Bases: RuntimeError

A third-party plugin failed to load: raised at load under TRACK_GEN_PLUGINS=strict, and by Registry.get() when a name provided by a quarantined plugin is requested (chained from the original error).

Parameters

A generator’s own parameters, set with generator_params (a config field or a preset argument) and reported by SpecInfo.params_schema; see Parameters: Params and Per-generator parameters: generator_params. At run time, TrackGenerator.params and GateGenerator.params return a ParamTable (see Runtime parameters and inputs), whose entries hold each dynamic field as a live per-env array.

class Params[source]

Bases: object

Base of every stage’s parameters: a frozen, keyword-only dataclass whose fields all have defaults.

Subclasses must be decorated with @dataclasses.dataclass(frozen=True, kw_only=True); fields use static() or dynamic() (a plain field counts as static). Reserved field names: validate, from_legacy, to_dict, from_dict, replace, schema, fields, describe and any name starting with _.

A static field annotated int, float or bool must hold a value of that type (int: an integral, not a bool; float: any real, not a bool; bool: a bool, or 0 / 1), checked with its lo / hi / choices.

validate(layout)[source]

Validate ONLY this object’s own fields; raise ValueError.

Called at generator construction, and only for the selected generator, before the generic lo / hi / choices checks. The default does nothing.

Parameters:

layout – The run’s Layout (num_envs, stride, device).

Return type:

None

classmethod from_legacy(cfg)[source]

Map a flat TrackGenConfig / GateGenConfig to this class.

Used when cfg.generator_params is None. The default returns cls(). A plugin migrating off a TrackGenConfig subclass overrides it to read its old flat fields (getattr(cfg, "acme_x", default)).

Return type:

Params

__init__()
Return type:

None

class NoParams[source]

Bases: Params

The Params of a stage without parameters.

__init__()
Return type:

None

classmethod __new__(*args, **kwargs)
static(default, *, lo=None, hi=None, choices=None, doc='', ui=None)[source]

A field baked in at construction: shapes, iteration counts, enums, Python branches.

Parameters:

  • default – The field’s default value (required: every Params field has one).

  • lo – Inclusive lower bound (JSON schema minimum), or None.

  • hi – Inclusive upper bound (JSON schema maximum), or None.

  • choices – The allowed values (JSON schema enum), or None.

  • doc (str) – A one-line description (JSON schema description).

  • ui – Free-form UI hint: a widget name such as "slider", or a mapping with widget / label / step / group keys (JSON schema x-ui).

dynamic(default, *, lo=None, hi=None, doc='', ui=None)[source]

A per-env field: an [E] device array the stage reads by pointer (param()).

The annotation selects the dtype: float -> wp.float32, int and bool -> wp.int32. The array is filled with the field’s value at construction and is writable between or inside graph replays with no re-capture (gen.params[stage_id][field], StageParams.set()). Arguments as in static() (no choices).

Parameters:

doc (str)

class ParamTable[source]

Bases: Mapping

stage_id -> StageParams (gen.params): "generator", plus every pipeline stage with a dynamic field (a randomizer, an elevation profile with its altitude knobs, gate_frames, …). Created by the generator facades, never by user code (the constructor is private and may change).

__init__(*, _token=None, stages)[source]
Parameters:

stages (Mapping[str, StageParams])

Return type:

None

set(stage_id, field, value, *, stream=None)[source]

self[stage_id].set(field, value, stream=stream).

Parameters:
Return type:

None

class StageParams[source]

Bases: Mapping

One stage’s dynamic parameters: field -> LIVE [E] array.

The arrays have a stable pointer for the generator’s lifetime; iteration yields the dynamic field names only. params is the resolved Params instance (its static values and the INITIAL dynamic values), or None for a legacy generator without a Params class.

Each generation latches the live arrays into the stage’s working copy for the envs it regenerates (every env, or mask[e] != 0), so a write takes effect at the next generation and envs left out of a masked regeneration keep the value their current output was built with.

Unlatched arrays: a field drawn by a randomize.* stage (drawn) is a single array OWNED by its randomizer and written by it every generation: self[field] returns it (observable), set() raises (set the randomizer’s lo / hi instead), and a direct write is overwritten by the next draw. Bezier’s style_sampling ranges draw their knobs the same way, by the stage itself (drawn["rad"] == "generator", the stage’s own id).

Curriculum bookkeeping: in_effect() returns the [E] values the stage READ for the current rows of each env (the latched working copy, or the drawn array), so a masked regeneration’s unmasked envs report what their rows were built with; bounds() the (lo, hi) a built-in kernel clamps a value to where it reads it (the value USED is clamp(in_effect, *bounds)).

Created by the generator facades, never by user code: the constructor is private (it raises TypeError) and may change; the read API is [field], set(), in_effect(), bounds(), params, stage_id and drawn.

__init__(*, _token=None, stage_id, params, arrays=None, num_envs=0, device=None, default_stream=None, drawn=None, bounds=None, effective=None, hints=None)[source]
Parameters:
Return type:

None

property drawn: Mapping[str, str]

a randomize.* stage targeting this stage, or stage_id itself for a field the stage draws itself (bezier’s style_sampling ranges). Every value is a key of gen.params; empty when nothing draws this stage’s fields.

Type:

field -> id of the stage that draws it every generation

property stage_id: str

The stage id ("generator" for the generator stage).

property params: Params | None

a legacy generator without one): its static values and the INITIAL values of its dynamic fields. A later set() (or a write to the live array) does not change it; the current values are the live arrays, self[field].

Type:

The resolved Params instance as of construction (None

in_effect(field, *, out=None, stream=None)[source]

The [E] values this stage READ for the current rows of each env: for a latched field the stage’s working copy (the latch writes it at the start of every generation, for every env or only the masked ones, so an env left out of a masked regeneration keeps the value its rows were built with); for a drawn field (a randomize.* target, or a knob bezier’s style_sampling draws) the drawn array itself. Same dtype as self[field].

Values read, not values used: an array written out of range is returned as written; a built-in stage clamps it to bounds() where it reads it, so the value it USED is clamp(in_effect(field), *bounds(field)) (equal to in_effect for every in-range write, scalar set() and randomizer draw). Exceptions: width.from_channel reads its scale / lo / hi as written (it clamps its output), a randomizer its lo / hi (it clamps its draw), bezier’s style_sampling draws are not bounded, and Course gates mode’s floor gate_frames.gate_width > 0 is checked by set() on a scalar only (an array value is clamped to the declared 0).

Timing: the values are written on the generation’s stream. A latched field is current once the generation’s first segment ran (the latch is its prologue); a drawn field once the segment holding its randomizer ran (before that, the previous generation’s draws). Read it on that stream (a later launch or capture node, or out= with stream=) or after waiting for the generation; between a set() and the next generation it holds the old values; before any generation, the initial value.

Parameters:

  • field (str) – A dynamic field of this stage.

  • out (array | None) – Optional [E] contiguous array of the field’s dtype on the generator’s device: one copy into it, enqueued on the resolved stream (as set(); capturable), and out is returned. Without it the INTERNAL array is returned (zero-copy, stable pointer): it aliases the stage’s state, so never write it (a write changes what unmasked envs report and what a randomizer’s unmasked envs redraw).

  • stream – Optional wp.Stream (or "torch") for the out= copy.

Raises:

  • KeyError – field is static or unknown.

  • ValueError – out of the wrong shape, dtype, device or layout.

Return type:

array

bounds(field)[source]

(lo, hi) a value of the dynamic field is held to: its declared bounds, narrowed per instance (bezier / hull’s min_num_points: (2, max_num_points); a randomizer’s lo / hi: the drawn field’s declared range); None for an open side. set() refuses a scalar outside them, and the built-in generator, elevation, normalize.extent and gate_frames stages clamp an array value to them where they read it (comparisons: a NaN passes; the exceptions are listed in in_effect()).

Raises:

KeyError – field is static or unknown.

Parameters:

field (str)

Return type:

tuple

set(field, value, *, stream=None)[source]

Write a dynamic field’s LIVE [E] array, in place.

An ENQUEUE on the resolved stream (stream= > ambient scope > the generator’s constructor stream > torch capture > Warp’s current stream): one fill launch for a scalar (checked on the host against lo/hi first), one copy for an array (no range check, no host sync). Inside a capture it is recorded: a scalar is baked into the caller’s graph, an array is read at replay. Takes effect at the next generation. Never allocates.

Parameters:

  • field (str) – A dynamic field of this stage.

  • value – A host scalar (broadcast to every env) or an [E] contiguous wp.array of the field’s dtype on the generator’s device.

  • stream – Optional wp.Stream (or "torch") to enqueue on.

Raises:

  • KeyError – field is static or unknown.

  • ValueError – wrong shape, dtype or device; a value that is neither a wp.array nor a scalar of the field’s type; a scalar outside lo/hi (compared in float32 for a float32 field; for a randomizer’s lo / hi: outside the declared range of the field it draws; for bezier / hull’s min_num_points: above max_num_points; see bounds()); a field drawn every generation (by a randomizer, or by the stage itself: bezier’s style_sampling ranges).

Return type:

None

An array is not range-checked: a built-in stage clamps each value to the field’s declared bounds where it reads it.

The built-in generators’ and elevation stages’ Params have dynamic fields too (the generators’ shape knobs and size, bezier / hull’s int32 min_num_points, z_min / z_max / z_max_step / z_noise_amplitude, the gate pipelines’ normalize.size and gate_frames.gate_width), with declared bounds: see Built-in generator parameters and Per-env shape knobs. in_effect and bounds are the curriculum bookkeeping of the tuning guide. Reach a built-in generator’s class with registry.get("generator", name).Params; .params still works for bezier, hull, polar, voronoi and checkpoint but emits a DeprecationWarning and is removed in 1.0.

Stages

The Stage protocol: a Stage subclass, registered as a class, that allocates through the StageBuilder it is handed and returns an enqueue(ctx) callable. Plugins register kind="generator" stages with output="centerline" (see The Stage contract; an ingestion generator reads per="raw" inputs and may declare its ports per params with declare, see Writing an ingestion generator), kind="gate_generator" stages with output="anchors" (see The anchors stage contract; they also declare provides, orderings or excluded_orderings, and preferred), and stages of the open pipeline kinds "relax", "width", "elevation", "validity", "order" and "stations", which read channels through StageBuilder.read and rewrite them into StageBuilder.write (see Writing a Stage). The other stages are TrackGen’s built-ins (see Pipeline stages). A stage’s CaptureClass is STATIC, BOUNDED (a loop through LaunchContext.bounded_loop, see Writing an iterative (BOUNDED) stage) or HOST (host work through LaunchContext.host() and StageBuilder.host_buffer, see Writing a host-driven (HOST) stage).

class Stage[source]

Bases: ABC

Base of a pipeline stage.

Class attributes a stage sets:

  • kind: the registry kind: "generator", "gate_generator", or a pipeline kind "relax", "width", "elevation", "validity", "order" or "stations".

  • name: the registered name; third parties use a dotted "package.name".

  • Params: the stage’s Params class (default NoParams).

  • capture: its CaptureClass (default STATIC; BOUNDED for a stage whose ctx.bounded_loop passes active=, HOST for one that uses ctx.host()); override capture_class() when it depends on the params.

  • reads: local port -> channel the stage reads (StageBuilder.read()). Must be empty for a generator.

  • writes: local port -> channel the stage writes (StageBuilder.write()); a generator writes "raw.xy" (required) and optionally "gen.valid" and its raw attributes ("raw.z", "raw.half_width", "raw.<package>.<name>"); a gate generator "anchors.xy" and "anchors.count" (required), the attributes it provides and its "anchors.<package>.<name>" attributes. A channel in both reads and writes, through ONE port, is a rewrite: write returns a NEW version of the channel (later stages see it), read the previous one, so the stage writes its whole output into the write buffer and never writes the read one.

  • provides: a gate generator’s per-anchor attributes, a subset of {"z", "tangent", "half_size", "roll"} equal to the anchors.<attr> channels it writes (the gate preset then drops the stages that would compute them); a centerline generator’s, a subset of {"z", "half_width"} equal to the raw.z / raw.half_width it writes (the presets drop the elevation / the width stage); empty for every other kind.

  • orderings / excluded_orderings (gate generators): the order stages it supports: a whitelist (orderings, a frozenset of order names), or, with orderings=None, every registered order but excluded_orderings. A generator that provides attributes declares orderings=frozenset({"raw"}) (they are carried to the gate slots by index).

  • preferred: a gate generator’s {"order": name, "normalize": "extent"} (the order the gate preset picks by default, and whether it normalizes the gates’ extent after ordering); a centerline generator’s {"relax": name} (the relax stage presets.track(relax="auto") picks, "none" for an ingestion generator whose geometry must not move).

  • seeded: False for a generator whose output does not depend on the seed (an ingestion generator): the conformance kit reports diversity as skipped; default True.

  • topology: "closed" or "any" (the topologies the stage supports; BatchSpec.closed=False is refused in 0.x, naming the closed-only stages).

  • output: "centerline" for a generator, "anchors" for a gate generator; None for the other kinds.

  • api_version: the plugin API version the stage targets (<= API_VERSION).

The orchestrator instantiates the class once per pipeline with no arguments.

capture_class(), capacity() and build() receive the params with their STATIC fields only: reading a dynamic field there raises AttributeError (a host scalar taken from it would be baked into the launches and ignore every set); read it on the device through StageBuilder.param().

Besides kind, name, Params, reads, writes, provides, capture, output and topology, a stage may set seeded (True by default; False for a generator whose output does not depend on the seed, see Writing an ingestion generator) and preferred (a gate generator’s "order" / "normalize", a centerline generator’s "relax").

Every built-in generator is a Stage class: registry.get("generator", "bezier") (hull, polar, voronoi, checkpoint: STATIC; repulsive: BOUNDED; external) and registry.get("gate_generator", name). Their module functions (generate_*_warp, *_alloc_scratch, the sample_* helpers and their *_per_env forms) are not public API; see Legacy _src helpers.

kind: ClassVar[str]
name: ClassVar[str]
Params

alias of NoParams

capture: ClassVar[CaptureClass] = 0
reads: ClassVar[Mapping[str, str]] = mappingproxy({})
writes: ClassVar[Mapping[str, str]] = mappingproxy({})
provides: ClassVar[frozenset] = frozenset({})
orderings: ClassVar[frozenset[str] | None] = None
excluded_orderings: ClassVar[frozenset] = frozenset({})
preferred: ClassVar[Mapping[str, str]] = mappingproxy({})
topology: ClassVar[str] = 'closed'
output: ClassVar[str | None] = None
api_version: ClassVar[int] = 1
seeded: ClassVar[bool] = True
classmethod capture_class(params)[source]

The stage’s capture class for params (default: capture).

Return type:

CaptureClass

classmethod declare(params)[source]

(writes, provides) for these params (default: the class’s writes and provides).

Read for centerline generators only in 0.x: a generator whose raw attributes depend on its params (an ingestion generator writes raw.z only when it is given altitudes) overrides it; the declaration is checked at bind time with the class rules (raw.xy mapped, provides naming exactly the catalogue attributes written) and routes the presets (dropped stages, carried attributes) and the planner. params holds the static fields only. Overriding it on another kind is refused at registration.

Return type:

tuple[Mapping[str, str], frozenset]

capacity(params, layout)[source]

Buffer capacities this stage needs. A generator may return {"raw": n} (points per env of its "raw.xy" output, an integer n >= 3); the default {} uses TrackGenConfig.num_points. A gate generator (and a stations stage) returns {"gates": A, "min_gates": k}: A anchor rows per env (the most real anchors it emits, >= 1; it must fit max_gates) and k the fewest (default 0, at most A). Any other key raises ValueError.

Parameters:

layout (Layout)

Return type:

Mapping[str, int]

abstractmethod build(b, params)[source]

Allocate everything through b, ONCE, and return enqueue(ctx).

enqueue emits only ctx.launch / ctx.copy / ctx.fill / ctx.bounded_loop / track_gen.ops work: no allocation, no host sync, no branch on device data or on ctx.capturing (a HOST stage reaches the host through ctx.host()).

Parameters:

b (StageBuilder)

Return type:

Callable[[LaunchContext], None]

class StageBuilder[source]

Bases: object

Handed to Stage.build(); created by the orchestrator, never by plugins.

This release implements read(), write(), temp_like(), scratch(), host_buffer(), param(), input(), salt(), fail_bit(), stride_of(), per_of() and dtype_of(). Every method allocates at build time only, and raises once build() returned.

E

The number of envs.

layout

The run’s Layout.

device

The wp.Device.

stage_id

The stage’s id in the pipeline ("generator" for the generator).

scale

TrackGenConfig.scale (a core field, not a generator Param).

seed

[E] int32: the effective per-env seed of the current generation (seed and epoch folded); read it in enqueue.

raw

Points per env of the "raw.xy" channel.

The arrays param() and input() return are READ-ONLY for the stage: never write them (not even as scratch), since a generation’s values come from the live arrays the caller writes.

__init__(*, _token=None, stage, params, layout, scale, seed, raw, outputs, dynamic, reserved_inputs=('regen_mask',), stage_id='generator', plan=None, capture=None)[source]
Parameters:
Return type:

None

write(port, dtype=None, *, per=None, interp=None, export=False)[source]

The output buffer of writes[port].

A generator writes "raw.xy" (wp.vec2f, per="raw", [E * raw]) and "gen.valid" (wp.int32, per="env", [E]), plus its raw attributes: "raw.z" / "raw.half_width" (float32, the catalogue attributes it provides) and its own "raw.<package>.<name>" channels, each carried to the centerline (cl.<package>.<name>) by the resamples with the interp policy its write declares. per sizes the buffer: "env" is [E], "cl" is [E * N_max], "raw" is [E * raw], "anchors" is [E * A] and "gates" is [E * max_gates] (gate pipelines). dtype / per are those of the channel the port DECLARES; a port the planner rewired (an elevation’s cl.z on gates) gets a buffer sized for the channel it is wired to (per_of(), stride_of()). The buffer is owned by the orchestrator and read directly by later stages (no copy). For a channel the stage also reads (a rewrite) it is a NEW version: read() keeps returning the previous one. A plugin channel (<domain>.<package>.<name>) takes the dtype / per of its first writer; a later writer must match. The prefix fixes per for raw.* ("raw") and anchors.* ("anchors"); cl.* / gates.* default to "env".

dtype and per may be omitted for a catalogue channel, or a plugin channel an earlier stage wrote: they are then the channel’s (b.write("z")). Given, they must match it. A plugin channel’s first writer passes dtype.

Parameters:

  • port (str) – A write port of the stage.

  • dtype – The Warp dtype (see above).

  • per (str | None) – The size (see above).

  • interp (str | None) – A raw attribute’s carry policy (required on every raw.<package>.<name>; "linear", "nearest", "previous" or "angle"): how the resamples carry it to their output points. Only on a generator’s raw writes.

  • export (bool | str) – True (or the channel’s key, e.g. "acme.surface_id") to export a per-point plugin attribute as Track.attributes[key] / GateSequence.attributes[key] (key: the channel name without its domain prefix). On per-point plugin channels only (raw.*, anchors.*, or cl.* / gates.* written with per="cl" / "gates"), by the channel’s first writer.

Raises:

  • ValueError – an undeclared port (lists the declared ones), a channel the stage may not write, a second write of the same port, interp= on a channel that is not a generator’s raw attribute, export= on a catalogue or per-env channel or naming another key, interp= / export= on a rewrite.

  • TypeError – a dtype / per that does not match the channel; a raw attribute without interp=, of a dtype other than float32, int32, vec2f or vec3f, or with a policy its dtype does not admit (int32: "nearest" / "previous"; "angle": float32).

Return type:

array

dtype_of(port)[source]

The Warp dtype of the channel port is WIRED to (a required read port "cl.*": the channel StageRef.ports names).

Raises:

ValueError – an undeclared port, or a plugin channel nobody wrote yet.

Parameters:

port (str)

per_of(port)[source]

The per ("env", "raw", "cl", "anchors" or "gates") of the channel port is WIRED to (after any rewiring: an elevation’s z port is "gates" in a gate pipeline).

Raises:

ValueError – an undeclared port, or a plugin channel nobody wrote yet.

Parameters:

port (str)

Return type:

str

stride_of(port)[source]

The row stride (points per env) of the channel port is WIRED to: raw for "raw", N_max for "cl", A (the anchors producer’s capacity()["gates"]) for "anchors", max_gates for "gates", 1 for "env". Index env e’s row as e * stride_of(port) + i.

Raises:

ValueError – an undeclared port, or a plugin channel nobody wrote yet.

Parameters:

port (str)

Return type:

int

read(port)[source]

The buffer of reads[port]: the latest version of its channel written by an EARLIER stage (for a channel this stage also writes, the version before its own). READ-ONLY for the stage.

Raises:

ValueError – an undeclared read port (lists the declared ones). A generator declares none.

Parameters:

port (str)

Return type:

array

temp_like(port)[source]

A private zero-initialised buffer with the dtype and length of port’s channel (e.g. the ping-pong partner of an in-place rewrite), allocated as scratch() "<port>.tmp<k>" (k = 0, 1, ... per port).

Raises:

ValueError – port is neither a read nor a write port of the stage.

Parameters:

port (str)

Return type:

array

scratch(name, dtype, shape)[source]

A private zero-initialised buffer, namespaced "scratch:<stage_id>/<name>". Never re-zeroed: write before you read. A duplicate name raises ValueError.

Parameters:

name (str)

Return type:

array

host_buffer(dtype, shape)[source]

A zero-initialised HOST array (pinned on a CUDA device) for a HOST stage’s HostOps (via= of readback / upload).

Allocated once, here; the stage keeps it in its enqueue closure.

Raises:

  • ValueError – the stage’s capture class is not CaptureClass.HOST.

  • RuntimeError – called after build() returned.

Return type:

array

param(field)[source]

The stage-side [E] array of a DYNAMIC Params field (float32 or int32).

This is the working copy, latched from the live array the user writes (gen.params[stage_id][field]) at the start of every generation, for the envs that generation regenerates. A field a randomize.* stage draws is the drawn array itself (unlatched: written by the randomizer earlier in the same generation). READ-ONLY for the stage: never write it.

Raises:

KeyError – field is static (“read it from params”) or unknown (lists the dynamic fields).

Parameters:

field (str)

Return type:

array

input(name, dtype, default=0, *, per='env')[source]

A persistent device input exposed as gen.inputs[name] (and Program.input_info[name]), filled with default.

  • per="env" (default): a stable LIVE [E] float32 / int32 array, LATCHED into the returned working copy at every generation like a dynamic param (a write takes effect for the envs the next generation regenerates).

  • per="raw" (centerline generators only): a stable [E * raw] float32 / int32 / vec2f / vec3f array, UNLATCHED: the returned array IS the input, read during the generation (the caller writes it on the generation’s stream, or orders the streams with events). The ingestion inputs (centerline.*) are of this kind.

The returned array is READ-ONLY for the stage, like param()’s. Name the input with a dotted path in the stage’s own namespace (centerline.* is the ingestion namespace track_gen.pipeline.load_centerlines() reads).

Parameters:

  • name (str) – The input’s name ("regen_mask" and existing names clash).

  • dtype – wp.float32 or wp.int32 (per="raw": also wp.vec2f / wp.vec3f).

  • default – The initial value of every element (a scalar; a vector input broadcasts it to every component).

  • per (str) – "env" or "raw".

Raises:

  • TypeError – another dtype.

  • ValueError – a name clash; an unknown per; per="raw" on a stage that is not a centerline generator.

Return type:

array

salt(stream)[source]

A uint32 salt for an independent random stream: hash32_host(fnv1a_32(f"{stage_id}/{stream}".encode())).

Frozen for plugin API version 1, so outputs stay reproducible across releases. Key the per-env RNG as wp.rand_init(int32(hash32(uint32(seed[e]) ^ salt) & 0x7FFFFFFF)). Idempotent per stream; two streams whose salts collide raise ValueError.

Parameters:

stream (str)

Return type:

int

fail_bit(name=None)[source]

The index (16..30) of a new fail bit owned by this validity stage.

In the kernel: fail_out[e] = fail_in[e] | (1 << bit) where the check fails (fail_in / fail_out: the read / write buffers of the stage’s valid.fail port). Bits are handed out in pipeline order; call it once per independent check.

Parameters:

name (str | None) – The bit’s key in Program.fail_bits / Track.fail_bits (default: the stage id). Unique in the pipeline.

Raises:
Return type:

int

class CaptureClass[source]

Bases: IntEnum

How a stage’s enqueue may be captured (ordered: STATIC < BOUNDED < HOST; a pipeline’s class is the largest of its stages’).

  • STATIC: a fixed launch sequence, capturable by anything.

  • BOUNDED: data-dependent convergence with a static bound, through LaunchContext.bounded_loop with active=: capturable by anything; its loop may stop early (an eager host poll on CPU, conditional graph nodes inside Warp-registered captures), so every iteration after active[0] reads 0 must be an exact no-op (an absorbing freeze; the conformance kit checks it).

  • HOST: host readback / control flow (through LaunchContext.host): runs eagerly on every call, between the graph launches of the rest of its pipeline (segmented execution); a caller capture raises NotCapturableError before any launch. It reaches the host only through ctx.host() (HostOps, over StageBuilder.host_buffer buffers); a direct .numpy(), wp.synchronize*() or per-call allocation still fails the conformance kit’s hot_path.host_sync / hot_path.allocation.

__new__(value)
__init__(*args, **kwds)
class track_gen.plugin.Layout

The run’s shape handed to Params.validate, Stage.capacity and the StageBuilder: the same class as track_gen.results.Layout (see Layout).