Sweeps

{"sweep": [a, b, c]} at a field means the document is three points: the same experiment three times, differing in that one field. A swept document compiles to a root plan whose children are one plan per point, and a plan's path in the tree is its path on disk.

A sweep is lowered on the raw JSON, before anything is compiled. The wrapper is replaced by each of its values in turn and each resulting document is compiled on its own, so a point is an ordinary document in every way — its own addresses, its own tokenization, its own digest — and nothing downstream of the compiler knows a sweep happened.

At a glance #

{
  "header": {
    "protocol_version": "3",
    "description": "AUTHORED FOR causalab-mini, not shipped with causalab. minimal_cpu.json with the patched position swept: the same interchange at the last token, the one before it, and the one before that. One document, three experiments, one model load \u2014 which is what the protocol's sweep wrappers are for. Used to test that a plan tree needs nothing beyond a root plan with three children to say so."
  },
  "model": {
    "key": "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "revision": "9fb191250dd56d0ba7ec9785a025ed29c03d5998",
    "dtype": "fp32"
  },
  "data": {
    "base": {
      "dataset": "weekdays/train",
      "field": "input"
    },
    "counterfactual": {
      "dataset": "weekdays/train",
      "field": "counterfactual_inputs[0]"
    }
  },
  "method": {
    "sites": {
      "target": {
        "component": "block_output",
        "layers": [
          0
        ]
      },
      "lm_head": {
        "component": "lm_head"
      }
    },
    "reads": {
      "v_cf": {
        "site": "target",
        "pos": -1,
        "model": "original",
        "input": "counterfactual"
      },
      "logits": {
        "site": "lm_head",
        "pos": -1,
        "model": "patched",
        "input": "base"
      }
    },
    "writes": {
      "patch": {
        "site": "target",
        "pos": {
          "sweep": [
            -1,
            -2,
            -3
          ]
        },
        "do": {
          "swap": "v_cf"
        }
      }
    },
    "intervened_models": {
      "patched": {
        "input": "base",
        "writes": [
          "patch"
        ]
      }
    },
    "metrics": {
      "iia": {
        "kind": "match",
        "of": "logits",
        "expected": "cf_answer",
        "token_form": "space_prefixed"
      },
      "logit_diff": {
        "kind": "logit_diff",
        "of": "logits",
        "a": "cf_answer",
        "b": "base_answer",
        "token_form": "space_prefixed"
      }
    },
    "save": [
      {
        "value": "iia",
        "model": "patched",
        "input": "base",
        "file_path": "iia.json"
      },
      {
        "value": "logit_diff",
        "model": "patched",
        "input": "base",
        "file_path": "logit_diff.json"
      }
    ]
  }
}

One wrapper, at method.writes.patch.pos: the write's position is {"sweep": [-1, -2, -3]} where an unswept document would have an integer. Everything else — the read at -1, the swap, the two metrics — is the same in all three points.

$ uv run causalab-mini validate documents/pos_sweep_cpu.json
ok: documents/pos_sweep_cpu.json is a valid protocol document

$ uv run causalab-mini validate documents/real/weekdays_layer_sweep.json
ok: documents/real/weekdays_layer_sweep.json is a valid plan-shaped document

$ uv run causalab-mini explain documents/pos_sweep_cpu.json
<root>: Plan
  pos=-1: Plan
    observe: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
          read  'v_cf' at block_output[0] pos={index:-1} via 'identity'
        forward 'patched' on 'base'  (4 rows x 11 tokens)
          write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'
  pos=-2: Plan
    observe: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
          read  'v_cf' at block_output[0] pos={index:-1} via 'identity'
        forward 'patched' on 'base'  (4 rows x 11 tokens)
          write 'patch' at block_output[0] pos={index:-2} swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'
  pos=-3: Plan
    observe: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
          read  'v_cf' at block_output[0] pos={index:-1} via 'identity'
        forward 'patched' on 'base'  (4 rows x 11 tokens)
          write 'patch' at block_output[0] pos={index:-3} swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'

Three points, one model load. The read is at -1 in all three — only the write moves — and each point's pos={index:-2} is an ordinary fixed position by the time the compiler sees it.

The wrapper #

A wrapper is an object whose keys are sweep and, optionally, as — nothing else. That is how the lowering tells a wrapper from an ordinary object with a field called sweep, and it is why a wrapper holds one list and not several fields.

A range is [start, stop] or [start, stop, step], in integers, and means what Python's range means — half-open, so [0, 16] is a 16-layer model's every layer and not seventeen of them.

