mixle.task.edge module

Edge distillation under hard device budgets.

mixle.task.tune_recipe() tunes one family’s knobs with a soft cost penalty. An edge deployment is a different problem: the budget is hard (bytes of flash, ops per inference, “no torch on the device”), and the largest gains often come from choosing the right kind of student – a hashed-feature MLP versus a structured probabilistic classifier (a learned Bayesian network: kilobytes, torch-free, exactly calibrated posteriors) – not from tweaking widths inside one kind.

This module makes both choices searchable, and makes the search itself a model:

• DeviceSpec / EdgeFootprint – the deployment budget and a student’s measured cost (serialized bytes, per-inference op count, torch dependence).

• EdgeSpace – one design space spanning family (structure) and each family’s training recipe (process), decoded from a unit cube so any optimizer can drive it.

• DesignModel – the model that writes the model: GP surrogates over design -> (quality, budget violations), proposing the next design by feasibility-weighted expected improvement (mixle.doe.propose_next_constrained()). It persists (to_json/from_json) and keys every row on a task_fingerprint(), so design knowledge accumulates across tasks: the surrogate learns task-similarity through the fingerprint coords, and a warm-started search – even on a different task – skips the designs that never work anywhere near it.

• distill_for_edge() – the front door: screen candidates at reduced fidelity, promote the promising ones to full training, return the best student that fits the device, with an adaptive selection Pareto front, an independent qualification score, and the updated DesignModel.

• distill_designer() – distill the accumulated design ledger into a torch-free structured student that predicts whether a design is worth training.

Teachers are consulted once per dataset (labels are cached), never per candidate.

class DeviceSpec(max_bytes=None, max_ops=None, max_inference_seconds=None, torch_free=False)

Bases: object

A hard deployment budget. None leaves an axis unconstrained.

max_bytes gates the exact exported artifact. max_inference_seconds gates a benchmark of the actual serving path. max_ops is retained as a structural compatibility constraint, but cannot be used without an actual latency bound. torch_free excludes torch-backed artifacts.

Parameters:

• max_bytes (int | None)

• max_ops (int | None)

• max_inference_seconds (float | None)

• torch_free (bool)

classmethod for_latency(max_ms, ops_per_second, *, max_bytes=None, torch_free=False)

A budget from a latency target: max_ops = ops_per_second * max_ms / 1000.

ops_per_second must come from a probe run on the target device for the student kind you deploy (measure_ops_per_second() measures it for a representative student) – throughput differs by orders of magnitude across devices and student kinds, so there is no portable built-in constant.

Parameters:

• max_ms (float)

• ops_per_second (float)

• max_bytes (int | None)

• torch_free (bool)

Return type:

DeviceSpec

violations(fp)

Normalized constraint values, feasible when <= 0 (the form constrained BO consumes).

Parameters:

fp (EdgeFootprint)

Return type:

list[float]

feasible(fp)

Return whether a measured footprint satisfies device constraints.

Parameters:

fp (EdgeFootprint)

Return type:

bool

class EdgeFootprint(bytes, ops, torch_free, inference_seconds=None, artifact_digest=None, ops_kind='structural_estimate', payload_bytes=None)

Bases: object

Exact exported bytes plus execution-path evidence.

ops is an explicitly labeled structural estimate. inference_seconds is measured by running the serving object’s real batch path.

bytes is what you flash to the device: weights and the manifest, adapter spec, labels and metadata. payload_bytes is the parameter payload alone (weights, biases, per-layer scales), which is the quantity a compression ratio is about – at small model sizes the constant-size manifest dominates bytes, so a 4x/8x weight claim is only checkable against payload_bytes. Gate deployment on bytes; quote compression from payload_bytes.

Parameters:

• bytes (int)

• ops (int)

• torch_free (bool)

• inference_seconds (float | None)

• artifact_digest (str | None)

• ops_kind (str)

• payload_bytes (int | None)

class EdgeSpace(families=('mlp', 'structured'), dim_choices=(64, 128, 256, 512), hidden_range=(4, 96), epochs_range=(40, 320), log10_lr_range=(-3.0, -1.0), bits_choices=(32, 8), ngram=3, components_range=(1, 4), bins_range=(2, 8), max_its_range=(10, 60), min_gain_range=(0.0, 5.0))

Bases: object

One unit-cube design space over family (structure) and each family’s recipe (process).

Coordinate 0 selects the family; 1..4 decode family-specifically. families defaults to what the input kind and device allow: hashed-MLP students need torch; structured students need fixed-schema records. Decode is deterministic, so a design point is a reproducible recipe.

Parameters:

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

• dim_choices (Sequence[int])

• hidden_range (tuple[int, int])

• epochs_range (tuple[int, int])

• log10_lr_range (tuple[float, float])

• bits_choices (Sequence[int])

• ngram (int)

• components_range (tuple[int, int])

• bins_range (tuple[int, int])

• max_its_range (tuple[int, int])

• min_gain_range (tuple[float, float])

dims()

Return the normalized design-space dimensionality.

Return type:

int

bounds()

Return normalized design-space bounds for DOE search.

Return type:

list[tuple[float, float]]

signature()

Fingerprint of the space so persisted design knowledge is only reused where it applies.

Return type:

str

decode(point)

Unit-cube point -> (family, recipe kwargs); the mlp recipe carries a bits precision.

Parameters:

point (ndarray)

Return type:

tuple[str, dict[str, Any]]

class DesignModel(signature, n_constraints, n_fingerprint=0)

Bases: object

A probabilistic model of the design space itself: design point -> (quality, budget violations).

Every evaluated design is a row; GP surrogates fitted on the rows drive propose() (the next design worth training, by feasibility-weighted expected improvement) and predict() (mean, sd, and probability-of-fitting-the-device for untrained designs). It serializes, so what was learned designing students for one task warm-starts the next – the design knowledge is itself a model artifact. (And distill_designer() compresses it into a student – models all the way down, each level a real artifact.)

Parameters:

• signature (str)

• n_constraints (int)

• n_fingerprint (int)

add(point, quality, violations, *, fingerprint=None, **tag)

Append one evaluated design point and its feasibility metadata.

Parameters:

• point (Any)

• quality (float)

• violations (Sequence[float])

• fingerprint (Sequence[float] | None)

• tag (Any)

Return type:

None

propose(bounds, *, seed=None, n_candidates=256, prefilter=None, max_tries=8, fingerprint=None)

The next design worth training: feasibility-weighted EI over everything seen so far.

prefilter closes the designer loop: pass a design judge – typically the compact student from distill_designer(), called as prefilter(point_tuple) -> label – and any proposal it labels "weak" is vetoed and re-drawn (fresh acquisition seed), up to max_tries. The distilled design knowledge thus skips known weak designs before a single training run is spent; if every retry is vetoed the last proposal is returned anyway (the judge advises, the surrogate decides). fingerprint conditions the proposal on the current task (see _fingerprint_bounds()); the returned point has design coords only.

Parameters:

• bounds (Sequence[tuple[float, float]])

• seed (Any)

• n_candidates (int)

• prefilter (Callable[[tuple], Any] | None)

• max_tries (int)

• fingerprint (Sequence[float] | None)

Return type:

ndarray

predict(points, *, fingerprint=None)

For untrained designs: predicted quality (mean, sd) and P(fits the device).

points carry design coords only; fingerprint (required when the ledger is fingerprinted) selects which task’s slice the prediction conditions on.

Parameters:

• points (Any)

• fingerprint (Sequence[float] | None)

Return type:

dict[str, ndarray]

to_json()

Serialize the design ledger for reuse across search runs.

Return type:

dict[str, Any]

classmethod from_json(d)

Reconstruct a design ledger from serialized JSON data.

Parameters:

d (dict[str, Any])

Return type:

DesignModel

class EdgeDistillResult(model, family, recipe, agreement, selection_agreement, test_size, footprint, feasible, pareto, design, trials=<factory>)

Bases: object

Outcome of edge search and independent deployment qualification.

selection_agreement and pareto describe adaptive evidence from the search panel. agreement is measured once for the frozen winner on a disjoint test panel and is the only reported deployment-quality estimate.

Parameters:

• model (TaskModel)

• family (str)

• recipe (dict[str, Any])

• agreement (float)

• selection_agreement (float)

• test_size (int)

• footprint (EdgeFootprint)

• feasible (bool)

• pareto (list[dict[str, Any]])

• design (DesignModel)

• trials (list[dict[str, Any]])

