remix_d4_part2_solution.py

# %%
# This document serves multiple purposes:
#   - A reproduction of the experiments in the Paren Balancer causal scrubbing post
#   - A causal scrubbing demo
#   - A set of exercises for remix (hence the `if "SOLUTION"` blocks, and instructions)

# %% [markdown]
"""
# Day 4b: Paren Balancer Causal Scrubbing

To start, please read this [LessWrong post](https://www.lesswrong.com/s/h95ayYYwMebGEYN5y/p/kjudfaQazMmC74SbF).

We will replicate the experiments today.

<!-- toc -->

## Learning Objectives

After today's material, you should be able to:

- Execute causal scrubbing experiments using the `causal_scrubbing` package.
- Apply more complex rewrites in Circuits.

Note: unlike Circuits which is in Rust, `causal_scrubbing` is in Python so feel free to use Go To Definition and look at the implementation to clarify what's going on.
"""
import os
import sys
import uuid
from typing import Optional, Tuple

import rust_circuit.optional as op
import numpy as np
import rust_circuit as rc
import torch
from rust_circuit.causal_scrubbing.experiment import (
    Experiment,
    ExperimentCheck,
    ExperimentEvalSettings,
    ScrubbedExperiment,
)
from rust_circuit.causal_scrubbing.hypothesis import (
    Correspondence,
    CondSampler,
    ExactSampler,
    FuncSampler,
    InterpNode,
    UncondSampler,
    chain_excluding,
    corr_root_matcher,
)
from rust_circuit.algebric_rewrite import (
    residual_rewrite,
    split_to_concat,
)
from rust_circuit.model_rewrites import To, configure_transformer
from rust_circuit.py_utils import I
from torch.nn.functional import binary_cross_entropy_with_logits

import remix_d4_part2_test as tests
from remix_d4_part2_setup import ParenDataset, ParenTokenizer, get_h00_open_vector
import remix_utils

MAIN = __name__ == "__main__"

if "SKIP":
    # CI takes longer than 20s timeout right now, just skip CI
    IS_CI = os.getenv("IS_CI")
    if IS_CI:
        sys.exit(0)

# %%
# define model
SEQ_LEN = 42
NUM_EXAMPLES = 4000
MODEL_ID = "jun9_paren_balancer"

PRINT_CIRCUITS = True
ACTUALLY_RUN = True
SLOW_EXPERIMENTS = True
DEFAULT_CHECKS: ExperimentCheck = True
#DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
DEVICE = "cpu"
EVAL_DEVICE = "cuda:0"
print("Preparation on device:", DEVICE, "and evaluation on device:", EVAL_DEVICE)
# If you have less memory, you will want to reduce this and also add a batch_size
MAX_MEMORY = 70_000_000_000
BATCH_SIZE = 2000

# %% [markdown]
"""
## Setup 
No exercises here! It may be helpful to read over the code, however.

### Circuit loading
If any of these operations confuse you, try printing out the circuit before and after!

Step 1: Initial loading
"""

circ_dict, _, model_info = remix_utils.load_paren_balancer()
circuit = circ_dict["t.bind_w"]

# %% [markdown]
"""
Step 2: We bind the model to an input by attaching a placeholder symbol input named "tokens" to the model. We then specify the attention mask that prevents attending to padding depends on this tokens array.

We use one hot tokens as this makes defining the attention mask a simple indexing operation.

The symbol has a random fixed uuid (fixed as this gives consistency when comparing in tests).
"""
toks_uuid = uuid.UUID("ce34280e-169f-40bd-b78e-8adeb4274aba")
tokens_arr = rc.Symbol((SEQ_LEN, ParenTokenizer.vocab_size), uuid=toks_uuid, name="tokens")
tok_embeds = rc.Einsum.from_fancy_string(
    "seqlen vocab_size, vocab_size hidden -> seqlen hidden",
    tokens_arr,
    circ_dict["t.w.tok_embeds"],
    name=f"tok_embeds",
)
attn_mask = rc.Add.minus(
    rc.Scalar(1),
    rc.Index(tokens_arr, I[:, ParenTokenizer.PAD_TOKEN]),
    name="pos_mask",
)
circuit = model_info.bind_to_input(circuit, tok_embeds, circ_dict["t.w.pos_embeds"], attn_mask)

# [markdown]
"""
Step 3: rewrites the circuit into a more convenient structure to work with using `configure_transformer`.
This flattens out the residual stream (as opposed to the nested layer structure originally) and, pushes down the weight bindings, and separates out each attention layer into a sum of heads.
"""

circuit = circuit.update(
    "t.bind_w",
    lambda c: configure_transformer(
        c,
        To.ATTN_HEAD_MLP_NORM,
        split_by_head_config="full",
        use_pull_up_head_split=True,
        use_flatten_res=True,
        flatten_components=True,
    ),
)

# [markdown]
"""
Some additional misc rewrites.

We substitute the inputs to be duplicated everywhere they apperar in the model instead of being in one outer module bind.

We also index as we only care about the classification at position 0, and use `rc.conform_all_modules` to replace any remaining symbolic shapes with their numeric values.
"""

circuit = circuit.cast_module().substitute()
circuit = rc.Index(circuit, I[0]).rename("logits_pos0")
circuit = rc.conform_all_modules(circuit)

# %% [markdown]
"""
Finally, some custom renames that make the circuit more intuitive.
"""
# TBD I think the next line is meant to be the following
# ```
# circuit = circuit.update("t.logits", lambda c: c.rename("logits"))
# ```
# but changing that may break things below because the head of the circuit will then become "logits_with_bias".
circuit = circuit.update("t.call", lambda c: c.rename("logits"))
circuit = circuit.update("t.call", lambda c: c.rename("logits_with_bias"))
circuit = circuit.update(rc.Regex(r"[am]\d(.h\d)?$"), lambda c: c.rename(c.name + ".inner"))
circuit = circuit.update("t.inp_tok_pos", lambda c: c.rename("embeds"))
circuit = circuit.update("t.a.mask", lambda c: c.rename("padding_mask"))

for l in range(model_info.params.num_layers):
    circuit = circuit.update(f"b{l}.m", lambda c: c.rename(f"m{l}"))
    circuit = circuit.update(f"b{l}.a.h0", lambda c: c.rename(f"a{l}.h0"))
    circuit = circuit.update(f"b{l}.a.h1", lambda c: c.rename(f"a{l}.h1"))

    next = "final" if l == model_info.params.num_layers - 1 else f"a{l+1}"
    circuit = circuit.update(f"b{l}", lambda c: c.rename(f"{next}.input"))

