def build_structured_dataset(source, groups, *, input_specs, target_specs, spec, sample_meta=None, available=None):
    """Validate and fingerprint named scientific samples without retaining a corpus.

    Parameters
    ----------
    source : sequence
        Stable provider of StructuredSample objects; lazy file access is allowed [-].
    groups : array_like
        Independent scientific group identity for every sample [-].
    input_specs : mapping
        Named input TensorSpec records with shape, dtype, units and masks [-].
    target_specs : mapping
        Named fixed-shape target TensorSpec records [-].
    spec : DatasetSpec
        Task, group convention and source-role/preprocessing provenance [-].
    sample_meta : mapping or None, optional
        Per-sample machine, time or identity arrays [any].
    available : array_like or None, optional
        Boolean sample-level input availability; True means available [-].

    Returns
    -------
    StructuredDatasetArtifact
        Validated source and schema with a content SHA-256 fingerprint [-].

    Processing steps
    ----------------
    1. Validate names, shapes, mask broadcasts and finite valid elements.
    2. Stream metadata and every input, target and validity array into SHA-256.
    3. Retain the stable source and fingerprint, without retaining sampled tensors.

    Convention
    ----------
    True means valid. Missing observations differ from physical zero. Only a
    leading input axis may be variable; target availability uses separate masks.

    Applicability
    -------------
    Machine-independent.

    Limitations
    -----------
    A source must have stable contents. DataLoader workers require a picklable
    source. Fixed output shapes are required; this is not an IMAS data model.

    Provenance
    ----------
    .. [vaft2016] VEST-Tokamak/vaft#2016, structured observation and provenance
       extension of the existing #669 grouped ML lifecycle.
    """
    data = StructuredDatasetArtifact(source=source, groups=groups, input_specs=input_specs,
        target_specs=target_specs, spec=spec, sample_meta=sample_meta or {}, available=available)
    digest = hashlib.sha256(canonical_json({"spec": spec.to_dict(), "schema": data.schema(),
        "groups": data.group_labels.tolist(), "metadata": dict(data.sample_meta),
        "available": data.available_mask}).encode())
    for i in range(data.n_samples):
        sample = data.sample(i)
        for values in (sample.inputs, sample.targets, sample.target_masks):
            for name, a in sorted(values.items()):
                digest.update(canonical_json([name, a.shape, a.dtype.str]).encode())
                digest.update(np.ascontiguousarray(a).tobytes())
    from dataclasses import replace
    return replace(data, fingerprint=digest.hexdigest())

def dataset_fingerprint(dataset):
    """SHA-256 content fingerprint of a dataset.

    Two datasets share a fingerprint only when their spec, feature schema,
    inputs, targets, groups and per-sample metadata are identical in value,
    dtype and shape.  A model records this fingerprint as part of its identity,
    so "which data trained it" is answerable from the artifact alone.

    Parameters
    ----------
    dataset : DatasetArtifact
        The dataset to fingerprint; its own ``fingerprint`` field is ignored [-].

    Returns
    -------
    str
        64-character lowercase hex digest [-].

    Processing steps
    ----------------
    1. Hash the canonical JSON of the spec (and the feature spec, if any).
    2. Hash inputs, targets, string group labels, the availability mask (when
       one is set) and each ``sample_meta`` array in sorted key order, each
       prefixed by its name, dtype and shape.

    Applicability
    -------------
    Machine-independent.
    """
    digest = hashlib.sha256()
    digest.update(canonical_json(dataset.spec.to_dict()).encode())
    if isinstance(dataset, FeatureDataset):
        digest.update(canonical_json(dataset.feature_spec.to_dict()).encode())
    _hash_array(digest, "inputs", np.asarray(dataset.inputs))
    if dataset.targets is not None:
        _hash_array(digest, "targets", np.asarray(dataset.targets))
    _hash_array(digest, "groups", dataset.group_labels.astype(str))
    if dataset.available is not None:
        _hash_array(digest, "available", np.asarray(dataset.available))
    for key in sorted(dataset.sample_meta):
        _hash_array(digest, f"meta:{key}", np.asarray(dataset.sample_meta[key]))
    return digest.hexdigest()

def export_model(model, directory, *, format="onnx"):
    """Export the network to an inference format and prove it agrees.

    The training artifact and the deployment artifact need not be the same
    thing: a model trained in PyTorch can be served by ONNX Runtime.  The
    export is accepted only after the exported graph reproduces the source
    model on a probe batch.

    Parameters
    ----------
    model : ModelArtifact
        Trained or loaded model [-].
    directory : str or path-like
        Output directory; must not exist or be empty [-].
    format : str, optional
        Target format; ``"onnx"`` is the only one implemented [-].

    Returns
    -------
    dict
        Export record written to ``export.json``: format, opset, parity
        difference and tolerance, runtimes, and the source model identity [-].

    Raises
    ------
    ModelContractError
        When the backend cannot export to ``format`` or the parity check fails.
    BackendUnavailableError
        When the exporter or runtime is not installed.

    Output semantics
    ----------------
    The exported graph maps *normalised* inputs to *normalised* outputs;
    ``normalization.npz`` is written beside it and must be applied around it,
    exactly as :func:`predict` does.  A ``sklearn`` one-class SVM's graph
    returns ``scores``, the estimator's decision function; the model's
    ``score`` is its negative.

    Applicability
    -------------
    Machine-independent.

    Provenance
    ----------
    .. [vaft669] VEST-Tokamak/vaft#669, section 12: training and inference
       artifacts are separate, with numerical validation between them.
    """
    if model.tensor_schema:
        raise ModelContractError("structured ONNX export is not qualified; use the native artifact")
    directory = Path(directory)
    if directory.exists() and any(directory.iterdir()):
        raise ModelContractError(f"{directory} is not empty")
    directory.mkdir(parents=True, exist_ok=True)
    backend = get_backend(model.model_spec.backend)
    output_shape = model.target_shape if model.supervised else model.input_shape
    record = backend.export(
        model.model_spec, model.state, model.input_shape, output_shape, directory / f"model.{format}", format
    )
    (directory / NORMALIZATION).write_bytes(_npz(model.normalization))
    record = {**record, "source": model.resolved_identity().to_dict()}
    (directory / "export.json").write_text(json.dumps(_jsonable(record), indent=2, sort_keys=True) + "\n")
    return record

def import_data_roles(dataset, assignments, allowed_operations, *, role_to_partition=None, provenance=None):
    """Import exact group roles without recomputing a scientific split.

    Parameters
    ----------
    dataset : DatasetArtifact or StructuredDatasetArtifact
        Existing stable, fingerprinted samples [-].
    assignments : mapping of str to str
        Actual independent group-to-role assignments, including absent groups [-].
    allowed_operations : mapping of str to sequence of str
        Explicit role permissions: train, preprocess, calibrate, select,
        evaluate or infer [-].
    role_to_partition : mapping or None, optional
        Explicit mapping to execution partitions; None preserves role names [-].
    provenance : mapping or None, optional
        Split/protocol revisions, origins, label methods and dependencies [-].

    Returns
    -------
    tuple
        Re-fingerprinted dataset and exact Split, without random reassignment [-].

    Processing steps
    ----------------
    1. Copy the declared assignments, permissions, mapping and provenance.
    2. Validate role coverage and record the full assignment digest.
    3. Recompute dataset identity and construct a fixed split for its groups.

    Convention
    ----------
    Roles are owner-defined; neither machine restrictions nor salts are inferred.

    Applicability
    -------------
    Machine-independent.

    Limitations
    -----------
    Re-fingerprinting streams structured samples. Sources must remain stable.
    Supplied historical provenance is a declaration, not retrospective proof.

    Provenance
    ----------
    .. [vaft2047] VEST-Tokamak/vaft#2047, exact assignment and usage provenance.
    """
    roles = {str(k): str(v) for k, v in assignments.items()}
    mapping = dict(role_to_partition) if role_to_partition is not None else {v: v for v in roles.values()}
    policy = {"schema_version": 1, "group_roles": roles,
        "allowed_operations": {str(k): sorted(set(v)) for k, v in allowed_operations.items()},
        "role_to_partition": mapping, "assignment_sha256": _hash(roles),
        "provenance": dict(provenance or {}), "history_basis": "owner_declaration"}
    spec = replace(dataset.spec, metadata={**dataset.spec.metadata, "data_use_policy": policy, "data_use_policy_sha256": _hash(policy)})
    data = replace(dataset, spec=spec)
    _policy(data)
    from .structured import StructuredDatasetArtifact, build_structured_dataset
    if isinstance(data, StructuredDatasetArtifact):
        data = build_structured_dataset(data.source, data.groups, input_specs=data.input_specs,
            target_specs=data.target_specs, spec=spec, sample_meta=data.sample_meta, available=data.available)
    else:
        from .dataset import dataset_fingerprint
        data = replace(data, fingerprint=dataset_fingerprint(data))
    fixed = {g: mapping[roles[g]] for g in sorted(set(data.group_labels))}
    if not fixed or any(not isinstance(v, str) or not v for v in fixed.values()):
        raise ModelContractError("execution partitions must be nonempty strings")
    names = sorted(set(fixed.values()))
    split = Split(spec=SplitSpec(fractions={p: 1 / len(names) for p in names}, fixed=fixed), assignment=fixed)
    return data, split

