Migrating from 1.x to 2.0

unitrack 2.0 is a hard break. There are no compatibility shims, and a tracker written against 1.x will not import unchanged: the primitives were re-cut around a typed data model and a composable stage tree. This page maps the 1.x surface onto 2.0 and then migrates a real 1.x tracker — an appearance-embedding matcher — end to end.

Why 2.0 is a hard break

Four capabilities the companion paper needs forced the redesign: parallel fusion of cost terms, Kalman motion state, first-class gates, and tracklet lifecycle. Each is a cross-cutting concern — it touches the data model and the way stages compose — not a leaf that bolts onto the side. A cross-cutting concern cannot be added to an abstraction that did not anticipate it; you either re-cut the abstraction or encode the change as mode flags, and 1.x's untyped field dicts and flat stage list left only the second route. A compatibility shim would therefore have frozen the very primitives that had to change. 2.0 re-cuts the core instead (decisions Q1–Q2). The five shifts below are what a migration touches.

What changed, conceptually

Five shifts account for almost every edit a migration requires.

1. Field extraction → a typed Detections record. In 1.x, MultiStageTracker(fields=[...]) took a list of TensorDictModules that selected and renamed tensors out of a raw input dict at the front of the pipeline. 2.0 has no field layer. You construct a Detections tensorclass whose field names match the state schema and pass it to step directly.

   Why: the field layer only renamed tensors inside an untyped dict, so a schema error surfaced — if at all — deep inside a stage at runtime. A tensorclass makes the schema a type: an ill-formed record fails to construct, and the now-fixed shapes are what let torch.compile and vmap specialize. Checking moves from runtime to construction, and the renaming layer leaves with the dict it served.

2. Flat stage list → composable stage tree. 1.x nested stages as plain lists (stages=[Gate(then=[Association(...)])]). 2.0 expresses the same structure as a typed tree of combinators — Pipe, Gated, Sequential, Parallel — validated at construction time.

   Why: parallel fusion and cascaded matching are compositions, and a flat list cannot compose — it can only enumerate, so each new combination becomes another global mode flag. A typed tree of combinators is closed under composition: any well-typed subtree is itself a stage, and the constructor rejects the ill-typed ones and names the offending path (PipelineTypeError). Errors a flat list would defer to runtime become unconstructable (Q4).

3. Gates are first-class, separate from costs. 1.x folded gating into the cost with GateCost(field).wrap(cost). 2.0 makes gates GateProducers (ClassGate, ScoreGate, SpatialGate2D, MotionGate, …), combines them with Sequential, and applies them with Gated.

   Why: binding the gate to the cost couples two independent concerns and forces a full (N, M) mask even for a per-detection score floor. Split out, a gate is an algebraic type — per-pair, per-side, or cost-bias — closed under conjunction, so gates compose with Sequential, attach to any cost through Gated, and let the executor allocate the pairwise mask only when a per-pair gate demands it (Q5).

4. States moved from the memory to the tracker, and Value became a triple. 1.x declared states on TrackletMemory(states={...}) as states.Value(dtype, shape=...). 2.0 declares them on Tracker(states={...}) as a State(schema, process, observation, init). A plain feature cache — what 1.x Value was — is Identity process + Replace observation + FromDetectionField initializer.

   Why: a dynamics-free Value can neither predict nor update, which Kalman motion requires. Factoring a state into Process × Observation × Initializer spans the range from a plain cache to a Kalman filter by composition rather than subclassing — the same three slots, filled differently. The states live on the Tracker because it assembles them into one fixed Tracklets type at construction, which gives torch.compile a stable shape and lets the constructor reject type errors before the first allocation (Q3, Q7).

5. SimpleTracker / StatefulTracker → MultiStream; step is pure. 1.x bundled tracker and memory in SimpleTracker and exposed read() / write(). In 2.0 a Tracker owns no per-stream state — step is a pure function over snapshots — and MultiStream holds one snapshot per stream key. MultiStream.step(key, detections, ctx) is what you call frame to frame.

   Why: step is now referentially transparent — snapshot in, snapshot out, no hidden read() / write(). That one property is what each wrapper needs: vmap vectorizes only a function with no aliased state, MultiStream forks a stream only when nothing is shared behind its back, and autodiff needs a functional graph to differentiate. 1.x's in-place memory denied all three at once; 2.0 earns all three from a single change, and the stream and id bookkeeping move out to MultiStream / BatchTracker (Q7, Q10, Q11).