distill_for_edge(teacher, train_data, val_data, device, *, test_data=None, test_frac=0.25, labels=None, train_labels=None, val_labels=None, test_labels=None, space=None, design=None, designer=None, n_init=4, n_iter=6, screen_fidelity=0.3, promote=2, seed=0, task='')

Search structure x process, then qualify the frozen winner independently.

Candidate screening, promotion, winner selection, and the returned Pareto front use only val_data (the search panel). The selected artifact is then scored exactly once on untouched test_data. If test_data is omitted, val_data is deterministically split into disjoint search and test panels. Accordingly, result.selection_agreement and every Pareto point are adaptive search evidence; result.agreement is independent deployment evidence.

The teacher labels each panel once, or callers may provide every matching label panel and pass teacher=None. Validation or test labels never expand the class vocabulary: when labels is omitted it is derived from training labels alone. Pass a previous compatible design to warm-start or a distilled designer to prefilter known-weak proposals. If nothing fits the device, the least-infeasible selected student is returned with feasible=False.

Parameters:

• teacher (Callable[[...], Any] | None)

• train_data (Sequence[Any])

• val_data (Sequence[Any])

• device (DeviceSpec)

• test_data (Sequence[Any] | None)

• test_frac (float)

• labels (Sequence[str] | None)

• train_labels (Sequence[Any] | None)

• val_labels (Sequence[Any] | None)

• test_labels (Sequence[Any] | None)

• space (EdgeSpace | None)

• design (DesignModel | None)

• designer (Callable[[tuple], Any] | None)

• n_init (int)

• n_iter (int)

• screen_fidelity (float)

• promote (int)

• seed (int)

• task (str)

Return type:

EdgeDistillResult

distill_designer(design, *, quality_quantile=0.5, seed=0)

Distill the design ledger into a compact torch-free student that judges designs: point -> good/weak.

The DesignModel’s rows are records (the design coordinates); a row is labeled good when it was feasible on the device and its quality reached the ledger’s quality_quantile. distill_structured_from_labels() – the same machinery the design model tunes – compresses that knowledge into a kilobyte Bayesian-network student usable as a zero-cost pre-filter for future searches. Teacher trains student; a model designs the students; the designer is distilled into a student: each level of the tower is a real artifact.

Parameters:

• design (DesignModel)

• quality_quantile (float)

• seed (int)

Return type:

TaskModel

footprint(student, *, probe_inputs=None, repeats=3)

Measure the exact exported artifact and optionally benchmark serving.

The artifact byte count comes from TaskModel.save(), so weights use their actual dtype and the manifest, adapter, labels, configuration, and metadata are included. payload_bytes reports the parameter payload on its own, because the manifest is a near-constant cost that swamps the weights on small students – an int4 student whose weights are 8x smaller than fp32 still exports a comparable manifest. ops remains a structural estimate and is never sufficient by itself for a hard latency gate.

Parameters:

• student (TaskModel)

• probe_inputs (Sequence[Any] | None)

• repeats (int)

Return type:

EdgeFootprint

task_fingerprint(data, labels)

A fixed small vector describing which task this is: the coords cross-task warm start keys on.

(log10 #examples, #labels, #fields, fraction of categorical fields, normalized label entropy) – low-overhead invariants of the dataset, O(1)-scaled so the design surrogate’s default lengthscale treats similar tasks as informative neighbors and dissimilar ones as weakly coupled.

Parameters:

• data (Sequence[Any])

• labels (Sequence[Any])

Return type:

list[float]

measure_inference_seconds(student, inputs, *, repeats=5)

Median measured wall-clock seconds per single-input inference for student on this host.

Runs student.batch(inputs) repeats times (after one untimed warm-up) and reports the median per-item time. Measured, not modeled – run it on the machine whose latency you care about (the deploy device, not the dev laptop) for a number that means anything there.

Parameters:

• student (TaskModel)

• inputs (Sequence[Any])

• repeats (int)

Return type:

float

measure_ops_per_second(student, inputs, *, repeats=5)

Measured throughput (footprint ops / measured second) for this student kind on this host.

The calibration constant that turns a latency budget into DeviceSpec max_ops (DeviceSpec.for_latency()): probe once per (device, student kind), reuse across searches.

Parameters:

• student (TaskModel)

• inputs (Sequence[Any])

• repeats (int)

Return type:

float