def load_dataset(directory):
    """Load a dataset written by :func:`save_dataset`, verifying its fingerprint.

    Parameters
    ----------
    directory : str or path-like
        Directory holding ``arrays.npz`` and ``dataset.json`` [-].

    Returns
    -------
    DatasetArtifact
        The dataset, a :class:`FeatureDataset` when one was saved [-].

    Raises
    ------
    ModelContractError
        When the files are missing or the recomputed fingerprint differs from
        the recorded one -- the arrays changed after they were saved.

    Applicability
    -------------
    Machine-independent.
    """
    directory = Path(directory)
    record_path = directory / _DATASET_RECORD
    if not record_path.is_file() or not (directory / _DATASET_ARRAYS).is_file():
        raise ModelContractError(f"{directory} does not hold a saved dataset")
    record = json.loads(record_path.read_text(encoding="utf-8"))
    with np.load(directory / _DATASET_ARRAYS, allow_pickle=False) as payload:
        arrays = {key: payload[key] for key in payload.files}
    spec_record = record["spec"]
    spec = DatasetSpec(
        name=spec_record["name"],
        task=spec_record["task"],
        group_key=spec_record.get("group_key", "group"),
        input_names=tuple(spec_record.get("input_names", ())),
        target_names=tuple(spec_record.get("target_names", ())),
        description=spec_record.get("description", ""),
        metadata=spec_record.get("metadata", {}),
    )
    groups = arrays["groups"]
    if np.dtype(record.get("group_dtype", groups.dtype)).kind in "iu":
        groups = groups.astype(record["group_dtype"])
    fields = dict(
        spec=spec,
        inputs=arrays["inputs"],
        groups=groups,
        targets=arrays.get("targets"),
        sample_meta={k[len("meta:"):]: v for k, v in arrays.items() if k.startswith("meta:")},
        available=arrays.get("available"),
    )
    fs = record.get("feature_spec")
    if fs is not None:
        dataset = FeatureDataset(
            feature_spec=FeatureSpec(
                feature_names=tuple(fs["feature_names"]),
                blocks={k: tuple(v) for k, v in fs.get("blocks", {}).items()},
                version=fs.get("version", ""),
            ),
            **fields,
        )
    else:
        dataset = DatasetArtifact(**fields)
    fingerprint = dataset_fingerprint(dataset)
    if fingerprint != record["fingerprint"]:
        raise ModelContractError(
            f"dataset in {directory} has fingerprint {fingerprint}, but was saved as {record['fingerprint']}"
        )
    return dataset.replace(fingerprint=fingerprint)

