netket.logging.WandbLog

class netket.logging.WandbLog

Bases: AbstractCallback

Logger that streams metrics and optional model checkpoints to Weights & Biases.

On the first call the logger lazily starts a wandb run so that constructing the logger does not create a run if it is never used.

The wandb package must be installed (pip install wandb).

Nested metric keys (e.g. Energy/Mean) are logged with / preserved, which wandb’s UI renders as collapsible groups. Complex-valued scalars are split into separate <key>/re and <key>/im entries.

NetKet loggers can be called more than once with the same step (for different keys), while wandb expects a single, monotonically increasing global step. For this reason, every top-level namespace (the first path segment of a key, e.g. Energy) gets its own hidden <namespace>/step metric, declared with wandb.define_metric(), and its metrics are plotted against it.

This class is a full AbstractCallback and can be passed either as out=logger or as callback=logger. When used as a callback the logger automatically captures the variational state snapshot taken just before the parameter update.

Parameters:

• project (str | None) – Name of the wandb project. If None wandb uses its own default (the WANDB_PROJECT environment variable, or "uncategorized").

• entity (str | None) – wandb entity (user or team) that owns the run.

• run_name (str | None) – Human-readable label attached to the new run.

• run_id (str | None) – If provided, the run with this id is resumed if it exists, or created with this id otherwise. Useful to resume an interrupted job in place.

• group (str | None) – Group name used to cluster related runs in the wandb UI.

• job_type (str | None) – Free-form job type of the run (e.g. "vmc"), useful to filter the runs of a group.

• tags (Collection[str] | None) – Optional collection of string tags to attach to the run.

• mode (str | None) – wandb mode, one of "online", "offline" or "disabled". If None, wandb’s default is used (which can be set with the WANDB_MODE environment variable).

• save_params (bool) – If True, periodically serialize model parameters as a versioned wandb artifact (MessagePack binary, same format as JsonLog).

• save_params_every (int) – Save parameters every this many optimisation steps. Only relevant when save_params=True.

• save_final_state (bool) – If True, the variational state at the end of the run is uploaded as a separate artifact.

• log_histograms (bool) – If True, log a histogram of every parameter every time the parameters are saved.

• metadata (dict | None) – Optional dict of hyper-parameters stored in the wandb config of the run.

• ignore (Collection[str] | None) – Optional collection of top-level log-dict keys to skip entirely. Useful for suppressing entries that cannot be serialised as wandb metrics (e.g. large arrays or custom objects).

Tip

Use metadata to attach a dict of hyper-parameters (learning rate, system size, model type, …) to the run. They are stored in the wandb config of the run, making it easy to filter and compare runs in the wandb UI without external bookkeeping.

Examples

Log an optimisation run to wandb.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> logger = nk.logging.WandbLog(
...     project="Ising1d",
...     run_name="RBM_alpha1",
...     tags=["RBM"],
...     metadata={"learning_rate": 0.01, "alpha": 1, "L": 20},
... )
>>> gs.run(n_iter=500, out=logger)

Log offline, to synchronise later with wandb sync.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> logger = nk.logging.WandbLog(project="Ising1d", mode="offline")
>>> gs.run(n_iter=500, out=logger)

Resume a previous run.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> logger = nk.logging.WandbLog(project="Ising1d", run_id="<run-id>")
>>> gs.run(n_iter=500, out=logger)

Using the logger as a callback.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> logger = nk.logging.WandbLog(project="Ising1d")
>>> gs.run(n_iter=500, callback=logger)

Inheritance

Attributes

callback_order

Methods

__call__(step, item, variational_state=None)

Call self as a function.

Parameters:

• step (int)

• item (dict[str, Any])

• variational_state (VariationalState | None)

before_parameter_update(step, log_data, driver)

Called after all update logic has been computed and the step has been accepted, but before the driver applies the parameter update.

At this point:

• The loss and its gradient have been computed by compute_loss_and_update().

• The step has been accepted (not rejected by on_compute_update_end()).

• driver.step_count still refers to the current step — it has not yet been incremented.

• The variational state parameters have not yet changed.

This is the right place to estimate additional observables, add data to log_data, or take a snapshot of the state for logging. Callbacks with a lower callback_order run first, so observables callbacks (order 0) are guaranteed to populate log_data before logger callbacks (order 10) read it.

flush(variational_state=None)

Flushes pending data and optionally saves model parameters.

Parameters:

variational_state (VariationalState | None) – if provided and save_params=True, the current model parameters are uploaded as an artifact.

on_compute_update_end(step, log_data, driver)

Callback called at the end of the compute update phase, after computing the loss and its gradient.

This is called before the parameters are updated, so it can be used to implement custom logic for rejecting a step based on the computed loss or gradient.

Return type:

bool

Returns:

A boolean indicating whether to reject the step (i.e. repeat it with the same parameters). If it returns None, it is treated as False.

on_compute_update_start(step, log_data, driver)

on_run_end(step, driver)

on_run_error(step, error, driver)

on_run_start(step, driver)

on_step_end(step, log_data, driver)

on_step_start(step, log_data, driver)

replace(**kwargs)

Replace the values of the fields of the object with the values of the keyword arguments. If the object is a dataclass, dataclasses.replace will be used. Otherwise, a new object will be created with the same type as the original object.

Return type:

TypeVar(P, bound= Pytree)

Parameters:

• self (P)

• kwargs (Any)