printer = rc.PrintHtmlOptions(
    shape_only_when_necessary=False,
    traversal=rc.restrict(
        rc.IterativeMatcher("embeds", "padding_mask", "final.norm", rc.Regex(r"^[am]\d(.h\d)?$")),
        term_if_matches=True,
    ),
)
# %%

if PRINT_CIRCUITS:
    printer.print(circuit)

circuit = rc.cast_circuit(circuit, rc.TorchDeviceDtypeOp(device=DEVICE))

# %% [markdown]
"""
### dataset and experiment code
We have a custom dataset class that precomputes some features of paren sequences, and handles pretty printing / etc.
"""

ds = ParenDataset.load(device=DEVICE)


def bce_with_logits_loss(logits: torch.Tensor, labels: torch.Tensor):
    """
    Computes the binary cross entropy loss for the provided labels.
    logits: [batch, 2]. Class 0 is unbalanced logit, class 1 is balanced logit.
    labels: [batch]. True if balanced.
    """
    targets = labels.to(dtype=logits.dtype, device=logits.device)
    logit_diff = logits[..., 1] - logits[..., 0]
    correct = (logit_diff > 0) == targets
    loss = binary_cross_entropy_with_logits(logit_diff, targets, reduction="none")
    return loss, correct


def paren_experiment(
    circuit: rc.Circuit,
    dataset: ParenDataset,
    corr: Correspondence,
    checks: ExperimentCheck = DEFAULT_CHECKS,
    random_seed=1,
    actually_run=ACTUALLY_RUN,
    num_examples=NUM_EXAMPLES,
    batch_size=BATCH_SIZE,
    **kwargs,
) -> Tuple[ScrubbedExperiment, Optional[float]]:
    ex = Experiment(
        circuit,
        dataset,
        corr,
        random_seed=random_seed,
        check=checks,
        **kwargs,
    )
    scrubbed = ex.scrub(num_examples, treeify=actually_run)
    overall_loss: Optional[float] = None
    if actually_run:
        logits = scrubbed.evaluate(
            ExperimentEvalSettings(
                optim_settings=rc.OptimizationSettings(max_memory=MAX_MEMORY, scheduling_naive=True),
                device_dtype=EVAL_DEVICE,
                optimize=True,
                batch_size=batch_size,
            ),
        )
        ref_ds = ParenDataset.unwrap(scrubbed.ref_ds)
        labels = ref_ds.is_balanced.value

        def loss_str(mask):
            loss, correct = bce_with_logits_loss(logits[mask], labels[mask])
            loss = loss.cpu()
            std_err = loss.std() / len(loss) ** 0.5

            return f"{loss.mean():.3f}  SE={std_err:.3f}  acc={correct.float().mean():.1%} "

        print(f"  overall:               {loss_str(slice(None))}")
        print(f"    on bal:              {loss_str(labels.to(dtype=torch.bool))}")
        print(f"    on unbal:            {loss_str(~labels.to(dtype=torch.bool))}")
        print(f"    on count failures:   {loss_str(~ref_ds.count_test.to(dtype=torch.bool))}")  # type: ignore
        print(f"    on horizon failures: {loss_str(~ref_ds.horizon_test.to(dtype=torch.bool))}")  # type: ignore
        overall_loss = bce_with_logits_loss(logits, labels)[0].mean().item()

    return scrubbed, overall_loss


def check_loss(loss: Optional[float], target: float, std_err: float):
    # we usually set tol to be 4 * SE
    assert loss is not None
    err = abs(loss - target)
    assert err < 4 * std_err, f"err too large! loss ({loss:.2f}) != target ({target:.2f}) ± 4*SE ({std_err:.2f})"
    if err > 2 * std_err:
        raise Warning(f"Err is kinda large! loss ({loss:.2f}) != target ({target:.2f}) ± 2*SE ({std_err:.2f})")


# %% [markdown]
"""
### Helpful tidbit on tests:
Most of the tests in this file raise assertion errors that contain extra data, for instance the objects from the comparison that failed. It can be convenient to catch this data to debug. For instance:

```
def check_eq(a, b):
    assert a == b, ("not equal!", a, b)

try:
    check_eq(0, 1)
except AssertionError as e:
    a, b = e.args[0][1], e.args[0][2]
    print(a, b)
```
"""


# %% [markdown]
"""
## Experiment 0: Baselines
To start with let's measure two baselines:
  - running the model normally
  - interchanging the logits randomly

Make causal scrubbing experiments that impliment both of these. In each case there should be a single interp node named "logits".

The tests do explictly check that the interp nodes in the correspondence are named correctly, in order to facilitate more helpful feedback.
"""

if "SOLUTION":
    corr0a = Correspondence()
    corr0a.add(InterpNode(cond_sampler=ExactSampler(), name="logits"), corr_root_matcher)
else:
    corr0a = Correspondence()
    """TODO: Your code here"""
if MAIN:
    tests.t_ex0a_corr(corr0a)
    print("\nEx0a: Exact sampling")
    ex0a, loss0a = paren_experiment(circuit, ds, corr0a)
    check_loss(loss0a, 0, 0.01)

if "SOLUTION":
    corr0b = Correspondence()
    corr0b.add(InterpNode(cond_sampler=UncondSampler(), name="logits"), corr_root_matcher)
else:
    corr0b = Correspondence()
    """TODO: Your code here"""
if MAIN:
    tests.t_ex0b_corr(corr0b)
    print("\nEx0b: Interchanging logits")
    ex0b, loss0b = paren_experiment(circuit, ds, corr0b)
    check_loss(loss0b, 4.30, 0.12)

# %% [markdown]
"""
## Experiment 1: The direct contributions to the output
Now, let's construct a basic experiment to determine the role that different heads play.
We'll start by testing the following claimed hypothesis:
 - Heads 1.0 and 2.0 compute the count test, and check that there are equal numbers of open and close parentheses
 - Head 2.1 computes the horizon test.

### Matchers

Define the following matchers. You only want to match _direct_ paths, that is paths through the residual stream and not through indirect paths. This can be accomplished by calling `rc.restrict` or the `chain_excluding` utilty included in causal scrubbing code.
    
"""
if "SOLUTION":
    # There are several ways to define this matcher, eg:
    # m_10 = rc.IterativeMatcher("final.input").chain(rc.restrict("a1.h0", end_depth=2))
    # or defining all_components = {"a0.h0", "a0.h1", "m0", ..., "m2"} and then
    # m_10 = chain_excluding(corr_root_matcher, "a1.h0", all_components - "a1.h0")
    # It's just down to personal preference / what is clearer in the circumstance.
    m_10 = chain_excluding(corr_root_matcher, "a1.h0", {"m2", "m1", "a2.h0", "a2.h1"})
    m_20 = chain_excluding(corr_root_matcher, "a2.h0", "m2")
    m_21 = chain_excluding(corr_root_matcher, "a2.h1", "m2")