def load_model(name, *, version=None, stage=None, registry=None, cache=None, env=None, fetch=False):
    """Resolve a published model and load it, ready for :func:`predict`.

    Parameters
    ----------
    name : str
        Registry model name [-].
    version : str or None, optional
        Exact version; give this or ``stage`` [-].
    stage : str or None, optional
        Lifecycle alias such as ``"production"`` [-].
    registry : str or path-like or None, optional
        ``vaft-nn`` checkout; ``$VAFT_NN_HOME`` when ``None`` [-].
    cache : str or path-like or None, optional
        Artifact cache root; ``$VAFT_NN_CACHE``, else the platform cache
        directory [-].
    env : mapping or None, optional
        Environment to read the variables from [-].
    fetch : bool, optional
        On a cache miss or mismatch, run :func:`fetch_model` and retry [-].

    Returns
    -------
    ModelArtifact or ModelBundle
        The model (or bundle of models), carrying the identity
        :func:`resolve_model` returned [-].

    Applicability
    -------------
    Machine-independent.
    """
    from .artifact import load_model_artifact
    from .bundle import _with_identity, load_model_bundle

    env = os.environ if env is None else env
    try:
        identity = resolve_model(name, version=version, stage=stage, registry=registry, cache=cache, env=env)
    except ModelResolutionError:
        if not fetch:
            raise
        fetched = fetch_model(name, version=version, stage=stage, registry=registry, cache=cache, env=env)
        identity = ModelIdentity(**{**fetched.to_dict(), "stage": stage})
    root = _registry_root(registry, env)
    manifest = root / "models" / name / "versions" / identity.version / "manifest.json"
    kind = json.loads(manifest.read_bytes()).get("kind", "model")
    # verify=False skips only the size/presence pass resolve_model just made;
    # every byte loaded is still hashed, and the manifest re-read here must be
    # the one whose hash the registry pins.
    if kind == "bundle":
        loaded = load_model_bundle(identity.source, manifest=manifest, verify=False)
    else:
        loaded = load_model_artifact(identity.source, manifest=manifest, verify=False)
    if loaded.identity.manifest_sha256 != identity.manifest_sha256:
        raise ModelResolutionError(f"the registry manifest of {name} {identity.version} changed after verification")
    if kind == "bundle":
        return _with_identity(loaded, identity)
    return loaded.replace(identity=identity)

def load_model_artifact(directory, *, manifest=None, verify=True):
    """Load an artifact directory, verifying every file against its manifest.

    Parameters
    ----------
    directory : str or path-like
        Directory holding the weights and normalisation files [-].
    manifest : str or path-like or None, optional
        Manifest to trust; ``directory/manifest.json`` when ``None``.  The
        resolver passes the registry's reviewed copy here, so a tampered cache
        cannot vouch for itself [-].
    verify : bool, optional
        Check the SHA-256 and size of every file before loading [-].

    Returns
    -------
    ModelArtifact
        The model with an identity naming the directory, version and manifest
        hash [-].

    Raises
    ------
    ModelResolutionError
        When the manifest is missing or a file is missing or differs.
    ModelContractError
        When the manifest's schema version is not one this VAFT reads.

    Applicability
    -------------
    Machine-independent.

    Provenance
    ----------
    .. [vaft669] VEST-Tokamak/vaft#669, section 13: resolution verifies identity,
       version and artifact hash.
    """
    directory = Path(directory)
    manifest_path = Path(manifest) if manifest is not None else directory / MANIFEST
    if not manifest_path.is_file():
        raise ModelResolutionError(f"no manifest at {manifest_path}")
    raw = manifest_path.read_bytes()
    record = json.loads(raw)
    if record.get("kind", "model") != "model":
        raise ModelContractError(f"{manifest_path} describes a {record.get('kind')!r}; use load_model_bundle")
    model = _load_artifact(directory, record, verify=verify)
    identity = ModelIdentity(
        name=model.name,
        state_sha256=model.state_sha256,
        version=record["version"],
        manifest_sha256=sha256_bytes(raw),
        resolved_by="path",
        source=str(directory),
    )
    return model.replace(identity=identity)

def load_model_bundle(directory, *, manifest=None, verify=True):
    """Load a bundle directory, verifying every file and member manifest.

    Parameters
    ----------
    directory : str or path-like
        Directory holding the bundle's files [-].
    manifest : str or path-like or None, optional
        Bundle manifest to trust; ``directory/manifest.json`` when ``None``.
        The resolver passes the registry's reviewed copy [-].
    verify : bool, optional
        Check every file's SHA-256 and size first [-].

    Returns
    -------
    ModelBundle
        The members, each with an identity naming the bundle, and the
        composition record [-].

    Raises
    ------
    ModelResolutionError
        When a file is missing or differs, or a member manifest differs from
        the hash the bundle pins.
    ModelContractError
        When the manifest is not a bundle manifest this VAFT reads.

    Applicability
    -------------
    Machine-independent.
    """
    directory = Path(directory)
    manifest_path = Path(manifest) if manifest is not None else directory / MANIFEST
    if not manifest_path.is_file():
        raise ModelResolutionError(f"no manifest at {manifest_path}")
    raw = manifest_path.read_bytes()
    record = json.loads(raw)
    if record.get("schema_version") != MANIFEST_SCHEMA_VERSION or record.get("kind") != "bundle":
        raise ModelContractError(f"{manifest_path} is not a schema-{MANIFEST_SCHEMA_VERSION} bundle manifest")
    if verify:
        _verify_files(record, directory)
    identity = ModelIdentity(
        name=record["model"],
        state_sha256=sha256_bytes(raw),
        version=record["version"],
        manifest_sha256=sha256_bytes(raw),
        resolved_by="path",
        source=str(directory),
    )
    members = {}
    for role, entry in record["members"].items():
        member_raw = (directory / entry["manifest"]).read_bytes()
        if sha256_bytes(member_raw) != entry["manifest_sha256"]:
            raise ModelResolutionError(f"member {role!r} manifest differs from the hash the bundle pins")
        member = _load_artifact(directory, json.loads(member_raw), verify=verify, prefix=f"{role}.")
        members[role] = member.replace(identity=_member_identity(member, role, entry["manifest_sha256"], identity))
    return ModelBundle(
        name=record["model"], members=members, composition=record.get("composition") or {}, identity=identity
    )

def predict(model, data):
    """Run a model on new data and stamp the result with its exact identity.

    Accepts a dataset (checked against the model's feature schema) or a bare
    array.  The result records the model's immutable identity -- name, exact
    version, manifest and weight hashes -- never an alias such as
    ``production``, so a stored result keeps meaning what it meant when it was
    produced.

    Parameters
    ----------
    model : ModelArtifact
        Trained, loaded or resolved model [-].
    data : DatasetArtifact or array_like
        Inputs of shape ``(n, *model.input_shape)`` [any].

    Returns
    -------
    InferenceResult
        Outputs aligned with the samples (``prediction``; ``reconstruction``
        and a per-sample ``score``; or ``score`` alone; plus
        ``<name>_exceeds`` for every calibration and, for a dataset with an
        availability mask, ``available``), the model identity, and the groups
        and per-sample metadata of a dataset input [any].

    Raises
    ------
    ModelContractError
        When the input shape, or a dataset's column names, differ from the
        model's.

    Output semantics
    ----------------
    ``prediction`` and ``reconstruction`` are in the units of the targets and
    inputs; an autoencoder's ``score`` is the mean squared reconstruction
    error in the model's normalised units, and a score model's is the
    detector's own anomaly measure (larger is more anomalous).  An
    unavailable sample has NaN outputs and exceeds no threshold: missing
    data is not "no anomaly".  A score or an exceedance is a model output, not a
    physical event: an ``IRE-like candidate``, never a reconnection onset
    (#669 section 16).

    Applicability
    -------------
    Machine-independent.
    """
    from .usage import _inference_use
    usage = _inference_use(data) if hasattr(data, "spec") else None
    from .structured import StructuredDatasetArtifact
    if isinstance(data, StructuredDatasetArtifact):
        from ._structured_lifecycle import predict_structured
        result = predict_structured(model, data)
        from dataclasses import replace
        return replace(result, provenance={**result.provenance, "data_use": usage}) if usage else result
    if model.tensor_schema:
        raise ModelContractError("structured models require a StructuredDatasetArtifact")
    if isinstance(data, DatasetArtifact):
        _check_columns(model, data)
        x = np.asarray(data.inputs)
        available = data.available
        groups, meta = data.groups, dict(data.sample_meta)
        provenance = {"input": "dataset", "dataset_fingerprint": data.fingerprint}
    else:
        x = np.asarray(data)
        available = None
        groups, meta = None, {}
        digest = hashlib.sha256(np.ascontiguousarray(x).tobytes()).hexdigest()
        provenance = {"input": "array", "input_sha256": digest}
    if usage:
        provenance["data_use"] = usage
    outputs = _forward(model, x, available)
    return InferenceResult(
        outputs=outputs,
        model_identity=model.resolved_identity(),
        groups=groups,
        sample_meta=meta,
        provenance=provenance,
    )

def record_data_use(dataset, split, operation, *, groups=None):
    """Check a declared use and return its provenance receipt.

    Parameters
    ----------
    dataset : DatasetArtifact or StructuredDatasetArtifact
        Dataset with an optional data_use_policy [-].
    split : Split
        Realized execution assignment to verify against the declaration [-].
    operation : str
        Requested train, preprocess, calibrate, select, evaluate or infer [-].
    groups : iterable of str or None, optional
        Groups actually consumed; None selects every dataset group [-].

    Returns
    -------
    dict or None
        Usage receipt; None means no policy was supplied, not validated history [-].

    Processing steps
    ----------------
    1. Verify the dataset's execution assignment matches its declared mapping.
    2. Check explicit permissions for every selected group.
    3. Return dataset, policy and assignment identities with observed group roles.

    Convention
    ----------
    Calling this reports a requested entrypoint use; it cannot inspect external
    training history or intercept scientific operations outside this API.

    Applicability
    -------------
    Machine-independent.

    Provenance
    ----------
    .. [vaft2047] VEST-Tokamak/vaft#2047, opt-in operation-role checks.
    """
    policy = _policy(dataset)
    if policy is None:
        return None
    if operation not in _OPERATIONS:
        raise ModelContractError("unknown data-use operation")
    roles = policy["group_roles"]
    if _hash(roles) != policy["assignment_sha256"]:
        raise ModelContractError("data assignment digest does not match declaration")
    for g in set(dataset.group_labels):
        if split.assignment.get(g) != policy["role_to_partition"][roles[g]]:
            raise ModelContractError("split differs from declared data roles")
    selected = sorted(set(dataset.group_labels if groups is None else map(str, groups)))
    if set(selected) - set(dataset.group_labels):
        raise ModelContractError("requested use contains groups outside the dataset")
    for g in selected:
        if operation not in policy["allowed_operations"][roles[g]]:
            raise ModelContractError(f"{operation} is forbidden for role {roles[g]!r}, group {g!r}")
    return {"operation": operation, "groups": selected, "group_roles": {g: roles[g] for g in selected},
        "dataset_fingerprint": dataset.fingerprint, "assignment_sha256": policy["assignment_sha256"],
        "policy_sha256": _hash(policy), "history_basis": "declared_use_at_entrypoint"}

def register_architecture(name, factory, *, kind, backend="torch"):
    """Make a task-defined network available to a backend by name.

    Task code (the SXR reconstruction transformer, a sequence model) defines
    its network next to its science and registers it here; a
    :class:`ModelSpec` then selects it with ``architecture=name``.  The
    factory is called only when a model is built, so registering imports no
    ML framework.

    Parameters
    ----------
    name : str
        Architecture name a ``ModelSpec`` refers to [-].
    factory : callable
        ``factory(spec, input_shape, output_shape)`` returning a network that
        maps a batch of shape ``(n, *input_shape)`` to ``(n, *output_shape)``
        without flattening it first [-].
    kind : {"supervised", "reconstruction", "score"}
        What the network learns [-].
    backend : str, optional
        Backend the factory builds for; only ``"torch"`` accepts registrations [-].

    Returns
    -------
    None
        The architecture is registered for this process [-].

    Raises
    ------
    ModelContractError
        On an unknown kind, a backend that does not accept registrations, a
        name that shadows a built-in, or a second registration of one name
        with a different factory.

    Applicability
    -------------
    Machine-independent.

    Limitations
    -----------
    The registry is per process.  A published model records its
    architecture's name, not its code; loading it needs the same task module
    imported, and the code version is traced through the VAFT revision in the
    manifest.
    """
    if kind not in KINDS:
        raise ModelContractError(f"kind must be one of {KINDS}; got {kind!r}")
    _check_name("architecture", name)
    target = get_backend(backend)
    if not target.extensible:
        raise ModelContractError(f"backend {backend!r} does not accept registered architectures")
    if kind not in target.trainable_kinds:
        raise ModelContractError(f"backend {backend!r} trains {target.trainable_kinds}, not {kind!r}")
    if name in target.builtin:
        raise ModelContractError(f"{name!r} is a built-in {backend} architecture")
    registered = _REGISTERED.setdefault(backend, {})
    existing = registered.get(name)
    if existing is not None and (existing.factory is not factory or existing.kind != kind):
        raise ModelContractError(f"architecture {name!r} is already registered for {backend!r}")
    registered[name] = Architecture(name=name, kind=kind, factory=factory)

def register_augmentation(name, function):
    """Make a train-time augmentation available by name.

    Selected with ``TrainingConfig(augmentation=name)``.  It is applied to the
    **train partition only**, once per epoch, after the group split -- so an
    augmented copy of a sample can never reach validation or test.

    Parameters
    ----------
    name : str
        Augmentation name [-].
    function : callable
        ``function(inputs, targets, rng)`` returning ``(inputs, targets)`` of
        the same shapes, in physical (un-normalised) units; ``targets`` is
        ``None`` for an unsupervised model and ``rng`` is a seeded
        ``numpy.random.Generator`` [-].

    Returns
    -------
    None
        The augmentation is registered for this process [-].

    Raises
    ------
    ModelContractError
        On a name already registered with a different function.

    Applicability
    -------------
    Machine-independent.
    """
    _check_name("augmentation", name)
    if AUGMENTATIONS.get(name, function) is not function:
        raise ModelContractError(f"augmentation {name!r} is already registered")
    AUGMENTATIONS[name] = function

def register_loss(name, factory):
    """Make a task-defined training loss available by name.

    Selected with ``ModelSpec(hyperparameters={"loss": name})``.  The factory
    receives the model spec and the train-partition normalisation, so a
    physics loss -- a line-of-sight residual through a projector acting on
    emissivity in physical units -- can undo the normalisation the network
    works in.

    Parameters
    ----------
    name : str
        Loss name [-].
    factory : callable
        ``factory(spec, normalization)`` returning ``loss(prediction, target,
        inputs)``, a differentiable scalar in the backend's tensor type.
        ``prediction``, ``target`` and ``inputs`` are normalised and keep
        their sample shapes; ``normalization`` maps ``input_mean``,
        ``input_std`` and, for a supervised model, ``target_mean`` and
        ``target_std`` to NumPy arrays that broadcast against one sample [-].

    Returns
    -------
    None
        The loss is registered for this process [-].

    Raises
    ------
    ModelContractError
        On a name that is already registered with a different factory, or
        that shadows the built-in ``"mse"``.

    Applicability
    -------------
    Machine-independent.

    Limitations
    -----------
    Only the iterative ``torch`` backend trains with a loss; the closed-form
    and scikit-learn backends refuse a ``loss`` hyperparameter.
    """
    _check_name("loss", name)
    if name == "mse":
        raise ModelContractError("'mse' is the built-in loss")
    if LOSSES.get(name, factory) is not factory:
        raise ModelContractError(f"loss {name!r} is already registered")
    LOSSES[name] = factory

def resolve_model(name, *, version=None, stage=None, registry=None, cache=None, env=None):
    """Resolve a published model to one exact, hash-verified local version.

    Parameters
    ----------
    name : str
        Registry model name, e.g. ``"ire_ae_vaft_v1"`` [-].
    version : str or None, optional
        Exact version; give this or ``stage`` [-].
    stage : str or None, optional
        Lifecycle alias such as ``"production"``, resolved to an exact version
        before anything is loaded [-].
    registry : str or path-like or None, optional
        ``vaft-nn`` checkout; ``$VAFT_NN_HOME`` when ``None`` [-].
    cache : str or path-like or None, optional
        Artifact cache root; ``$VAFT_NN_CACHE``, else the platform cache
        directory [-].
    env : mapping or None, optional
        Environment to read the two variables from; ``os.environ`` when ``None`` [-].

    Returns
    -------
    ModelIdentity
        Name, exact version, manifest SHA-256, requested stage, lifecycle status
        and the verified artifact directory (``source``) [-].

    Raises
    ------
    ModelResolutionError
        When the registry, model, version, stage, manifest or a cached file is
        missing, or a hash differs from the one the registry pins.
    ModelContractError
        When both or neither of ``version`` and ``stage`` are given.

    Processing steps
    ----------------
    1. Read ``models/<name>/releases.yaml`` from the registry.
    2. Map ``stage`` to its version, or take ``version``; refuse a stage that
       points at a deprecated version.
    3. Check the registry manifest's SHA-256 against the pinned one.
    4. Check every file the manifest lists in ``<cache>/<name>/<version>/``
       for presence, size and SHA-256.

    Convention
    ----------
    A stage alias is mutable and a version is not.  The returned identity
    always carries the exact version and manifest hash; ``stage`` only records
    what was asked for.

    Applicability
    -------------
    Machine-independent.

    Provenance
    ----------
    .. [vaft669] VEST-Tokamak/vaft#669, sections 9 and 13: lifecycle aliases
       resolve to immutable artifacts, and resolution verifies the hash.
    """
    env = os.environ if env is None else env
    _, version, entry, _, manifest, manifest_sha = _registry_entry(name, version, stage, registry, env)
    directory = _cache_root(cache, env) / name / version
    problems = _cache_problems(directory, manifest)
    if problems:
        raise ModelResolutionError(
            f"cached artifact {directory} is not {name} {version}: " + "; ".join(problems)
            + f". Fetch it with vaft.process.ml.fetch_model({name!r}, version={version!r}), or: "
            + f"gh release download {entry.get('release') or f'{name}-v{version}'} -R {DEFAULT_REPOSITORY} -D {directory}"
        )
    return _identity(name, version, stage, entry, manifest, manifest_sha, directory)

def save_dataset(dataset, directory):
    """Materialise a dataset on disk so training can run elsewhere.

    Building a corpus-scale dataset (reading thousands of shots) and training
    on it are separate steps, often on separate machines; this is the hand-off.

    Parameters
    ----------
    dataset : DatasetArtifact
        Fingerprinted dataset from :func:`build_dataset` [-].
    directory : str or path-like
        Target directory; must not exist or be empty [-].

    Returns
    -------
    str
        The dataset fingerprint, which :func:`load_dataset` re-checks [-].

    Raises
    ------
    ModelContractError
        When the dataset has no fingerprint, the directory is not empty, or an
        array has dtype ``object``.

    Applicability
    -------------
    Machine-independent.
    """
    if not dataset.fingerprint:
        raise ModelContractError("dataset has no fingerprint; build it with build_dataset")
    directory = Path(directory)
    if directory.exists() and any(directory.iterdir()):
        raise ModelContractError(f"{directory} is not empty")
    directory.mkdir(parents=True, exist_ok=True)
    arrays = {"inputs": np.asarray(dataset.inputs), "groups": dataset.group_labels.astype(str)}
    if dataset.targets is not None:
        arrays["targets"] = np.asarray(dataset.targets)
    if dataset.available is not None:
        arrays["available"] = np.asarray(dataset.available)
    for key, value in dataset.sample_meta.items():
        arrays[f"meta:{key}"] = np.asarray(value)
    for key, value in arrays.items():
        if value.dtype == object:
            raise ModelContractError(f"{key} has dtype object and cannot be stored without pickle")
    np.savez(directory / _DATASET_ARRAYS, **arrays)
    record = {
        "fingerprint": dataset.fingerprint,
        "spec": dataset.spec.to_dict(),
        "feature_spec": dataset.feature_spec.to_dict() if isinstance(dataset, FeatureDataset) else None,
        "group_dtype": str(np.asarray(dataset.groups).dtype),
    }
    (directory / _DATASET_RECORD).write_text(json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return dataset.fingerprint

def save_model_artifact(model, directory, *, version, model_card=None, applicability=None):
    """Write a model as an immutable, hash-pinned artifact directory.

    The directory is what gets attached to a ``vaft-nn`` release; its
    ``manifest.json`` is what gets committed to the registry.

    Parameters
    ----------
    model : ModelArtifact
        Trained model; ``model.name`` must be lower_snake_case [-].
    directory : str or path-like
        Target directory; must not exist or be empty [-].
    version : str
        Semantic version of this artifact, e.g. ``"0.1.0"`` [-].
    model_card : mapping or None, optional
        Intended use, training population, validation evidence and known
        limits, stored in the manifest [-].
    applicability : mapping or None, optional
        The input domain the model was trained and validated on (parameter
        ranges, shot ranges, diagnostics), stored in the manifest [-].

    Returns
    -------
    ModelArtifact
        The same model carrying its saved identity: version, manifest SHA-256
        and weight SHA-256 [-].

    Raises
    ------
    ModelContractError
        On an invalid name or version, or a non-empty target directory.

    Processing steps
    ----------------
    1. Write the backend weights and the normalisation moments.
    2. Hash each file and write ``manifest.json`` (sorted keys, two-space
       indent) pinning every hash and size.
    3. Hash the manifest; that hash is the artifact's identity.

    Applicability
    -------------
    Machine-independent.

    Provenance
    ----------
    .. [vaft669] VEST-Tokamak/vaft#669, sections 10-11: the ``vaft-nn`` release
       layout and manifest contract.
    """
    _check_name_version(model.name, version)
    directory = _empty_directory(directory)
    raw = _write_artifact(model, directory, str(version), "", model_card, applicability)
    identity = ModelIdentity(
        name=model.name,
        state_sha256=model.state_sha256,
        version=str(version),
        manifest_sha256=sha256_bytes(raw),
        resolved_by="saved",
        source=str(directory),
    )
    return model.replace(identity=identity)