Tracking¶

tinax.tracking wraps Weights & Biases run tracking. Constructing WandbRun is the explicit start of the lifecycle: it imports wandb on demand, calls wandb.init, and stores the returned run. log and log_artifact require an unfinished run, finish is idempotent, and the context manager finishes the run on exit. Importing tinax.tracking imports neither wandb nor JAX.

from tinax.tracking import WandbRun

with WandbRun(project="demo", name="baseline", config={"learning_rate": 0.1}) as run:
    run.log({"loss": 0.5}, step=0)

WandbRun ¶

WandbRun(
    *,
    project: str,
    name: str | None = None,
    config: Mapping[str, Any] | None = None,
    tags: Sequence[str] | None = None,
    group: str | None = None,
)

An explicitly started Weights & Biases run with explicit finish semantics.

Constructing the handle starts the run by calling wandb.init and storing the returned run object; finish or leaving the context manager ends it. log and log_artifact require an unfinished run. The wandb package is imported only when a handle is constructed.

Parameters:

Name	Type	Description	Default
`project`	`str`	Weights & Biases project name. Must be non-empty.	required
`name`	`str \| None`	Optional run name.	`None`
`config`	`Mapping[str, Any] \| None`	Optional run configuration mapping with string keys.	`None`
`tags`	`Sequence[str] \| None`	Optional sequence of non-empty tag strings.	`None`
`group`	`str \| None`	Optional group name for related runs.	`None`

Raises:

Type	Description
`TypeError`	If a string argument is not a string, `config` is not a mapping with string keys, or `tags` is not a sequence of strings.
`ValueError`	If `project` is empty or a tag is empty.

log ¶

log(
    metrics: Mapping[str, float], *, step: int | None = None
) -> None

Log one mapping of scalar metrics to the started run.

Parameters:

Name	Type	Description	Default
`metrics`	`Mapping[str, float]`	Non-empty mapping of metric names to finite `int` or `float` values (booleans are rejected).	required
`step`	`int \| None`	Optional non-negative step associated with the metrics.	`None`

Raises:

Type	Description
`TypeError`	If `metrics` is not a mapping with string keys and numeric values, or `step` is not an integer (booleans rejected).
`ValueError`	If `metrics` is empty, a metric value is not finite, or `step` is negative.
`RuntimeError`	If the run was already finished.

log_artifact ¶

log_artifact(
    path: str | PathLike[str],
    *,
    name: str,
    artifact_type: str = "model",
    metadata: Mapping[str, Any] | None = None,
    aliases: Sequence[str] | None = None,
) -> None

Log one existing local file to the started run as a Weights & Biases artifact.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Path to an existing local file.	required
`name`	`str`	Non-empty artifact name.	required
`artifact_type`	`str`	Non-empty artifact type label.	`'model'`
`metadata`	`Mapping[str, Any] \| None`	Optional mapping stored as artifact metadata.	`None`
`aliases`	`Sequence[str] \| None`	Optional sequence of non-empty alias strings.	`None`

Raises:

Type	Description
`TypeError`	If `path` is not a string or `os.PathLike`, `name` or `artifact_type` is not a string, `metadata` is not a mapping or `None`, or an alias is not a string.
`ValueError`	If `name` or `artifact_type` is empty, or an alias is empty.
`FileNotFoundError`	If no file exists at `path`.
`RuntimeError`	If the run was already finished.

finish ¶

finish() -> None

Finish the run and release its resources; subsequent calls are no-ops.