else:
    m_10: rc.IterativeMatcher
    m_20: rc.IterativeMatcher
    m_21: rc.IterativeMatcher
    """TODO: Your code here"""

if MAIN:
    tests.t_m_10(m_10)
    tests.t_m_20(m_20)
    tests.t_m_21(m_21)

# %%  [markdown]
"""
### Cond Samplers 

Let's now define some samplers that sample according to whether or not a datum passes the count test and the horizon test. We first need to define some helper functions, then wrap them in `FuncSampler`s. You will need to use `ParenDataset.tokens_flat`.

Note: these are also predefined properties on `ParenDataset`. If you are short on time, feel free to use them and skip this exercise. The version inside `ParenDataset` also uses caching to make things faster, so if speed gets to be annoyingly slow for later experiments you could switch over).

<details>
<summary>
Using functions that come with `ParenDataset`
</summary>

You can use the utilities in `ParenDataset` by changing your code to the following:
```
count_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).count_test)

horizon_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).horizon_test)
```
</details>
"""


# %%
def passes_count(d: ParenDataset) -> torch.Tensor:
    """
    Returns a bool tensor of shape [len(d)]
    Result is true when the corresponding datum has equal numbers of open and close parens
    """
    # not used in solution, as implemented in ParenDataset
    """TODO: Your code here"""
    raise NotImplementedError


def passes_horizon(d: ParenDataset) -> torch.Tensor:
    """
    Returns a bool tensor of shape [len(d)]
    Result is true when the corresponding datum passes the right to left horizon test as described in the [writeup](https://www.lesswrong.com/s/h95ayYYwMebGEYN5y/p/kjudfaQazMmC74SbF#Algorithm).
    """
    # not used in solution, as implemented in ParenDataset
    """TODO: Your code here"""
    raise NotImplementedError


if "SOLUTION":
    count_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).count_test)
    horizon_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).horizon_test)
else:
    count_cond = FuncSampler(lambda d: passes_count(ParenDataset.unwrap(d)))
    horizon_cond = FuncSampler(lambda d: passes_horizon(ParenDataset.unwrap(d)))

if MAIN:
    tests.t_count_cond(count_cond)
    tests.t_horizon_cond(horizon_cond)

# %% [markdown]
"""
This first correspondence should have 4 nodes:
 - The root node, named "logits", with an ExactSampler (any sampler that agrees on the labels will be equivalent,
 but an exact sampler is somewhat more computationally efficient).
 - Three nodes for the three heads of interest, named "10", "20", and "21". The first two should use the cond sampler provided (`cs_for_h10_and_h20`, which will be `count_cond` for now), the third should use the `horizon_cond` sampler.

The exact interp node names are checked by the test, which allows it to give more meaningful feedback. 
"""


def make_ex1_corr(cs_for_h10_and_h20: CondSampler) -> Correspondence:
    """SOLUTION"""
    corr = Correspondence()
    i_logits = InterpNode(name="logits", cond_sampler=ExactSampler())
    corr.add(i_logits, corr_root_matcher)
    corr.add(i_logits.make_descendant(cs_for_h10_and_h20, "10"), m_10)
    corr.add(i_logits.make_descendant(cs_for_h10_and_h20, "20"), m_20)
    corr.add(i_logits.make_descendant(horizon_cond, "21"), m_21)
    return corr


# %%
if MAIN:
    print("\nEx1a: Just the count cond")
    tests.t_ex1_corr(make_ex1_corr, count_cond)
    ex1a, loss1a = paren_experiment(circuit, ds, make_ex1_corr(count_cond))
    check_loss(loss1a, 0.52, 0.04)

    if PRINT_CIRCUITS:
        ex1a.print()


# %% [markdown]
"""
As discussed in the writeup, we can more accurately capture the equivalence classes of 1.0 and 2.0's output by including if the first parenthesis is open or closed.

This is a natural feature for these heads to use: a sequence will always be unbalanced if it starts with a close parenthesis, and as these heads depend strongly on the residual stream at position 1 anyhow (as we will show in experiement 2) the information is readily accessible.

(reminder: position 0 is always the [START] token, so position 1 is the first parentheses. All sequences in our dataset have >= 2 parentheses in them, so you can assume position 1 is either an open or close paren.)

Define some new cond samplers that incorporate this feature. The first one checks that the first paren is an open parentheses, and the next one tests the input passes count test AND the first paren is open.
We won't use the pure `start_open` test quite yet, but we will soon and it's nice to define it here.

Again, click on the hint below if you are short on time and just want to use the versions native to the `ParenDataset` class.

<details>
<summary>
Using functions that come with `ParenDataset`
</summary>

You can use the utilities in `ParenDataset` by changing your code to the following:
```
start_open_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).starts_with_open)

count_open_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).count_test & ParenDataset.unwrap(d).starts_with_open)
```
</details>
"""


def passes_starts_open(d: ParenDataset) -> torch.Tensor:
    """
    Returns a bool tensor of shape [len(d)].
    Result is true when the corresponding datum starts with '('.
    """
    """TODO: Your code here"""
    raise NotImplementedError


def passes_count_open(d: ParenDataset) -> torch.Tensor:
    """
    Returns a bool tensor of shape [len(d)].
    Result is true when the corresponding datum starts with '(' and there are equal numbers of open and close parens in the entire sequence.
    """
    """TODO: Your code here"""
    raise NotImplementedError


if "SOLUTION":
    start_open_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).starts_with_open)
    count_open_cond = FuncSampler(lambda d: ParenDataset.unwrap(d).count_test & ParenDataset.unwrap(d).starts_with_open)
else:
    start_open_cond = FuncSampler(lambda d: passes_starts_open(ParenDataset.unwrap(d)))
    count_open_cond = FuncSampler(lambda d: passes_count_open(ParenDataset.unwrap(d)))