A sixth, smaller change: lifecycle and visibility are now explicit constructor arguments (NoLifecycle / StandardLifecycle, IncludeAll / ConfirmedOnly) rather than implicit behavior of the tracker wrapper. Why: birth, death, and re-ID are themselves research variables — SORT confirmation, occlusion windows, per-class rules. Making the lifecycle and visibility policies constructor parameters keeps the tracker open to new rules yet closed to edits for them: a new policy is a new argument, not a fork (Q8).

Symbol map

Several costs were also renamed in the move. The two most likely to bite: 1.x Softmax is now BiSoftmax, and 1.x BoxIoU (which computed CIoU) is now BoxCIoU, with BoxIoU reserved for plain IoU.

Worked migration: an appearance-embedding tracker

The 1.x tracker below matches detections by ReID embedding alone, gated by category and a score floor, with the Jonker–Volgenant solver. It is a typical 1.x builder: field modules at the front, a single Gate → Association stage, and per-field Value states on the memory.

Before (1.x)

import torch
import unitrack as ut
from tensordict.nn import TensorDictModule
from torch import nn

REID, SCORE, CATEGORY, LABEL = "reid", "score", "category", "label"


def _build_field(name, key=None, module=None):
    key = (name,) if key is None else (key,) if isinstance(key, str) else key
    return TensorDictModule(module or nn.Identity(), in_keys=[key], out_keys=[name])


class SelectAndFilter(nn.Module):
    def __init__(self, key_score, key_label, min_score=0.0):
        super().__init__()
        self.key_score, self.key_label, self.min_score = key_score, key_label, min_score

    def forward(self, ctx, cs, ds):
        cs_mask = torch.ones(cs.batch_size[:1], dtype=torch.bool, device=cs.device)
        ds_mask = ds.get(self.key_score) > self.min_score
        return cs_mask, ds_mask


def build_embedding_tracker(
    *,
    reid_key="reid",
    threshold=0.9,
    cost_module=ut.costs.Cosine,
    category_gate=True,
    min_score=0.0,
):
    cost = cost_module(field=REID)
    if category_gate:
        cost = ut.costs.GateCost(CATEGORY).wrap(cost)
    return ut.SimpleTracker(
        tracker=ut.MultiStageTracker(
            fields=[
                _build_field(REID, reid_key),
                _build_field(SCORE, "score"),
                _build_field(CATEGORY, "category"),
            ],
            stages=[
                ut.stages.Gate(
                    gate=SelectAndFilter(SCORE, LABEL, min_score=min_score),
                    then=[
                        ut.stages.Association(
                            cost=cost,
                            assignment=ut.assignment.Jonker(threshold=threshold),
                        )
                    ],
                ),
            ],
        ),
        memory=ut.TrackletMemory(
            states={
                SCORE: ut.states.Value(torch.float),
                CATEGORY: ut.states.Value(torch.long),
                REID: ut.states.Value(torch.float, shape=(256,)),
            }
        ),
    )

After (2.0)

import torch
import unitrack as ut
from unitrack.assignment import Associate, Jonker
from unitrack.costs import Cosine
from unitrack.data import TensorSpec
from unitrack.gates import ClassGate, ScoreGate
from unitrack.lifecycle import IncludeAll, NoLifecycle
from unitrack.pipeline import Gated, Pipe, Sequential
from unitrack.states import FromDetectionField, Identity, Replace, State

REID, SCORE, CATEGORY = "reid", "score", "category"


def _feature_state(name: str, shape: tuple[int, ...], dtype: torch.dtype) -> State:
    """Pure feature cache: no motion model, replace on match, seed from detection."""
    return State(
        schema=TensorSpec(shape=shape, dtype=dtype),
        process=Identity(name),
        observation=Replace(name),
        init=FromDetectionField(name),
    )


def build_embedding_tracker(
    *,
    threshold: float = 0.1,
    cost_cls: type = Cosine,
    category_gate: bool = True,
    min_score: float = 0.0,
    reid_dim: int = 256,
) -> ut.Tracker:
    gates = [ScoreGate(SCORE, threshold=min_score)]
    if category_gate:
        gates.insert(0, ClassGate(CATEGORY))
    root = Gated(
        gate=Sequential(gates),
        then=Pipe(cost=cost_cls(REID), assoc=Associate(Jonker(threshold=threshold))),
    )
    return ut.Tracker(
        root=root,
        states={
            REID: _feature_state(REID, (reid_dim,), torch.float32),
            SCORE: _feature_state(SCORE, (), torch.float32),
            CATEGORY: _feature_state(CATEGORY, (), torch.int64),
        },
        lifecycle=NoLifecycle(),
        visibility=IncludeAll(),
    )