Where a point lands, and what it is called #

A point's name is the swept field's own name — the last string key on the path to the wrapper — and the value it took. That name is the plan's name in the tree and its directory on disk:

$ uv run causalab-mini run documents/pos_sweep_cpu.json --device-map cpu --out out/pos
out/b-pos/document.json
out/b-pos/run.json
out/b-pos/pos=-1/document.json
out/b-pos/pos=-1/iia.json
out/b-pos/pos=-1/logit_diff.json
out/b-pos/pos=-2/document.json
out/b-pos/pos=-2/iia.json
out/b-pos/pos=-2/logit_diff.json
out/b-pos/pos=-3/document.json
out/b-pos/pos=-3/iia.json
out/b-pos/pos=-3/logit_diff.json

out/pos/
├── document.json          the document as written, sweep and all
├── run.json               engine, versions, package digest, served dtype
├── pos=-1/
│   ├── document.json      this point's document, the sweep lowered out
│   ├── iia.json
│   └── logit_diff.json
├── pos=-2/
└── pos=-3/

The root carries the document you wrote and the run record; each point carries the document it became. Nothing on disk is named by position in a list, so adding a value to a sweep does not rename anyone else's results.

A one-element list collapses in the name #

Several of the things worth sweeping are lists that happen to hold one element: a layer band, a single head, one SAE latent. In a point's name a one-element list is written as its element; anything longer is written as compact JSON. Three documents, three shapes of value:

$ uv run causalab-mini explain documents/v2/head_patching.json | grep ': Plan'
<root>: Plan
  heads=0: Plan
  heads=1: Plan
  heads=2: Plan
  heads=3: Plan

$ uv run causalab-mini explain documents/v2/neuron_patching.json | grep ': Plan'
<root>: Plan
  units=[0,1,2,3]: Plan
  units=[4,5,6,7]: Plan
  units=[8,9,10,11]: Plan

$ uv run causalab-mini explain documents/real/weekdays_layer_sweep.json | grep ': Plan'
<root>: Plan
  layers=0,pos=-4: Plan
  layers=0,pos=-1: Plan
…

heads: {"sweep": [[0], [1], [2], [3]]} gives heads=0, because a site's heads field is a list and each point wants one of them. units sweeps lists of four and keeps the list. Only the name collapses — the document at each point holds the list it was given.

Two axes #

Several swept fields are their cross product, one point per combination, labelled in the order the fields appear in the document. documents/real/weekdays_layer_sweep.json has two: a range over the 16 layers of a site, and a two-value sweep of the position — which is why its points are called layers=12,pos=-1.

Three wrappers are in that document, in three places. The site's layers holds {"sweep": {"range": [0, 16]}}; the read v_cf and the write patch each hold {"sweep": [-4, -1], "as": "pos"} as their position. The whole document is on Layer sweep.

Every field on a named axis needs a value per point, and a mismatch is refused naming both places:

$ uv run causalab-mini validate swept_mismatch.json
Traceback (most recent call last):
  …
causalab_mini.plan.sweep.SweepError: interventions.patching.writes.patch.pos: axis 'pos' has 2 values at interventions.patching.reads.v_cf.pos and 3 here; fields that move together need a value each per point

What 32 points look like #

Mean over the 42 rows, one cell per point. Two metrics at two positions, 16 layers deep:

What that says about the model is on Layer sweep. What it says about sweeps is that a grid of experiments is a grid of directories, and the directory names are the coordinates.

Asking a swept run for a result #

Result names come from the document, so they are unique within a plan and repeated across a sweep's points. Asking the root for one is an error, not a first match:

>>> root.result("iia")
causalab_mini.plan.plan.PlanError: result 'iia': 3 in this plan. Results here: ['iia', 'logit_diff']

>>> float(root.steps["pos=-2"].result("iia").mean())
0.0

Naming the point is what disambiguates, and the point's name is the same string as its directory. all_results() refuses on a sweep for the same reason; ask a point instead.

What is refused #

The first one is the load-bearing refusal: the engine loads the model once and every point runs against it, so a point may not ask for a different one. That is also why the CLI reads the first point to find out what to load.

Below that, Spec and Document each describe one point, so either of them meeting a wrapper means the lowering did not run. Both front ends check for it rather than letting a {"sweep": …} object be validated as whatever field it is sitting on.

The protocol allows more here than mini implements — at_once, cohorts, swept parameter bundles. Those are refused by name rather than loaded and ignored.