Evidence of Learned Look-Ahead#

Open In Colab

Setup#

[ ]:

MODE = "local"  # "colab" | "colab-dev" | "local"

[2]:

if MODE == "colab":
    !pip install -q lczerolens
elif MODE == "colab-dev":
    !rm -r lczerolens
    !git clone https://github.com/Xmaster6y/lczerolens -b main
    !pip install -q ./lczerolens

Cloning into 'lczerolens'...
remote: Enumerating objects: 867, done.
remote: Counting objects: 100% (367/367), done.
remote: Compressing objects: 100% (253/253), done.
remote: Total 867 (delta 203), reused 139 (delta 108), pack-reused 500 (from 1)
Receiving objects: 100% (867/867), 1.77 MiB | 6.19 MiB/s, done.
Resolving deltas: 100% (419/419), done.
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
  Building wheel for lczerolens (pyproject.toml) ... done

[1]:

!wget https://figshare.com/ndownloader/files/46473526?private_link=adc80845c00b67c8fce5 -O interesting_puzzles.pkl
!wget https://figshare.com/ndownloader/files/46473529?private_link=adc80845c00b67c8fce5 -O lc0.onnx

# When `wget` fail, e.g., "403 Forbidden"
# !pip install gdown
# !gdown https://drive.google.com/uc?id=1GT6I7FAgxWIxA-tzsifBQx0MkKZcR_qz -O interesting_puzzles.pkl
# !gdown https://drive.google.com/uc?id=1PB097ZKd_zTaPHxLK29WKUWmv6KcZ15T -O lc0.onnx

--2024-12-21 16:48:54--  https://figshare.com/ndownloader/files/46473526?private_link=adc80845c00b67c8fce5
Resolving figshare.com (figshare.com)... 18.202.168.57, 52.211.27.26, 2a05:d018:1f4:d000:dd63:62c3:56b8:c387, ...
Connecting to figshare.com (figshare.com)|18.202.168.57|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://s3-eu-west-1.amazonaws.com/pfigshare-u-files/46473526/interesting_puzzles.pkl?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIYCQYOYV5JSSROOA/20241221/eu-west-1/s3/aws4_request&X-Amz-Date=20241221T164854Z&X-Amz-Expires=10&X-Amz-SignedHeaders=host&X-Amz-Signature=ca1918d8041e0e440c9fb9a2ea3c5594aa20b4e0a4bd1e68bdb7203f55bc6375 [following]
--2024-12-21 16:48:54--  https://s3-eu-west-1.amazonaws.com/pfigshare-u-files/46473526/interesting_puzzles.pkl?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIYCQYOYV5JSSROOA/20241221/eu-west-1/s3/aws4_request&X-Amz-Date=20241221T164854Z&X-Amz-Expires=10&X-Amz-SignedHeaders=host&X-Amz-Signature=ca1918d8041e0e440c9fb9a2ea3c5594aa20b4e0a4bd1e68bdb7203f55bc6375
Resolving s3-eu-west-1.amazonaws.com (s3-eu-west-1.amazonaws.com)... 3.5.68.202, 52.218.105.210, 52.92.20.144, ...
Connecting to s3-eu-west-1.amazonaws.com (s3-eu-west-1.amazonaws.com)|3.5.68.202|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 11192747 (11M) [application/octet-stream]
Saving to: ‘interesting_puzzles.pkl’

interesting_puzzles 100%[===================>]  10.67M  8.51MB/s    in 1.3s

2024-12-21 16:48:56 (8.51 MB/s) - ‘interesting_puzzles.pkl’ saved [11192747/11192747]

--2024-12-21 16:48:57--  https://figshare.com/ndownloader/files/46473529?private_link=adc80845c00b67c8fce5
Resolving figshare.com (figshare.com)... 18.202.168.57, 52.211.27.26, 2a05:d018:1f4:d000:dd63:62c3:56b8:c387, ...
Connecting to figshare.com (figshare.com)|18.202.168.57|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://s3-eu-west-1.amazonaws.com/pfigshare-u-files/46473529/lc0.onnx?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIYCQYOYV5JSSROOA/20241221/eu-west-1/s3/aws4_request&X-Amz-Date=20241221T164858Z&X-Amz-Expires=10&X-Amz-SignedHeaders=host&X-Amz-Signature=a2d12b032c45401539ab24e25ce7f1a4b1d70ccbb4f3c154c51bc540032b7552 [following]
--2024-12-21 16:48:58--  https://s3-eu-west-1.amazonaws.com/pfigshare-u-files/46473529/lc0.onnx?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIYCQYOYV5JSSROOA/20241221/eu-west-1/s3/aws4_request&X-Amz-Date=20241221T164858Z&X-Amz-Expires=10&X-Amz-SignedHeaders=host&X-Amz-Signature=a2d12b032c45401539ab24e25ce7f1a4b1d70ccbb4f3c154c51bc540032b7552
Resolving s3-eu-west-1.amazonaws.com (s3-eu-west-1.amazonaws.com)... 3.5.68.202, 52.218.105.210, 52.92.20.144, ...
Connecting to s3-eu-west-1.amazonaws.com (s3-eu-west-1.amazonaws.com)|3.5.68.202|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 378669573 (361M) [application/octet-stream]
Saving to: ‘lc0.onnx’

lc0.onnx            100%[===================>] 361.13M  28.8MB/s    in 14s

2024-12-21 16:49:12 (25.4 MB/s) - ‘lc0.onnx’ saved [378669573/378669573]

Checking Assets#

[3]:

import pickle

import chess

[4]:

with open("interesting_puzzles.pkl", "rb") as f:
    puzzles = pickle.load(f)
puzzles.head()