What dissolved, line by line

• _build_field and the fields=[…] list. The front-of-pipeline field selection has no analogue in 2.0. You no longer rename tensors inside the tracker; you name them when you build the Detections record (next section). The reid_key parameter — which existed only to rename an input tensor to "reid" — disappears with it.
• GateCost(CATEGORY).wrap(cost). Category gating is now a ClassGate, a GateProducer that emits a per-pair equality mask. It composes with the score gate through Sequential and is applied to the cost stage by Gated, instead of being woven into the cost object.
• SelectAndFilter. Its score branch (ds_mask = score > min_score) becomes ScoreGate(SCORE, threshold=min_score). Its cs_mask was always all-True, i.e. a no-op on the tracklet side, so nothing replaces it. If you did filter tracklets by status here, that is now a Filter(StatusFilter(…), on="cs") node.
• states.Value(dtype, shape=…). Each becomes a State. A Value was a cache with no dynamics, which is exactly Identity (no predict) + Replace (overwrite matched rows) + FromDetectionField (seed new rows). TensorSpec carries the same (shape, dtype) the 1.x Value did.
• SimpleTracker and TrackletMemory. The builder now returns a bare Tracker. The per-stream snapshot that SimpleTracker and TrackletMemory used to own moves into MultiStream, created at the call site.

Feeding the tracker

In 1.x, an input module (multiformer's TrackerInput) produced a dict of tensors, the fields=[…] modules selected from it, and SimpleTracker handled read/write. In 2.0 you build a Detections record and call MultiStream.step. The record needs a reserved index field (int64, shape (M,)) that threads your detection ordering through to the MatchOutcome; the remaining fields are the ones your states read.

from unitrack.data import Detections, FrameContext

tracker = build_embedding_tracker()
ms = ut.MultiStream(tracker)

for frame_idx, frame in enumerate(stream):  # your per-frame source
    reid = frame["reid"]  # (M, 256) float32
    score = frame["score"]  # (M,)     float32
    category = frame["category"]  # (M,)     int64
    m = score.shape[0]

    det = Detections(
        index=torch.arange(m, dtype=torch.int64, device=score.device),
        reid=reid,
        score=score,
        category=category,
        batch_size=[m],
    )
    res = ms.step(0, det, FrameContext.make(frame_idx, stream_key=0))

    # res.ids: int64 tracklet IDs visible this frame (per the Visibility policy).
    # res.match.matched_pairs: (K, 2) [tracklet_row, detection_row]; recover the
    # original detection via det.index[pair[1]], and stable identity via
    # res.snapshot.id. See the MatchOutcome reference for the residual indices.
    for tracklet_row, det_row in res.match.matched_pairs.tolist():
        det_id = det.index[det_row].item()
        ...

Preprocessing that produced detection tensors in 1.x is unaffected. multiformer's MaskToBoxes, for instance, is a plain tensor op upstream of the tracker; its (N, 4) output now becomes a Detections field instead of an entry in the dict the field modules read.

Behavioral notes

• Re-derive thresholds; do not copy the 1.x number. In 2.0, Cosine returns 1 − cosine_similarity (a distance), Jonker(threshold=t) masks pairs whose cost is strictly above t, and Associate keeps matched pairs with cost <= t. So a 1.x similarity floor of 0.9 corresponds to a 2.0 distance ceiling of 0.1 — hence threshold=0.1 in the migrated builder. Check the cost's range and direction before reusing a tuned value.
• NoLifecycle() + IncludeAll() reproduces the 1.x embedding tracker. It had no birth/death logic — every tracklet matched and every ID was visible. To add SORT-style confirmation, switch to StandardLifecycle(min_hits, max_age) and ConfirmedOnly(); the recipes/sort recipe shows the full form.
• Detections is immutable and validated. Construction checks that index is present, int64, and shape (M,). Use Detections.empty() for an empty-frame placeholder.

See also

• notebooks/tutorials/migration — this migration run as an interactive notebook: it ports the tracker above and then showcases the new 2.0 possibilities (parallel fusion, cascaded matching, lifecycle, and differentiable matching) on detections from a real lightweight detector.
• recipes/overlap_tracker — the closest sibling: a single-stage class- and score-gated matcher, ported from the 1.x models.overlap builder.
• recipes/sort — a stateful IoU + Kalman tracker, showing StandardLifecycle, ConfirmedOnly, and a Kalman state.