if MAIN:
    tests.t_start_open_cond(start_open_cond)
    tests.t_count_open_cond(count_open_cond)


# %%
if MAIN:
    print("\nEx1b: Without a0")
    tests.t_ex1_corr(make_ex1_corr, count_open_cond)
    ex1b, loss1b = paren_experiment(circuit, ds, make_ex1_corr(count_open_cond))
    check_loss(loss1b, 0.30, 0.04)

    if PRINT_CIRCUITS:
        ex1b.print()

# [markdown]
"""
Bonus: Can you improve on the loss by specifying other direct paths, or choosing better features to ensure
agreement along?
"""
# %%
if "SOLUTION":
    if MAIN:
        print("\nEx1 bonus 1: With a0")
        corr = make_ex1_corr(count_open_cond)

        i_00 = corr.get_by_name("logits").make_descendant(count_open_cond, "00")
        m_00 = chain_excluding(
            corr_root_matcher,
            "a0.h0",
            {"m2", "m1", "m0", "a2.h0", "a2.h1", "a1.h0", "a1.h1"},
        )

        corr.add(i_00, m_00)
        ex1_bonus1 = paren_experiment(circuit, ds, corr)


# %% [markdown]
r"""
## Experiment 2: Diving into heads 1.0 and 2.0
We are going split up experiment 2 into four parts:
 - Part 1: 1.0 and 2.0 only depend on their input at position 1
 - Part 2 (ex2a in writeup): 1.0 and 2.0 only depend on:
    - the output of 0.0 (which computes $p$, the proportion of open parentheses) and
    - the embeds (which encode if the first paren is open)  
 - Part 3: Projecting the output of 0.0 onto a single direction
 - Part 4 (ex2b in writeup): Estimate the output of 0.0 with a function $\phi(p)$

### Part 1: Splitting up the input to 1.0 and 2.0 by sequence position
One of the claims we'd like to test is that only the input at position 1 (the first paren position) matters for both heads 1.0 and 2.0.
Currently, however, there is no node of our circuit corresponding to "the input at position 1". Let's change that!

Write a `separate_pos1` function that will transform a circuit `node` into:
```
'node_concat' Concat
  'node.at_pos0' Index [0:1, :]
    'node'
  'node.at_pos1' Index [1:2, :]
    'node'
  'node.at_pos2_42' Index [2:42, :]
    'node'
```
This can be acomplished by calling `split_to_concat()` (from algebraic_rewrites.py) and `.rename`-ing the result. If your test is not passing, take a look at your circuit and see how `split_to_concat()` names things by default.
"""


def separate_pos1(c: rc.Circuit) -> rc.Circuit:
    "SOLUTION"
    out = split_to_concat(
        c,
        0,
        [0, 1, torch.arange(2, 42, device=DEVICE)],
        partitioning_idx_names=["pos0", "pos1", "pos2_42"],
        use_dot=True,
    ).rename(f"{c.name}_concat")
    assert out.device == DEVICE
    return out


# TODO: add test for just this function

# %% [markdown]
"""
Then split the input node, but only along paths that are reached through head 2.0 and 1.0. (We don't want to split the input to 2.1 in particular, as we'll split that differently later.)

You'll probably want to do this with two different `.update` calls. One that splits the input to 2.0 and one that splits the input to 1.0.

If you are struggling to get the exact correct circuit, you are free to import it from the solution file and print it out. You can also try `print(rc.diff_circuits(your_circuit, our_circuit))` though.

Yes, the tests are very strict about naming things exactly correctly.
This is partially because it is convenient for tests, but also because names are really important!
Good names reduce confusion about what that random node of the circuit actually means.
Mis-naming nodes is also a frequent cause of bugs, e.g. a matcher that traverses a path that it wasn't supposed to.

The tests are also, unfortunately strict about the exact way that the idexes are specified.
This is because `rc.Index(arr, 3:8) != rc.Index(arr, t.arange(3,8))` even though they are functionally equivilant.
"""

ex2_part1_circuit = circuit
if "SOLUTION":
    ex2_part1_circuit = ex2_part1_circuit.update(rc.IterativeMatcher("a2.h0").chain("a2.input"), separate_pos1)
    ex2_part1_circuit = ex2_part1_circuit.update(rc.IterativeMatcher("a1.h0").chain("a1.input"), separate_pos1)

if MAIN and PRINT_CIRCUITS:
    subcirc = ex2_part1_circuit.get_unique(rc.IterativeMatcher("a2.h0").chain("a2.input_concat"))
    printer.print(subcirc)

if MAIN:
    tests.t_ex2_part1_circuit(ex2_part1_circuit)

# %% [markdown]
"""
Now we can test the claim that both 1.0 and 2.0 only cares about position 1!

We'll need new matchers, which just match the pos1 input.
"""

if "SOLUTION":
    m_10_p1 = m_10.chain("a1.input.at_pos1")
    m_20_p1 = m_20.chain("a2.input.at_pos1")
else:
    m_10_p1: rc.IterativeMatcher
    m_20_p1: rc.IterativeMatcher
    """TODO: Your code here"""

if MAIN:
    tests.t_m_10_p1(m_10_p1)
    tests.t_m_20_p1(m_20_p1)

# %% [markdown]
"""
Then create a correspondence that extends the one returned by `make_ex1_corr(count_open_cond)` so that both 1.0 and 2.0 only use information from position 1. `Correspondence.get_by_name` is useful here.

Have your new nodes be named "10_p1" and "20_p1".
"""


def make_ex2_part1_corr() -> Correspondence:
    "SOLUTION"
    corr = make_ex1_corr(count_open_cond)
    i_10_p1 = corr.get_by_name("10").make_descendant(name="10_p1", cond_sampler=count_open_cond)
    corr.add(i_10_p1, m_10_p1)
    i_20_p1 = corr.get_by_name("20").make_descendant(name="20_p1", cond_sampler=count_open_cond)
    corr.add(i_20_p1, m_20_p1)
    return corr


if MAIN:
    tests.t_make_ex2_part1_corr(make_ex2_part1_corr())
    print("\nEx 2 part 1: 1.0/2.0 depend on position 1 input")
    ex2_p1, loss2_p1 = paren_experiment(ex2_part1_circuit, ds, make_ex2_part1_corr())


# %% [markdown]
"""
### Part 2
We now construct experiment 2a from the writeup. We will be strict about where 1.0 and 2.0 learn the features they depend on. We claim that the 'count test' is determined by head 0.0 checking the exact proportion of open parens in the sequence and outputting this into the residual stream at position 1.

We thus need to also split up the output of attention head 0.0, so we can specify it only cares about the output of this head at position 1. Again, let's only split it for the branch of the circuit we are working with: copies of 0.0 that are upstream of either `m_10_p1` or `m_20_p1`. 
"""

ex2_part2_circuit = ex2_part1_circuit
if "SOLUTION":
    ex2_part2_circuit = ex2_part2_circuit.update(m_10_p1.chain("a0.h0"), separate_pos1)
    ex2_part2_circuit = ex2_part2_circuit.update(m_20_p1.chain("a0.h0"), separate_pos1)

if MAIN and PRINT_CIRCUITS:
    printer.print(ex2_part2_circuit.get_unique(m_10_p1))

if MAIN:
    tests.t_ex2_part2_circuit(ex2_part2_circuit)

# %% [markdown]
"""
First, make a new cond sampler that samples an input that agrees on what is called $p_1^($ in the writeup. This can be done with a `FuncSampler` based on a function with the following equivalence classes:
 - one class for _all_ inputs that start with a close parenthesis
 - one class for every value of $p$ (proportion of open parentheses in the entire sequence)

Note the actual values returned aren't important, just the equivalence clases. Also, remember to check the `ParenDataset` class for useful properties. For example, you might consider using `d.starts_with_open` and `d.p_open_after`.
"""


def p1_if_starts_open(d: ParenDataset):
    """Returns a tensor of size [len(ds)]. The value represents p_1 if the sequence starts open, and is constant otherwise"""
    "SOLUTION"
    return torch.where(
        d.starts_with_open,
        d.p_open_after[:, 1],
        torch.tensor(-1.0, device=DEVICE, dtype=torch.float32),
    )


p1_open_cond = FuncSampler(lambda d: p1_if_starts_open(ParenDataset.unwrap(d)))

# %% [markdown]
"""And some matchers"""

if "SOLUTION":
    m_10_p1_h00 = m_10_p1.chain("a0.h0.at_pos1")
    m_20_p1_h00 = m_20_p1.chain("a0.h0.at_pos1")
else:
    m_10_p1_h00: rc.IterativeMatcher
    m_20_p1_h00: rc.IterativeMatcher
    """TODO: Your code here"""

# TBD: add a test here?

# %% [markdown]
"""
Now make the correspondence!

You should add 4 nodes to the correspondence from part 1:
 - "10_p1_00"
 - "20_p1_00"
 - "10_p1_emb"
 - "20_p1_emb"

Refer to the writeup if you are unsure what to do with the `emb` nodes. This is [Experiment 2a](https://www.lesswrong.com/s/h95ayYYwMebGEYN5y/p/kjudfaQazMmC74SbF#2a__Dependency_on_0_0) in the writeup.
"""


def make_ex2_part2_corr() -> Correspondence:
    "SOLUTION"
    corr = make_ex2_part1_corr()

    i_10_p1 = corr.get_by_name("10_p1")
    i_20_p1 = corr.get_by_name("20_p1")

    # indirect paths to 0.0
    i_10_p1_00 = i_10_p1.make_descendant(name="10_p1_00", cond_sampler=p1_open_cond)
    i_20_p1_00 = i_20_p1.make_descendant(name="20_p1_00", cond_sampler=p1_open_cond)

    corr.add(i_10_p1_00, m_10_p1_h00)
    corr.add(i_20_p1_00, m_20_p1_h00)

    i_10_p1_emb = i_10_p1.make_descendant(name="10_p1_emb", cond_sampler=start_open_cond)
    i_20_p1_emb = i_20_p1.make_descendant(name="20_p1_emb", cond_sampler=start_open_cond)

    corr.add(i_10_p1_emb, chain_excluding(m_10_p1, "embeds", "a0.h0"))
    corr.add(i_20_p1_emb, chain_excluding(m_20_p1, "embeds", "a0.h0"))

    return corr


if MAIN:
    tests.t_ex2_part2_corr(make_ex2_part2_corr())
    print("\nEx 2 part 2 (2a in writeup): 1.0/2.0 depend on position 0.0 and emb")
    ex2a, loss2a = paren_experiment(ex2_part2_circuit, ds, make_ex2_part2_corr())
    check_loss(loss2a, 0.55, 0.04)


# %% [markdown]
"""
### Part 3: Projecting 0.0 onto a single direction
#### Circuit rewrite
Another claim we would like to test is that only the output of 0.0 written in a particular direction is important.

To do this we will rewrite the output of 0.0 as the sum of two terms: the [projection](https://en.wikipedia.org/wiki/Vector_projection) and rejection (aka the perpendicular component) along this direction.
"""

# %%

h00_open_vector = get_h00_open_vector(MODEL_ID).to(DEVICE)


def project_into_direction(c: rc.Circuit, v: torch.Tensor = h00_open_vector) -> rc.Circuit:
    """
    Rename `c` to `f"{c.name}_orig"`.
    Then return a circuit that computes (the renamed) `c`: [seq_len, 56] projected onto the direction of vector `v`: [56].
    Call the resulting circuit `{c.name}_projected`.
    """
    "SOLUTION"
    v_dir_array = rc.Array(v / torch.linalg.norm(v), "dir")
    return rc.Einsum.from_einsum_string(
        "s e, e, f -> s f",
        c.rename(c.name + "_orig"),
        v_dir_array,
        v_dir_array,
        name=f"{c.name}_projected",
    )


if MAIN:
    tests.t_project_into_direction(project_into_direction)


def get_ex2_part3_circuit(c: rc.Circuit, project_fn=project_into_direction):
    """
    Uses `residual_rewrite` to write head 0.0 at position 1 (when reached by either `m_10_p1_h00` or `m_20_p1_h00`), as a sum of the projection and the rejection along h00_open_vector. The head retains its same name, with children named `{head.name}_projected` and `{head.name}_projected_residual`.
    """
    "SOLUTION"
    split = lambda h00: residual_rewrite(h00, project_fn(h00), "projected")[0]
    return c.update((m_10_p1_h00 | m_20_p1_h00).chain("a0.h0"), split)


ex2_part3_circuit = get_ex2_part3_circuit(ex2_part2_circuit)

if MAIN and PRINT_CIRCUITS:
    proj_printer = printer.evolve(
        traversal=rc.new_traversal(term_early_at={"a0.h0.at_pos0", "a0.h0.at_pos2_42", "a0.h0_orig"})
    )
    subcirc = ex2_part3_circuit.get_unique(m_10_p1.chain("a0.h0_concat"))
    proj_printer.print(subcirc)

if MAIN:
    tests.t_ex2_part3_circuit(get_ex2_part3_circuit)

# %% [markdown]
"""
Now make the correspondence. Be sure to avoid the residual node!

This correspondence requires adding two new nodes:
 - "10_p1_00_projected"
 - "20_p1_00_projected"
"""


def make_ex2_part3_corr() -> Correspondence:
    "SOLUTION"
    corr = make_ex2_part2_corr()
    corr.add(
        corr.get_by_name("10_p1_00").make_descendant(p1_open_cond, "10_p1_00_projected"),
        chain_excluding(m_10_p1_h00, "a0.h0_projected", f"a0.h0_projected_residual"),
    )
    corr.add(
        corr.get_by_name("20_p1_00").make_descendant(p1_open_cond, "20_p1_00_projected"),
        chain_excluding(m_20_p1_h00, "a0.h0_projected", f"a0.h0_projected_residual"),
    )
    return corr


if MAIN:
    tests.t_ex2_part3_corr(make_ex2_part3_corr())
    print("\nEx 2 part 3: Projecting h00 into one direction")
    ex2_p3, loss2_p3 = paren_experiment(ex2_part3_circuit, ds, make_ex2_part3_corr())


# %% [markdown]
r"""
### Part 4: The $\phi$ function
"""


# %%
def compute_phi_circuit(tokens: rc.Circuit):
    """
    tokens: [seq_len, vocab_size] array of one hot tokens representing a sequence of parens
    (see ParenTokenizer for the one_hot ordering)

    Returns a circuit that computes phi: tokens -> R^56
    phi = h00_open_vector(2p - 1)
    where p = proportion of parens in `tokens` that are open.

    Returns a circuit with name 'a0.h0_phi'.
    """
    "SOLUTION"
    num_opens = rc.Index(tokens, [slice(None), ParenTokenizer.OPEN_TOKEN], name="is_open").sum(
        axis=-1, name="num_opens"
    )
    num_closes = rc.Index(tokens, [slice(None), ParenTokenizer.CLOSE_TOKEN], name="is_close").sum(
        axis=-1, name="num_closes"
    )
    num_parens = rc.Add(num_opens, num_closes, name="num_parens")
    p_open = num_opens.mul(rc.reciprocal(num_parens), name="p_open")
    p_open_mul_factor = p_open.mul_scalar(2).sub(rc.Scalar(1), name="p_open_mul_factor")
    dir = rc.Array(h00_open_vector, "open direction")
    return dir.mul(p_open_mul_factor, name="a0.h0_phi")


if MAIN:
    tests.t_compute_phi_circuit(compute_phi_circuit)

# %%


def get_ex2_part4_circuit(
    orig_circuit: rc.Circuit = ex2_part2_circuit,
    compute_phi_circuit_fn=compute_phi_circuit,
):
    """
    Split the output of head 0.0 at position 1, when reached through the appropriate paths, into a phi estimate
    and the residual of this estimate.

    The resulting subcircuit should have name 'a0.h0' with children 'a0.h0_phi' and 'a0.h0_phi_residual'.
    """
    "SOLUTION"
    split_by_phi = lambda c: residual_rewrite(c, compute_phi_circuit_fn(c.get_unique("tokens")), "phi")[0]
    return orig_circuit.update((m_10_p1_h00 | m_20_p1_h00).chain("a0.h0"), split_by_phi)


ex2_part4_circuit = get_ex2_part4_circuit()

if MAIN and PRINT_CIRCUITS:
    proj_printer = printer.evolve(
        traversal=rc.new_traversal(term_early_at={"a0.h0.at_pos0", "a0.h0.at_pos2_42", "a0.h0_orig"})
    )
    subcirc = ex2_part4_circuit.get_unique(m_10_p1.chain("a0.h0_concat"))
    proj_printer.print(subcirc)

if MAIN:
    tests.t_ex2_part4_circuit(get_ex2_part4_circuit)

# %% [markdown]
"""
And now make the correspondence -- it should be very similar to the one from part 3. Build on top of the one from part **2**, with new node names "10_p1_00_phi" and "20_p1_00_phi".
"""


def make_ex2_part4_corr() -> Correspondence:
    "SOLUTION"
    corr = make_ex2_part2_corr()
    corr.add(
        corr.get_by_name("10_p1_00").make_descendant(p1_open_cond, "10_p1_00_phi"),
        chain_excluding(m_10_p1_h00, "a0.h0_phi", f"a0.h0_phi_residual"),
    )
    corr.add(
        corr.get_by_name("20_p1_00").make_descendant(p1_open_cond, "20_p1_00_phi"),
        chain_excluding(m_20_p1_h00, "a0.h0_phi", f"a0.h0_phi_residual"),
    )
    return corr


if MAIN:
    tests.t_ex2_part4_corr(make_ex2_part4_corr())
    print("Ex2 part 4 (2b in writeup): replace a0 by phi(p)")
    ex2b, loss2b = paren_experiment(ex2_part4_circuit, ds, make_ex2_part4_corr())
    check_loss(loss2b, 0.53, 0.04)

# %% [markdown]
"""
Congradulations! This is the end of the main part of today's content. Below is some additional content that covers experiments 3 and 4 from the writeup, although with less detailed testing and instructions.
"""

# %%
if "SKIP":
    ######### Bonus experiments! ########
    # some bonus tests that don't appear in the writeup
    # these test the direct dependances of how information flows from a0 -> 2.0
    # note we don't have a dependance on the tokens here for the 'starts open'
    # which is probably a shortcoming of these hypotheses

    def make_ex2_explicit_paths_corr(from_10=True, to_m0_m1_a0=True, all_to_a0=True) -> Correspondence:
        corr = make_ex1_corr(count_open_cond)
        i_20 = corr.get_by_name("20")
        i_20_p1 = i_20.make_descendant(name="20_p1", cond_sampler=count_open_cond)
        corr.add(i_20_p1, m_20_p1)

        if to_m0_m1_a0:
            i_20_p1_a0 = i_20_p1.make_descendant(name="20_p1_a0", cond_sampler=count_open_cond)
            i_20_p1_m0 = i_20_p1.make_descendant(name="20_p1_m0", cond_sampler=count_open_cond)
            i_20_p1_m1 = i_20_p1.make_descendant(name="20_p1_m1", cond_sampler=count_open_cond)

            m_20_p1_m0 = chain_excluding(m_20_p1, "m0", "m1")
            m_20_p1_m1 = m_20_p1.chain("m1")
            corr.add(
                i_20_p1_a0,
                chain_excluding(m_20_p1, "a0.h0", {"m0", "a1.h0", "a1.h1", "m1"}),
            )
            corr.add(i_20_p1_m0, m_20_p1_m0)
            corr.add(i_20_p1_m1, m_20_p1_m1)

            if all_to_a0:
                # m0 depends on a0
                i_20_p1_m0_a0 = i_20_p1_m0.make_descendant(name="20_p1_m0_00", cond_sampler=p1_open_cond)

                # m1 depends on a0, m0
                i_20_p1_m1_a0 = i_20_p1_m1.make_descendant(name="20_p1_m1_00", cond_sampler=p1_open_cond)
                i_20_p1_m1_m0 = i_20_p1_m1.make_descendant(name="20_p1_m1_m0", cond_sampler=p1_open_cond)
                # m1.m0 depends on a0
                i_20_p1_m1_m0_a0 = i_20_p1_m1_m0.make_descendant(name="20_p1_m1_m0_00", cond_sampler=p1_open_cond)

                corr.add(i_20_p1_m0_a0, m_20_p1_m0.chain("a0.h0"))
                corr.add(i_20_p1_m1_a0, chain_excluding(m_20_p1_m1, "a0.h0", "m0"))
                corr.add(i_20_p1_m1_m0, m_20_p1_m1.chain("m0"))
                corr.add(i_20_p1_m1_m0_a0, m_20_p1_m1.chain("m0").chain("a0.h0"))

        if from_10:
            i_10 = corr.get_by_name("10")
            i_10_p1 = i_10.make_descendant(name="10_p1", cond_sampler=count_open_cond)
            corr.add(i_10_p1, m_10_p1)

            if to_m0_m1_a0:
                i_10_p1_a0 = i_10_p1.make_descendant(name="10_p1_a0", cond_sampler=count_open_cond)
                i_10_p1_m0 = i_10_p1.make_descendant(name="10_p1_m0", cond_sampler=count_open_cond)

                corr.add(i_10_p1_a0, chain_excluding(m_10_p1, "a0.h0", {"m0"}))
                m10_p1_m0 = m_10_p1.chain("m0")
                corr.add(i_10_p1_m0, m10_p1_m0)

                if all_to_a0:
                    i_10_p1_m0_a0 = i_10_p1_m0.make_descendant(name="10_p1_m0_00", cond_sampler=p1_open_cond)
                    corr.add(i_10_p1_m0_a0, m10_p1_m0.chain("a0.h0"))

        return corr

    # %%
    if MAIN:
        print("Ex 2_bonus_a: 2.0 depends on only pos1")
        ex2_bonus_a = paren_experiment(ex2_part2_circuit, ds, make_ex2_explicit_paths_corr(False, False, False))

        print("Ex 2_bonus_b: 2.0 depends on only pos1, which depends on a0 + m0 + m1")
        ex2_bonus_b = paren_experiment(ex2_part2_circuit, ds, make_ex2_explicit_paths_corr(False, True, False))

        print("Ex 2_bonus_c: 2.0 depends on only pos1, which depends on a0 + m0 + m1, which bottom out at a0")
        ex2_bonus_c = paren_experiment(ex2_part2_circuit, ds, make_ex2_explicit_paths_corr(False, True, True))

        if PRINT_CIRCUITS:
            ex2_part2_circuit.print(printer)


# %%
"""
## Experiment 3: Head 2.1

This experiment is officially optional content, and has less instructions / tests along the way. I expect it to be valuable practice for expressing informal hypotheses within the framework, but it may be somewhat more difficult to figure out exactly what the tests expect!
"""


def separate_all_seqpos(c: rc.Circuit) -> rc.Circuit:
    """
    Separate c into all possible sequence positions.
    c is renamed to `{c.name}_concat`, with children `{c.name}.at_pos{i}`
    """
    "SOLUTION"
    return split_to_concat(
        c,
        0,
        range(42),
        partitioning_idx_names=[f"pos{i}" for i in range(42)],
        use_dot=True,
    ).rename(f"{c.name}_concat")


if MAIN:
    tests.t_separate_all_seqpos(separate_all_seqpos)

ex3_circuit = circuit
if "SOLUTION":
    ex3_circuit = ex3_circuit.update(rc.IterativeMatcher("a2.h1").chain("a2.input"), separate_all_seqpos)

if MAIN and PRINT_CIRCUITS:
    printer.print(ex3_circuit.get_unique(rc.IterativeMatcher("a2.input_concat")))

if MAIN:
    tests.t_ex3_circuit(ex3_circuit)


# %% [markdown]
"""
When adjusted = True, use the `ds.adjusted_p_open_after` attribute instead of `ds.p_open_after` to compute the horizon test.

One possible gotcha in this section is late-binding-closures messing with the values of i. I think if you follow the outline you should be fine, but if you get strange bugs it's one possibility.
"""


def to_horizon_vals(d: ParenDataset, i: int, adjusted: bool = False) -> torch.Tensor:
    """
    Returns a value for the horizon_i test dividing up the input datums into 5 equivalence classes.
    The actual numerical return values don't have inherent meaning, but are defined as follows:
        0 on padding,
        positive on plausibly-balanced positions,
        negative on unbalance-evidence positions,
        1 / -1 on END_TOKENS,
        2 / -2 on non-end tokens
    """
    "SOLUTION"
    ps: torch.Tensor = d.adjusted_p_open_after if adjusted else d.p_open_after
    toks = d.tokens_flat.value
    assert toks.shape == ps.shape
    conds = [
        toks[:, i] == ParenTokenizer.PAD_TOKEN,
        (toks[:, i] == ParenTokenizer.END_TOKEN) & (toks[:, i - 1] == ParenTokenizer.OPEN_TOKEN),
        (toks[:, i] == ParenTokenizer.END_TOKEN) & (toks[:, i - 1] == ParenTokenizer.CLOSE_TOKEN),
        ps[:, i] > 0.5,
        ps[:, i] <= 0.5,
    ]
    vals = np.select(conds, [0, -1, 1, -2, 2], default=np.nan)
    if np.any(np.isnan(vals)):
        print(i)
        print(d[np.isnan(vals)][:10])
        print(ps[np.isnan(vals)][:10])
        raise AssertionError
    return torch.tensor(vals)


if MAIN:
    tests.t_to_horizon_vals(to_horizon_vals)


def get_horizon_cond(i: int, adjusted: bool) -> FuncSampler:
    """Func sampler for horizon_i"""
    "SOLUTION"
    return FuncSampler(lambda d: to_horizon_vals(ParenDataset.unwrap(d), i, adjusted))


def get_horizon_all_cond(adjusted: bool) -> FuncSampler:
    """Func sampler horizon_all"""
    "SOLUTION"

    def to_horizon_all_vals(d: ParenDataset):
        stacked = torch.stack([to_horizon_vals(d, i, adjusted) for i in range(0, 42)], dim=-1)
        return (stacked >= 0).all(dim=-1)

    return FuncSampler(lambda d: to_horizon_all_vals(ParenDataset.unwrap(d)))


if MAIN:
    tests.t_get_horizon_all_cond(get_horizon_all_cond)


def make_ex3_corr(adjusted: bool = False, corr=None) -> Correspondence:
    """
    `adjusted`: uses `adjusted_p_open_after` based conditions if True, `p_open_after` otherwise.
    `corr`: The starting corr. Uses experiemnt 1b corr by default (that is `make_ex1_corr(count_open_cond)`).

    Makes the following modifications:
     - Changes the cond sampler on node `21` to be the horizon_all cond sampler.
     - Adds one node for each sequence position, called `21_p{i}` with horizon_i cond sampler.
     - Also adds a node `pos_mask`, ensuring the pos_mask of head 2.1 is sampled from an input with the same input length.
    """
    "SOLUTION"
    corr = op.unwrap_or(corr, make_ex1_corr(count_open_cond))

    i_21 = corr.get_by_name("21")
    i_21.cond_sampler = get_horizon_all_cond(adjusted)
    for i in range(0, 42):
        i_pos_i = i_21.make_descendant(name=f"21_p{i}", cond_sampler=get_horizon_cond(i, adjusted))
        corr.add(i_pos_i, m_21.chain(f"a2.input.at_pos{i}"))

    i_pos_mask = i_21.make_descendant(
        name="pos_mask",
        cond_sampler=FuncSampler(lambda d: ParenDataset.unwrap(d).input_lengths),
    )

    m_pos_mask = chain_excluding(m_21, "pos_mask", "a2.input")
    corr.add(i_pos_mask, m_pos_mask)

    return corr


# %% [markdown]
"""
Note, we have a mini-replication crisis and with the current code can't replicate the exact numbers from the writeup. The handling of the position mask is somewhat different, although much more sane imo. I haven't had time to diagnose the exact cause of the difference.

In any case, expect your loss to be closer to ~1.23 for ex3a and ~1.17 for ex3b
"""
if MAIN:
    print("splitting up 2.1 input by seqpos")

    tests.t_make_ex3_corr(make_ex3_corr)
    print("\nEx3a: first with real open proportion")
    ex3a, loss3a = paren_experiment(ex3_circuit, ds, make_ex3_corr(adjusted=False))
    check_loss(loss3a, 1.124, 0.1)  # sad sad sad

    print("\nEx3b: now with adjusted open proportion")
    ex3b, loss3b = paren_experiment(ex3_circuit, ds, make_ex3_corr(adjusted=True))
    check_loss(loss3b, 1.140, 0.1)  # sad sad sad

    if PRINT_CIRCUITS:
        ex3b.print()

# %%

if "SKIP":

    def make_ex3c_corr(adjusted: bool):
        corr = make_ex3_corr(adjusted)
        for i in range(42):
            i_pi = corr.get_by_name(f"21_p{i}")
            m_pi = corr[i_pi]
            corr.add(
                i_pi.make_descendant(get_horizon_cond(i, adjusted), f"21_p{i}_a0"),
                chain_excluding(m_pi, "a0.h0", {"m0", "m1"}),
            )
            corr.add(
                i_pi.make_descendant(get_horizon_cond(i, adjusted), f"21_p{i}_m0"),
                chain_excluding(m_pi, "m0", "m1"),
            )
            corr.add(
                i_pi.make_descendant(get_horizon_cond(i, adjusted), f"21_p{i}_m1"),
                m_pi.chain("m1"),
            )

        return corr

    if MAIN and SLOW_EXPERIMENTS:
        print("Ex3c: Extending each seqpos to {a0.0, m0, m1}")
        ex3c = paren_experiment(ex3_circuit, ds, make_ex3c_corr(True))

    def make_ex3d_corr(adjusted: bool):
        corr = make_ex3_corr(adjusted)
        for i in range(42):
            i_pi = corr.get_by_name(f"21_p{i}")
            m_pi = corr[i_pi]

            corr.add(
                i_pi.make_descendant(ExactSampler(), f"21_p{i}_a0"),
                chain_excluding(m_pi, "a0.h0", {"m0", "m1"}),
            )
            corr.add(
                i_pi.make_descendant(ExactSampler(), f"21_p{i}_m0"),
                chain_excluding(m_pi, "m0", "m1"),
            )
            corr.add(i_pi.make_descendant(ExactSampler(), f"21_p{i}_m1"), m_pi.chain("m1"))

        return corr

    if MAIN and SLOW_EXPERIMENTS:
        print("Ex3d: Extending each seqpos to a0.0+m0+m1")
        ex3d = paren_experiment(ex3_circuit, ds, make_ex3d_corr(True))


# %% [markdown]
"""
## Experiment 4: All together
Now, combine experiments 2 (phi rewrite) and 3 (with adj. proportion)! Expected loss is ~1.64
"""

ex4_circuit = circuit
if "SOLUTION":
    ex4_circuit = ex2_part4_circuit.update(rc.IterativeMatcher("a2.h1").chain("a2.input"), separate_all_seqpos)
    ex4_corr = make_ex3_corr(adjusted=True, corr=make_ex2_part4_corr())
else:
    ex4_corr = Correspondence()

if MAIN and PRINT_CIRCUITS:
    printer.print(ex4_circuit)

if MAIN:
    print("\nEx4: Ex2b (1.0 and 2.0 phi rewrite) + Ex3b (2.1 split by seqpos with p_adj)")
    ex4, loss4 = paren_experiment(ex4_circuit, ds, ex4_corr)
    check_loss(loss4, 1.7, 0.1)  # sad sad sad

# %%