[4]:
PuzzleId FEN Moves Rating RatingDeviation Popularity NbPlays Themes GameUrl OpeningTags principal_variation full_pv_probs full_model_moves full_wdl sparring_full_pv_probs sparring_full_model_moves sparring_wdl different_targets corrupted_fen
22 001w5 1rb2rk1/q5P1/4p2p/3p3p/3P1P2/2P5/2QK3P/3R2R1 b... f8f7 c2h7 g8h7 g7g8q 1073 77 91 189 advancedPawn attraction mate mateIn2 middlegam... https://lichess.org/0e1vxAEn/black#57 NaN [c2h7, g8h7, g7g8q] [0.7599087357521057, 1.0, 0.9156637191772461] [c2h7, g8h7, g7g8q] [0.561260461807251, 0.18697379529476166, 0.251... [0.006373010575771332, 1.0, 0.15674205124378204] [d1a1, g8h7, g7g8q] [0.0053426711820065975, 0.9625791907310486, 0.... False 1r4k1/q4rP1/2b1p2p/3p3p/3P1P2/2P5/2QK3P/3R2R1 ...
111 006wz 2r5/4ppkp/5bp1/1p6/1P6/P3B3/2r2PPP/1R1R2K1 b -... f6b2 b1b2 c2b2 e3d4 f7f6 d4b2 1515 75 73 527 attraction crushing endgame fork long sacrifice https://lichess.org/qT0W6o27/black#43 NaN [b1b2, c2b2, e3d4, f7f6, d4b2] [0.8258911967277527, 0.4339713156223297, 0.935... [b1b2, c2b2, e3d4, f7f6, d4b2] [0.816787600517273, 0.16938228905200958, 0.013... [0.04489869996905327, 0.881671667098999, 0.104... [d1d7, c2b2, e3d4, e7e5, d4b2] [0.0032313962001353502, 0.9889336824417114, 0.... False 2r5/4pp1p/6p1/1p3k2/1P6/P3B3/1br2PPP/1R1R2K1 w...
116 00761 3r2k1/1b3pbR/p2P2P1/3p2N1/2p5/2P2N2/PP6/2K5 b ... f7g6 h7g7 g8g7 g5e6 g7g8 e6d8 1512 75 94 17597 attraction crushing endgame exposedKing fork l... https://lichess.org/vu70Maig/black#55 NaN [h7g7, g8g7, g5e6, g7g8, e6d8] [0.7690635919570923, 0.81597501039505, 0.94436... [h7g7, g8g7, g5e6, g7f6, e6d8] [0.787416398525238, 0.18706896901130676, 0.025... [0.014131303876638412, 0.8102595806121826, 0.4... [f3d4, g8g7, g5e6, g7f6, e6d8] [0.005333698820322752, 0.9863110184669495, 0.0... False 2br2k1/6bR/p2P2p1/3p2N1/2p5/2P2N2/PP6/2K5 w - ...
170 00AoZ 8/1R6/p1pk4/6bp/1QP5/P7/KP6/3r2q1 b - - 2 44 g1c5 b7d7 d6d7 b4c5 1023 77 93 3354 advantage deflection endgame short https://lichess.org/356BAYqk/black#87 NaN [b7d7, d6d7, b4c5] [0.9532407522201538, 0.8907768130302429, 0.966... [b7d7, d6d7, b4c5] [0.8298567533493042, 0.16442830860614777, 0.00... [0.03029155358672142, 0.9142314195632935, 0.81... [b4b3, d6d7, b4c5] [0.01180175319314003, 0.947725236415863, 0.040... False 8/1R6/p1pk4/2q4p/1QP5/P6b/KP6/3r4 w - - 3 45
182 00Bg4 3r2k1/1q3ppp/p2rp3/Qp1B4/7P/P4P2/1PP3P1/1K1R3R... d6d5 a5d8 d5d8 d1d8 1374 85 75 303 backRankMate endgame mate mateIn2 short xRayAt... https://lichess.org/6qWf8wOP/black#41 NaN [a5d8, d5d8, d1d8] [0.9316630959510803, 1.0, 0.9474834203720093] [a5d8, d5d8, d1d8] [0.9491999745368958, 0.04444213956594467, 0.00... [0.008508995175361633, 1.0, 0.772297203540802] [d1d5, d5d8, d1d8] [0.007286431733518839, 0.983742892742157, 0.00... False 3r4/1q3ppp/p3p3/Qp1r4/5k1P/P4P2/1PP3P1/1K1R3R ... 
[5]:

from lczerolens import LczeroModel

model = LczeroModel.from_path("lc0.onnx")
model

[5]:

GraphModule(
  (attn_body/transpose): OnnxTranspose()
  (initializers): Module()
  (attn_body/reshape): OnnxReshape()
  (attn_body/shape): OnnxShape()
  (attn_body/batch): OnnxSlice()
  (attn_body/pos_encoding_shape): OnnxConcat()
  (attn_body/expand): OnnxExpand()
  (attn_body/padded_input): OnnxConcat()
  (attn_body/reshape2): OnnxReshape()
  (attn_body/matmul): OnnxMatMul()
  (attn_body/add): OnnxBinaryMathOperation()
  (attn_body/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (attn_body/mish/tanh): OnnxFunction()
  (attn_body/mish): OnnxBinaryMathOperation()
  (attn_body/ma_gating/rehape1): OnnxReshape()
  (ip_mul_gate): OnnxBinaryMathOperation()
  (ip_add_gate): OnnxBinaryMathOperation()
  (attn_body/ma_gating/rehape2): OnnxReshape()
  (encoder0/mha/Q/w): OnnxMatMul()
  (encoder0/mha/Q/b): OnnxBinaryMathOperation()
  (encoder0/mha/Q/reshape): OnnxReshape()
  (encoder0/mha/Q/transpose): OnnxTranspose()
  (encoder0/mha/K/w): OnnxMatMul()
  (encoder0/mha/K/b): OnnxBinaryMathOperation()
  (encoder0/mha/K/reshape): OnnxReshape()
  (encoder0/mha/K/transpose): OnnxTranspose()
  (encoder0/mha/V/w): OnnxMatMul()
  (encoder0/mha/V/b): OnnxBinaryMathOperation()
  (encoder0/mha/V/reshape): OnnxReshape()
  (encoder0/mha/V/transpose): OnnxTranspose()
  (encoder0/mha/QK/matmul): OnnxMatMul()
  (encoder0/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder0/smolgen/compress): OnnxMatMul()
  (encoder0/smolgen/compress/reshape): OnnxReshape()
  (encoder0/smolgen/dense1/w): OnnxMatMul()
  (encoder0/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder0/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder0/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder0/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder0/smolgen/dense2/w): OnnxMatMul()
  (encoder0/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder0/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder0/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder0/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder0/smolgen/gen_from/reshape): OnnxReshape()
  (encoder0/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder0/smolgen/out/reshape): OnnxReshape()
  (encoder0/smolgen_weights): OnnxBinaryMathOperation()
  (encoder0/mha/QK/softmax): Softmax(dim=3)
  (encoder0/mha/QKV/matmul): OnnxMatMul()
  (encoder0/mha/out/transpose): OnnxTranspose()
  (encoder0/mha/out/reshape): OnnxReshape()
  (encoder0/mha/out/dense/w): OnnxMatMul()
  (encoder0/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder0/alpha*input): OnnxBinaryMathOperation()
  (encoder0/mha/out/skip): OnnxBinaryMathOperation()
  (encoder0/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder0/ffn/dense1/w): OnnxMatMul()
  (encoder0/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder0/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder0/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder0/ffn/dense2/w): OnnxMatMul()
  (encoder0/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder0/alpha*out1): OnnxBinaryMathOperation()
  (encoder0/ffn/skip): OnnxBinaryMathOperation()
  (encoder0/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder1/mha/Q/w): OnnxMatMul()
  (encoder1/mha/Q/b): OnnxBinaryMathOperation()
  (encoder1/mha/Q/reshape): OnnxReshape()
  (encoder1/mha/Q/transpose): OnnxTranspose()
  (encoder1/mha/K/w): OnnxMatMul()
  (encoder1/mha/K/b): OnnxBinaryMathOperation()
  (encoder1/mha/K/reshape): OnnxReshape()
  (encoder1/mha/K/transpose): OnnxTranspose()
  (encoder1/mha/V/w): OnnxMatMul()
  (encoder1/mha/V/b): OnnxBinaryMathOperation()
  (encoder1/mha/V/reshape): OnnxReshape()
  (encoder1/mha/V/transpose): OnnxTranspose()
  (encoder1/mha/QK/matmul): OnnxMatMul()
  (encoder1/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder1/smolgen/compress): OnnxMatMul()
  (encoder1/smolgen/compress/reshape): OnnxReshape()
  (encoder1/smolgen/dense1/w): OnnxMatMul()
  (encoder1/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder1/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder1/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder1/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder1/smolgen/dense2/w): OnnxMatMul()
  (encoder1/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder1/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder1/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder1/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder1/smolgen/gen_from/reshape): OnnxReshape()
  (encoder1/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder1/smolgen/out/reshape): OnnxReshape()
  (encoder1/smolgen_weights): OnnxBinaryMathOperation()
  (encoder1/mha/QK/softmax): Softmax(dim=3)
  (encoder1/mha/QKV/matmul): OnnxMatMul()
  (encoder1/mha/out/transpose): OnnxTranspose()
  (encoder1/mha/out/reshape): OnnxReshape()
  (encoder1/mha/out/dense/w): OnnxMatMul()
  (encoder1/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder1/alpha*input): OnnxBinaryMathOperation()
  (encoder1/mha/out/skip): OnnxBinaryMathOperation()
  (encoder1/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder1/ffn/dense1/w): OnnxMatMul()
  (encoder1/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder1/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder1/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder1/ffn/dense2/w): OnnxMatMul()
  (encoder1/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder1/alpha*out1): OnnxBinaryMathOperation()
  (encoder1/ffn/skip): OnnxBinaryMathOperation()
  (encoder1/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder2/mha/Q/w): OnnxMatMul()
  (encoder2/mha/Q/b): OnnxBinaryMathOperation()
  (encoder2/mha/Q/reshape): OnnxReshape()
  (encoder2/mha/Q/transpose): OnnxTranspose()
  (encoder2/mha/K/w): OnnxMatMul()
  (encoder2/mha/K/b): OnnxBinaryMathOperation()
  (encoder2/mha/K/reshape): OnnxReshape()
  (encoder2/mha/K/transpose): OnnxTranspose()
  (encoder2/mha/V/w): OnnxMatMul()
  (encoder2/mha/V/b): OnnxBinaryMathOperation()
  (encoder2/mha/V/reshape): OnnxReshape()
  (encoder2/mha/V/transpose): OnnxTranspose()
  (encoder2/mha/QK/matmul): OnnxMatMul()
  (encoder2/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder2/smolgen/compress): OnnxMatMul()
  (encoder2/smolgen/compress/reshape): OnnxReshape()
  (encoder2/smolgen/dense1/w): OnnxMatMul()
  (encoder2/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder2/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder2/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder2/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder2/smolgen/dense2/w): OnnxMatMul()
  (encoder2/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder2/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder2/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder2/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder2/smolgen/gen_from/reshape): OnnxReshape()
  (encoder2/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder2/smolgen/out/reshape): OnnxReshape()
  (encoder2/smolgen_weights): OnnxBinaryMathOperation()
  (encoder2/mha/QK/softmax): Softmax(dim=3)
  (encoder2/mha/QKV/matmul): OnnxMatMul()
  (encoder2/mha/out/transpose): OnnxTranspose()
  (encoder2/mha/out/reshape): OnnxReshape()
  (encoder2/mha/out/dense/w): OnnxMatMul()
  (encoder2/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder2/alpha*input): OnnxBinaryMathOperation()
  (encoder2/mha/out/skip): OnnxBinaryMathOperation()
  (encoder2/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder2/ffn/dense1/w): OnnxMatMul()
  (encoder2/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder2/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder2/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder2/ffn/dense2/w): OnnxMatMul()
  (encoder2/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder2/alpha*out1): OnnxBinaryMathOperation()
  (encoder2/ffn/skip): OnnxBinaryMathOperation()
  (encoder2/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder3/mha/Q/w): OnnxMatMul()
  (encoder3/mha/Q/b): OnnxBinaryMathOperation()
  (encoder3/mha/Q/reshape): OnnxReshape()
  (encoder3/mha/Q/transpose): OnnxTranspose()
  (encoder3/mha/K/w): OnnxMatMul()
  (encoder3/mha/K/b): OnnxBinaryMathOperation()
  (encoder3/mha/K/reshape): OnnxReshape()
  (encoder3/mha/K/transpose): OnnxTranspose()
  (encoder3/mha/V/w): OnnxMatMul()
  (encoder3/mha/V/b): OnnxBinaryMathOperation()
  (encoder3/mha/V/reshape): OnnxReshape()
  (encoder3/mha/V/transpose): OnnxTranspose()
  (encoder3/mha/QK/matmul): OnnxMatMul()
  (encoder3/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder3/smolgen/compress): OnnxMatMul()
  (encoder3/smolgen/compress/reshape): OnnxReshape()
  (encoder3/smolgen/dense1/w): OnnxMatMul()
  (encoder3/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder3/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder3/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder3/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder3/smolgen/dense2/w): OnnxMatMul()
  (encoder3/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder3/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder3/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder3/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder3/smolgen/gen_from/reshape): OnnxReshape()
  (encoder3/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder3/smolgen/out/reshape): OnnxReshape()
  (encoder3/smolgen_weights): OnnxBinaryMathOperation()
  (encoder3/mha/QK/softmax): Softmax(dim=3)
  (encoder3/mha/QKV/matmul): OnnxMatMul()
  (encoder3/mha/out/transpose): OnnxTranspose()
  (encoder3/mha/out/reshape): OnnxReshape()
  (encoder3/mha/out/dense/w): OnnxMatMul()
  (encoder3/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder3/alpha*input): OnnxBinaryMathOperation()
  (encoder3/mha/out/skip): OnnxBinaryMathOperation()
  (encoder3/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder3/ffn/dense1/w): OnnxMatMul()
  (encoder3/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder3/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder3/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder3/ffn/dense2/w): OnnxMatMul()
  (encoder3/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder3/alpha*out1): OnnxBinaryMathOperation()
  (encoder3/ffn/skip): OnnxBinaryMathOperation()
  (encoder3/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder4/mha/Q/w): OnnxMatMul()
  (encoder4/mha/Q/b): OnnxBinaryMathOperation()
  (encoder4/mha/Q/reshape): OnnxReshape()
  (encoder4/mha/Q/transpose): OnnxTranspose()
  (encoder4/mha/K/w): OnnxMatMul()
  (encoder4/mha/K/b): OnnxBinaryMathOperation()
  (encoder4/mha/K/reshape): OnnxReshape()
  (encoder4/mha/K/transpose): OnnxTranspose()
  (encoder4/mha/V/w): OnnxMatMul()
  (encoder4/mha/V/b): OnnxBinaryMathOperation()
  (encoder4/mha/V/reshape): OnnxReshape()
  (encoder4/mha/V/transpose): OnnxTranspose()
  (encoder4/mha/QK/matmul): OnnxMatMul()
  (encoder4/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder4/smolgen/compress): OnnxMatMul()
  (encoder4/smolgen/compress/reshape): OnnxReshape()
  (encoder4/smolgen/dense1/w): OnnxMatMul()
  (encoder4/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder4/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder4/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder4/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder4/smolgen/dense2/w): OnnxMatMul()
  (encoder4/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder4/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder4/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder4/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder4/smolgen/gen_from/reshape): OnnxReshape()
  (encoder4/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder4/smolgen/out/reshape): OnnxReshape()
  (encoder4/smolgen_weights): OnnxBinaryMathOperation()
  (encoder4/mha/QK/softmax): Softmax(dim=3)
  (encoder4/mha/QKV/matmul): OnnxMatMul()
  (encoder4/mha/out/transpose): OnnxTranspose()
  (encoder4/mha/out/reshape): OnnxReshape()
  (encoder4/mha/out/dense/w): OnnxMatMul()
  (encoder4/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder4/alpha*input): OnnxBinaryMathOperation()
  (encoder4/mha/out/skip): OnnxBinaryMathOperation()
  (encoder4/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder4/ffn/dense1/w): OnnxMatMul()
  (encoder4/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder4/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder4/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder4/ffn/dense2/w): OnnxMatMul()
  (encoder4/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder4/alpha*out1): OnnxBinaryMathOperation()
  (encoder4/ffn/skip): OnnxBinaryMathOperation()
  (encoder4/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder5/mha/Q/w): OnnxMatMul()
  (encoder5/mha/Q/b): OnnxBinaryMathOperation()
  (encoder5/mha/Q/reshape): OnnxReshape()
  (encoder5/mha/Q/transpose): OnnxTranspose()
  (encoder5/mha/K/w): OnnxMatMul()
  (encoder5/mha/K/b): OnnxBinaryMathOperation()
  (encoder5/mha/K/reshape): OnnxReshape()
  (encoder5/mha/K/transpose): OnnxTranspose()
  (encoder5/mha/V/w): OnnxMatMul()
  (encoder5/mha/V/b): OnnxBinaryMathOperation()
  (encoder5/mha/V/reshape): OnnxReshape()
  (encoder5/mha/V/transpose): OnnxTranspose()
  (encoder5/mha/QK/matmul): OnnxMatMul()
  (encoder5/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder5/smolgen/compress): OnnxMatMul()
  (encoder5/smolgen/compress/reshape): OnnxReshape()
  (encoder5/smolgen/dense1/w): OnnxMatMul()
  (encoder5/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder5/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder5/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder5/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder5/smolgen/dense2/w): OnnxMatMul()
  (encoder5/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder5/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder5/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder5/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder5/smolgen/gen_from/reshape): OnnxReshape()
  (encoder5/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder5/smolgen/out/reshape): OnnxReshape()
  (encoder5/smolgen_weights): OnnxBinaryMathOperation()
  (encoder5/mha/QK/softmax): Softmax(dim=3)
  (encoder5/mha/QKV/matmul): OnnxMatMul()
  (encoder5/mha/out/transpose): OnnxTranspose()
  (encoder5/mha/out/reshape): OnnxReshape()
  (encoder5/mha/out/dense/w): OnnxMatMul()
  (encoder5/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder5/alpha*input): OnnxBinaryMathOperation()
  (encoder5/mha/out/skip): OnnxBinaryMathOperation()
  (encoder5/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder5/ffn/dense1/w): OnnxMatMul()
  (encoder5/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder5/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder5/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder5/ffn/dense2/w): OnnxMatMul()
  (encoder5/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder5/alpha*out1): OnnxBinaryMathOperation()
  (encoder5/ffn/skip): OnnxBinaryMathOperation()
  (encoder5/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder6/mha/Q/w): OnnxMatMul()
  (encoder6/mha/Q/b): OnnxBinaryMathOperation()
  (encoder6/mha/Q/reshape): OnnxReshape()
  (encoder6/mha/Q/transpose): OnnxTranspose()
  (encoder6/mha/K/w): OnnxMatMul()
  (encoder6/mha/K/b): OnnxBinaryMathOperation()
  (encoder6/mha/K/reshape): OnnxReshape()
  (encoder6/mha/K/transpose): OnnxTranspose()
  (encoder6/mha/V/w): OnnxMatMul()
  (encoder6/mha/V/b): OnnxBinaryMathOperation()
  (encoder6/mha/V/reshape): OnnxReshape()
  (encoder6/mha/V/transpose): OnnxTranspose()
  (encoder6/mha/QK/matmul): OnnxMatMul()
  (encoder6/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder6/smolgen/compress): OnnxMatMul()
  (encoder6/smolgen/compress/reshape): OnnxReshape()
  (encoder6/smolgen/dense1/w): OnnxMatMul()
  (encoder6/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder6/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder6/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder6/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder6/smolgen/dense2/w): OnnxMatMul()
  (encoder6/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder6/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder6/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder6/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder6/smolgen/gen_from/reshape): OnnxReshape()
  (encoder6/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder6/smolgen/out/reshape): OnnxReshape()
  (encoder6/smolgen_weights): OnnxBinaryMathOperation()
  (encoder6/mha/QK/softmax): Softmax(dim=3)
  (encoder6/mha/QKV/matmul): OnnxMatMul()
  (encoder6/mha/out/transpose): OnnxTranspose()
  (encoder6/mha/out/reshape): OnnxReshape()
  (encoder6/mha/out/dense/w): OnnxMatMul()
  (encoder6/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder6/alpha*input): OnnxBinaryMathOperation()
  (encoder6/mha/out/skip): OnnxBinaryMathOperation()
  (encoder6/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder6/ffn/dense1/w): OnnxMatMul()
  (encoder6/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder6/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder6/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder6/ffn/dense2/w): OnnxMatMul()
  (encoder6/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder6/alpha*out1): OnnxBinaryMathOperation()
  (encoder6/ffn/skip): OnnxBinaryMathOperation()
  (encoder6/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder7/mha/Q/w): OnnxMatMul()
  (encoder7/mha/Q/b): OnnxBinaryMathOperation()
  (encoder7/mha/Q/reshape): OnnxReshape()
  (encoder7/mha/Q/transpose): OnnxTranspose()
  (encoder7/mha/K/w): OnnxMatMul()
  (encoder7/mha/K/b): OnnxBinaryMathOperation()
  (encoder7/mha/K/reshape): OnnxReshape()
  (encoder7/mha/K/transpose): OnnxTranspose()
  (encoder7/mha/V/w): OnnxMatMul()
  (encoder7/mha/V/b): OnnxBinaryMathOperation()
  (encoder7/mha/V/reshape): OnnxReshape()
  (encoder7/mha/V/transpose): OnnxTranspose()
  (encoder7/mha/QK/matmul): OnnxMatMul()
  (encoder7/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder7/smolgen/compress): OnnxMatMul()
  (encoder7/smolgen/compress/reshape): OnnxReshape()
  (encoder7/smolgen/dense1/w): OnnxMatMul()
  (encoder7/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder7/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder7/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder7/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder7/smolgen/dense2/w): OnnxMatMul()
  (encoder7/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder7/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder7/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder7/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder7/smolgen/gen_from/reshape): OnnxReshape()
  (encoder7/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder7/smolgen/out/reshape): OnnxReshape()
  (encoder7/smolgen_weights): OnnxBinaryMathOperation()
  (encoder7/mha/QK/softmax): Softmax(dim=3)
  (encoder7/mha/QKV/matmul): OnnxMatMul()
  (encoder7/mha/out/transpose): OnnxTranspose()
  (encoder7/mha/out/reshape): OnnxReshape()
  (encoder7/mha/out/dense/w): OnnxMatMul()
  (encoder7/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder7/alpha*input): OnnxBinaryMathOperation()
  (encoder7/mha/out/skip): OnnxBinaryMathOperation()
  (encoder7/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder7/ffn/dense1/w): OnnxMatMul()
  (encoder7/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder7/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder7/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder7/ffn/dense2/w): OnnxMatMul()
  (encoder7/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder7/alpha*out1): OnnxBinaryMathOperation()
  (encoder7/ffn/skip): OnnxBinaryMathOperation()
  (encoder7/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder8/mha/Q/w): OnnxMatMul()
  (encoder8/mha/Q/b): OnnxBinaryMathOperation()
  (encoder8/mha/Q/reshape): OnnxReshape()
  (encoder8/mha/Q/transpose): OnnxTranspose()
  (encoder8/mha/K/w): OnnxMatMul()
  (encoder8/mha/K/b): OnnxBinaryMathOperation()
  (encoder8/mha/K/reshape): OnnxReshape()
  (encoder8/mha/K/transpose): OnnxTranspose()
  (encoder8/mha/V/w): OnnxMatMul()
  (encoder8/mha/V/b): OnnxBinaryMathOperation()
  (encoder8/mha/V/reshape): OnnxReshape()
  (encoder8/mha/V/transpose): OnnxTranspose()
  (encoder8/mha/QK/matmul): OnnxMatMul()
  (encoder8/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder8/smolgen/compress): OnnxMatMul()
  (encoder8/smolgen/compress/reshape): OnnxReshape()
  (encoder8/smolgen/dense1/w): OnnxMatMul()
  (encoder8/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder8/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder8/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder8/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder8/smolgen/dense2/w): OnnxMatMul()
  (encoder8/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder8/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder8/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder8/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder8/smolgen/gen_from/reshape): OnnxReshape()
  (encoder8/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder8/smolgen/out/reshape): OnnxReshape()
  (encoder8/smolgen_weights): OnnxBinaryMathOperation()
  (encoder8/mha/QK/softmax): Softmax(dim=3)
  (encoder8/mha/QKV/matmul): OnnxMatMul()
  (encoder8/mha/out/transpose): OnnxTranspose()
  (encoder8/mha/out/reshape): OnnxReshape()
  (encoder8/mha/out/dense/w): OnnxMatMul()
  (encoder8/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder8/alpha*input): OnnxBinaryMathOperation()
  (encoder8/mha/out/skip): OnnxBinaryMathOperation()
  (encoder8/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder8/ffn/dense1/w): OnnxMatMul()
  (encoder8/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder8/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder8/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder8/ffn/dense2/w): OnnxMatMul()
  (encoder8/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder8/alpha*out1): OnnxBinaryMathOperation()
  (encoder8/ffn/skip): OnnxBinaryMathOperation()
  (encoder8/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder9/mha/Q/w): OnnxMatMul()
  (encoder9/mha/Q/b): OnnxBinaryMathOperation()
  (encoder9/mha/Q/reshape): OnnxReshape()
  (encoder9/mha/Q/transpose): OnnxTranspose()
  (encoder9/mha/K/w): OnnxMatMul()
  (encoder9/mha/K/b): OnnxBinaryMathOperation()
  (encoder9/mha/K/reshape): OnnxReshape()
  (encoder9/mha/K/transpose): OnnxTranspose()
  (encoder9/mha/V/w): OnnxMatMul()
  (encoder9/mha/V/b): OnnxBinaryMathOperation()
  (encoder9/mha/V/reshape): OnnxReshape()
  (encoder9/mha/V/transpose): OnnxTranspose()
  (encoder9/mha/QK/matmul): OnnxMatMul()
  (encoder9/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder9/smolgen/compress): OnnxMatMul()
  (encoder9/smolgen/compress/reshape): OnnxReshape()
  (encoder9/smolgen/dense1/w): OnnxMatMul()
  (encoder9/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder9/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder9/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder9/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder9/smolgen/dense2/w): OnnxMatMul()
  (encoder9/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder9/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder9/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder9/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder9/smolgen/gen_from/reshape): OnnxReshape()
  (encoder9/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder9/smolgen/out/reshape): OnnxReshape()
  (encoder9/smolgen_weights): OnnxBinaryMathOperation()
  (encoder9/mha/QK/softmax): Softmax(dim=3)
  (encoder9/mha/QKV/matmul): OnnxMatMul()
  (encoder9/mha/out/transpose): OnnxTranspose()
  (encoder9/mha/out/reshape): OnnxReshape()
  (encoder9/mha/out/dense/w): OnnxMatMul()
  (encoder9/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder9/alpha*input): OnnxBinaryMathOperation()
  (encoder9/mha/out/skip): OnnxBinaryMathOperation()
  (encoder9/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder9/ffn/dense1/w): OnnxMatMul()
  (encoder9/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder9/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder9/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder9/ffn/dense2/w): OnnxMatMul()
  (encoder9/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder9/alpha*out1): OnnxBinaryMathOperation()
  (encoder9/ffn/skip): OnnxBinaryMathOperation()
  (encoder9/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder10/mha/Q/w): OnnxMatMul()
  (encoder10/mha/Q/b): OnnxBinaryMathOperation()
  (encoder10/mha/Q/reshape): OnnxReshape()
  (encoder10/mha/Q/transpose): OnnxTranspose()
  (encoder10/mha/K/w): OnnxMatMul()
  (encoder10/mha/K/b): OnnxBinaryMathOperation()
  (encoder10/mha/K/reshape): OnnxReshape()
  (encoder10/mha/K/transpose): OnnxTranspose()
  (encoder10/mha/V/w): OnnxMatMul()
  (encoder10/mha/V/b): OnnxBinaryMathOperation()
  (encoder10/mha/V/reshape): OnnxReshape()
  (encoder10/mha/V/transpose): OnnxTranspose()
  (encoder10/mha/QK/matmul): OnnxMatMul()
  (encoder10/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder10/smolgen/compress): OnnxMatMul()
  (encoder10/smolgen/compress/reshape): OnnxReshape()
  (encoder10/smolgen/dense1/w): OnnxMatMul()
  (encoder10/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder10/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder10/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder10/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder10/smolgen/dense2/w): OnnxMatMul()
  (encoder10/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder10/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder10/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder10/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder10/smolgen/gen_from/reshape): OnnxReshape()
  (encoder10/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder10/smolgen/out/reshape): OnnxReshape()
  (encoder10/smolgen_weights): OnnxBinaryMathOperation()
  (encoder10/mha/QK/softmax): Softmax(dim=3)
  (encoder10/mha/QKV/matmul): OnnxMatMul()
  (encoder10/mha/out/transpose): OnnxTranspose()
  (encoder10/mha/out/reshape): OnnxReshape()
  (encoder10/mha/out/dense/w): OnnxMatMul()
  (encoder10/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder10/alpha*input): OnnxBinaryMathOperation()
  (encoder10/mha/out/skip): OnnxBinaryMathOperation()
  (encoder10/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder10/ffn/dense1/w): OnnxMatMul()
  (encoder10/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder10/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder10/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder10/ffn/dense2/w): OnnxMatMul()
  (encoder10/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder10/alpha*out1): OnnxBinaryMathOperation()
  (encoder10/ffn/skip): OnnxBinaryMathOperation()
  (encoder10/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder11/mha/Q/w): OnnxMatMul()
  (encoder11/mha/Q/b): OnnxBinaryMathOperation()
  (encoder11/mha/Q/reshape): OnnxReshape()
  (encoder11/mha/Q/transpose): OnnxTranspose()
  (encoder11/mha/K/w): OnnxMatMul()
  (encoder11/mha/K/b): OnnxBinaryMathOperation()
  (encoder11/mha/K/reshape): OnnxReshape()
  (encoder11/mha/K/transpose): OnnxTranspose()
  (encoder11/mha/V/w): OnnxMatMul()
  (encoder11/mha/V/b): OnnxBinaryMathOperation()
  (encoder11/mha/V/reshape): OnnxReshape()
  (encoder11/mha/V/transpose): OnnxTranspose()
  (encoder11/mha/QK/matmul): OnnxMatMul()
  (encoder11/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder11/smolgen/compress): OnnxMatMul()
  (encoder11/smolgen/compress/reshape): OnnxReshape()
  (encoder11/smolgen/dense1/w): OnnxMatMul()
  (encoder11/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder11/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder11/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder11/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder11/smolgen/dense2/w): OnnxMatMul()
  (encoder11/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder11/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder11/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder11/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder11/smolgen/gen_from/reshape): OnnxReshape()
  (encoder11/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder11/smolgen/out/reshape): OnnxReshape()
  (encoder11/smolgen_weights): OnnxBinaryMathOperation()
  (encoder11/mha/QK/softmax): Softmax(dim=3)
  (encoder11/mha/QKV/matmul): OnnxMatMul()
  (encoder11/mha/out/transpose): OnnxTranspose()
  (encoder11/mha/out/reshape): OnnxReshape()
  (encoder11/mha/out/dense/w): OnnxMatMul()
  (encoder11/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder11/alpha*input): OnnxBinaryMathOperation()
  (encoder11/mha/out/skip): OnnxBinaryMathOperation()
  (encoder11/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder11/ffn/dense1/w): OnnxMatMul()
  (encoder11/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder11/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder11/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder11/ffn/dense2/w): OnnxMatMul()
  (encoder11/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder11/alpha*out1): OnnxBinaryMathOperation()
  (encoder11/ffn/skip): OnnxBinaryMathOperation()
  (encoder11/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder12/mha/Q/w): OnnxMatMul()
  (encoder12/mha/Q/b): OnnxBinaryMathOperation()
  (encoder12/mha/Q/reshape): OnnxReshape()
  (encoder12/mha/Q/transpose): OnnxTranspose()
  (encoder12/mha/K/w): OnnxMatMul()
  (encoder12/mha/K/b): OnnxBinaryMathOperation()
  (encoder12/mha/K/reshape): OnnxReshape()
  (encoder12/mha/K/transpose): OnnxTranspose()
  (encoder12/mha/V/w): OnnxMatMul()
  (encoder12/mha/V/b): OnnxBinaryMathOperation()
  (encoder12/mha/V/reshape): OnnxReshape()
  (encoder12/mha/V/transpose): OnnxTranspose()
  (encoder12/mha/QK/matmul): OnnxMatMul()
  (encoder12/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder12/smolgen/compress): OnnxMatMul()
  (encoder12/smolgen/compress/reshape): OnnxReshape()
  (encoder12/smolgen/dense1/w): OnnxMatMul()
  (encoder12/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder12/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder12/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder12/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder12/smolgen/dense2/w): OnnxMatMul()
  (encoder12/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder12/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder12/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder12/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder12/smolgen/gen_from/reshape): OnnxReshape()
  (encoder12/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder12/smolgen/out/reshape): OnnxReshape()
  (encoder12/smolgen_weights): OnnxBinaryMathOperation()
  (encoder12/mha/QK/softmax): Softmax(dim=3)
  (encoder12/mha/QKV/matmul): OnnxMatMul()
  (encoder12/mha/out/transpose): OnnxTranspose()
  (encoder12/mha/out/reshape): OnnxReshape()
  (encoder12/mha/out/dense/w): OnnxMatMul()
  (encoder12/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder12/alpha*input): OnnxBinaryMathOperation()
  (encoder12/mha/out/skip): OnnxBinaryMathOperation()
  (encoder12/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder12/ffn/dense1/w): OnnxMatMul()
  (encoder12/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder12/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder12/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder12/ffn/dense2/w): OnnxMatMul()
  (encoder12/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder12/alpha*out1): OnnxBinaryMathOperation()
  (encoder12/ffn/skip): OnnxBinaryMathOperation()
  (encoder12/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder13/mha/Q/w): OnnxMatMul()
  (encoder13/mha/Q/b): OnnxBinaryMathOperation()
  (encoder13/mha/Q/reshape): OnnxReshape()
  (encoder13/mha/Q/transpose): OnnxTranspose()
  (encoder13/mha/K/w): OnnxMatMul()
  (encoder13/mha/K/b): OnnxBinaryMathOperation()
  (encoder13/mha/K/reshape): OnnxReshape()
  (encoder13/mha/K/transpose): OnnxTranspose()
  (encoder13/mha/V/w): OnnxMatMul()
  (encoder13/mha/V/b): OnnxBinaryMathOperation()
  (encoder13/mha/V/reshape): OnnxReshape()
  (encoder13/mha/V/transpose): OnnxTranspose()
  (encoder13/mha/QK/matmul): OnnxMatMul()
  (encoder13/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder13/smolgen/compress): OnnxMatMul()
  (encoder13/smolgen/compress/reshape): OnnxReshape()
  (encoder13/smolgen/dense1/w): OnnxMatMul()
  (encoder13/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder13/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder13/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder13/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder13/smolgen/dense2/w): OnnxMatMul()
  (encoder13/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder13/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder13/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder13/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder13/smolgen/gen_from/reshape): OnnxReshape()
  (encoder13/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder13/smolgen/out/reshape): OnnxReshape()
  (encoder13/smolgen_weights): OnnxBinaryMathOperation()
  (encoder13/mha/QK/softmax): Softmax(dim=3)
  (encoder13/mha/QKV/matmul): OnnxMatMul()
  (encoder13/mha/out/transpose): OnnxTranspose()
  (encoder13/mha/out/reshape): OnnxReshape()
  (encoder13/mha/out/dense/w): OnnxMatMul()
  (encoder13/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder13/alpha*input): OnnxBinaryMathOperation()
  (encoder13/mha/out/skip): OnnxBinaryMathOperation()
  (encoder13/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder13/ffn/dense1/w): OnnxMatMul()
  (encoder13/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder13/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder13/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder13/ffn/dense2/w): OnnxMatMul()
  (encoder13/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder13/alpha*out1): OnnxBinaryMathOperation()
  (encoder13/ffn/skip): OnnxBinaryMathOperation()
  (encoder13/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder14/mha/Q/w): OnnxMatMul()
  (encoder14/mha/Q/b): OnnxBinaryMathOperation()
  (encoder14/mha/Q/reshape): OnnxReshape()
  (encoder14/mha/Q/transpose): OnnxTranspose()
  (encoder14/mha/K/w): OnnxMatMul()
  (encoder14/mha/K/b): OnnxBinaryMathOperation()
  (encoder14/mha/K/reshape): OnnxReshape()
  (encoder14/mha/K/transpose): OnnxTranspose()
  (encoder14/mha/V/w): OnnxMatMul()
  (encoder14/mha/V/b): OnnxBinaryMathOperation()
  (encoder14/mha/V/reshape): OnnxReshape()
  (encoder14/mha/V/transpose): OnnxTranspose()
  (encoder14/mha/QK/matmul): OnnxMatMul()
  (encoder14/mha/QK/scale): OnnxBinaryMathOperation()
  (encoder14/smolgen/compress): OnnxMatMul()
  (encoder14/smolgen/compress/reshape): OnnxReshape()
  (encoder14/smolgen/dense1/w): OnnxMatMul()
  (encoder14/smolgen/dense1/b): OnnxBinaryMathOperation()
  (encoder14/smolgen/dense1/swish/sigmoid): Sigmoid()
  (encoder14/smolgen/dense1/swish): OnnxBinaryMathOperation()
  (encoder14/smolgen/ln1): LayerNorm((256,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder14/smolgen/dense2/w): OnnxMatMul()
  (encoder14/smolgen/dense2/b): OnnxBinaryMathOperation()
  (encoder14/smolgen/dense2/swish/sigmoid): Sigmoid()
  (encoder14/smolgen/dense2/swish): OnnxBinaryMathOperation()
  (encoder14/smolgen/ln2): LayerNorm((6144,), eps=0.0010000000474974513, elementwise_affine=True)
  (encoder14/smolgen/gen_from/reshape): OnnxReshape()
  (encoder14/smolgen/smol_weight_gen): OnnxMatMul()
  (encoder14/smolgen/out/reshape): OnnxReshape()
  (encoder14/smolgen_weights): OnnxBinaryMathOperation()
  (encoder14/mha/QK/softmax): Softmax(dim=3)
  (encoder14/mha/QKV/matmul): OnnxMatMul()
  (encoder14/mha/out/transpose): OnnxTranspose()
  (encoder14/mha/out/reshape): OnnxReshape()
  (encoder14/mha/out/dense/w): OnnxMatMul()
  (encoder14/mha/out/dense/b): OnnxBinaryMathOperation()
  (encoder14/alpha*input): OnnxBinaryMathOperation()
  (encoder14/mha/out/skip): OnnxBinaryMathOperation()
  (encoder14/ln1): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (encoder14/ffn/dense1/w): OnnxMatMul()
  (encoder14/ffn/dense1/b): OnnxBinaryMathOperation()
  (encoder14/ffn/dense1/sqrrelu/relu): ReLU()
  (encoder14/ffn/dense1/sqrrelu/sqr): OnnxBinaryMathOperation()
  (encoder14/ffn/dense2/w): OnnxMatMul()
  (encoder14/ffn/dense2/b): OnnxBinaryMathOperation()
  (encoder14/alpha*out1): OnnxBinaryMathOperation()
  (encoder14/ffn/skip): OnnxBinaryMathOperation()
  (encoder14/ln2): LayerNorm((768,), eps=9.999999974752427e-07, elementwise_affine=True)
  (policy/dense1/matmul): OnnxMatMul()
  (policy/dense1/add): OnnxBinaryMathOperation()
  (policy/dense1/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (policy/dense1/mish/tanh): OnnxFunction()
  (policy/dense1/mish): OnnxBinaryMathOperation()
  (policy/Q/matmul): OnnxMatMul()
  (policy/Q/add): OnnxBinaryMathOperation()
  (policy/Q/reshape): OnnxReshape()
  (policy/K/matmul): OnnxMatMul()
  (policy/K/add): OnnxBinaryMathOperation()
  (policy/K/reshape): OnnxReshape()
  (policy/K/transpose): OnnxTranspose()
  (policy/matmul): OnnxMatMul()
  (policy/scale): OnnxBinaryMathOperation()
  (policy/promotion/slice): OnnxSlice()
  (policy/promotion/matmul): OnnxMatMul()
  (policy/promotion/transpose): OnnxTranspose()
  (policy/promotion/split): OnnxSplit13()
  (policy/promotion/add): OnnxBinaryMathOperation()
  (policy/promotion/transpose2): OnnxTranspose()
  (policy/promotion/reshape): OnnxReshape()
  (policy/promotion/slice2): OnnxSlice()
  (policy/promotion/reshape2): OnnxReshape()
  (policy/promotion/concat): OnnxConcat()
  (policy/promotion/reshape3): OnnxReshape()
  (policy/promotion/add2): OnnxBinaryMathOperation()
  (policy/promotion/reshape4): OnnxReshape()
  (policy/concat): OnnxConcat()
  (policy/reshape): OnnxReshape()
  (output/policy): OnnxGather()
  (value/embed/matmul): OnnxMatMul()
  (value/embed/add): OnnxBinaryMathOperation()
  (value/embed/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (value/embed/mish/tanh): OnnxFunction()
  (value/embed/mish): OnnxBinaryMathOperation()
  (value/reshape): OnnxReshape()
  (value/dense1/matmul): OnnxMatMul()
  (value/dense1/add): OnnxBinaryMathOperation()
  (value/dense1/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (value/dense1/mish/tanh): OnnxFunction()
  (value/dense1/mish): OnnxBinaryMathOperation()
  (value/dense2/matmul): OnnxMatMul()
  (value/dense2/add): OnnxBinaryMathOperation()
  (output/wdl): Softmax(dim=1)
  (mlh/embed/matmul): OnnxMatMul()
  (mlh/embed/add): OnnxBinaryMathOperation()
  (mlh/embed/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (mlh/embed/mish/tanh): OnnxFunction()
  (mlh/embed/mish): OnnxBinaryMathOperation()
  (mlh/reshape): OnnxReshape()
  (mlh/dense1/matmul): OnnxMatMul()
  (mlh/dense1/add): OnnxBinaryMathOperation()
  (mlh/dense1/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (mlh/dense1/mish/tanh): OnnxFunction()
  (mlh/dense1/mish): OnnxBinaryMathOperation()
  (mlh/dense2/matmul): OnnxMatMul()
  (mlh/dense2/add): OnnxBinaryMathOperation()
  (mlh/dense2/mish/softplus): Softplus(beta=1.0, threshold=20.0)
  (mlh/dense2/mish/tanh): OnnxFunction()
  (mlh/dense2/mish): OnnxBinaryMathOperation()
  (output/mlh): OnnxCopyIdentity()
)

[6]:

import IPython

from lczerolens.board import LczeroBoard

puzzle = puzzles.loc[19612]
board = LczeroBoard(puzzle.FEN)
moves = puzzle.Moves.split()
board.push_uci(moves[0])
corrupted_board = LczeroBoard(puzzle.corrupted_fen)
display(board)
display(corrupted_board)
../../_images/notebooks_tutorials_evidence-of-learned-look-ahead_9_0.svg
../../_images/notebooks_tutorials_evidence-of-learned-look-ahead_9_1.svg 
[7]:

out = model(*[board, corrupted_board])
out["wdl"]

[7]:

tensor([[0.3586, 0.4711, 0.1703],
        [0.0132, 0.1089, 0.8779]], grad_fn=<SoftmaxBackward0>)

Visualising Attention#

[8]:

import chess.svg
import matplotlib
import matplotlib.pyplot as plt


COLOR_MAP = matplotlib.colormaps["RdYlBu_r"].resampled(1000)
ALPHA = 1.0
NORM = matplotlib.colors.Normalize(vmin=0, vmax=1, clip=False)


def render_heatmap(
    board,
    heatmap,
    square=None,
    vmin=None,
    vmax=None,
    arrows=None,
    normalise="none",
    save_to=None,
):
    """
    Render a heatmap on the board.
    """
    if normalise == "abs":
        a_max = heatmap.abs().max()
        if a_max != 0:
            heatmap = heatmap / a_max
        vmin = -1
        vmax = 1
    if vmin is None:
        vmin = heatmap.min()
    if vmax is None:
        vmax = heatmap.max()
    norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=False)

    color_dict = {}
    for square_index in range(64):
        color = COLOR_MAP(norm(heatmap[square_index]))
        color = (*color[:3], ALPHA)
        color_dict[square_index] = matplotlib.colors.to_hex(color, keep_alpha=True)
    fig = plt.figure(figsize=(1, 6))
    ax = plt.gca()
    ax.axis("off")
    fig.colorbar(
        matplotlib.cm.ScalarMappable(norm=norm, cmap=COLOR_MAP),
        ax=ax,
        orientation="vertical",
        fraction=1.0,
    )
    if square is not None:
        try:
            check = chess.parse_square(square)
        except ValueError:
            check = None
    else:
        check = None
    if arrows is None:
        arrows = []

    return chess.svg.board(
        board,
        check=check,
        fill=color_dict,
        size=350,
        arrows=arrows,
    )

[9]:

layer = 9
head = 5

with model.trace(board):
    attention = getattr(model, f"encoder{layer}/mha/QK/softmax").output[0, head].save()
attention.shape

[9]:

torch.Size([64, 64])

[10]:

square = chess.F4

boardsvg = render_heatmap(board, attention[square].detach())
display(IPython.display.HTML((boardsvg)))

. . r . . r . k
. . . . . . p p
p . q . . p . .
. p . . . . . .
. . . R . N . .
P Q . . P . . P
. P . . . P P .
. . . . . . K .
../../_images/notebooks_tutorials_evidence-of-learned-look-ahead_14_1.png

Probing Analysis#

Activation Patching#

[11]:

from lczerolens.lenses import ActivationLens

MODULE = "encoder13/ln2"
act_lens = ActivationLens(MODULE)

clean_acts = act_lens.analyse(board, model)
corrupted_acts = act_lens.analyse(corrupted_board, model)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-11-7b3ad2d816ce> in <cell line: 6>()
      4 act_lens = ActivationLens(MODULE)
      5
----> 6 clean_acts = act_lens.analyse(board, model)
      7 corrupted_acts = act_lens.analyse(corrupted_board, model)

TypeError: ActivationLens.analyse() missing 1 required keyword-only argument: 'model'

[ ]:

corrupted_acts[MODULE].shape

[ ]:

with model.trace(board):
    out = model.output.save()

[ ]: