{
  "tokens": [
    "the",
    "cat",
    "sleeps"
  ],
  "X": [
    [
      1.0,
      0.0,
      1.0,
      0.0
    ],
    [
      0.0,
      1.0,
      0.0,
      1.0
    ],
    [
      1.0,
      1.0,
      0.0,
      0.0
    ]
  ],
  "head1_parameters": {
    "W_Q": [
      [
        1.0,
        0.0,
        0.0,
        0.0
      ],
      [
        0.0,
        1.0,
        0.0,
        0.0
      ]
    ],
    "W_K": [
      [
        0.0,
        0.0,
        1.0,
        0.0
      ],
      [
        0.0,
        0.0,
        0.0,
        1.0
      ]
    ],
    "W_V": [
      [
        1.0,
        0.0,
        1.0,
        0.0
      ],
      [
        0.0,
        1.0,
        0.0,
        2.0
      ]
    ]
  },
  "head2_parameters": {
    "W_Q": [
      [
        0.0,
        0.0,
        1.0,
        0.0
      ],
      [
        0.0,
        0.0,
        0.0,
        1.0
      ]
    ],
    "W_K": [
      [
        1.0,
        0.0,
        0.0,
        0.0
      ],
      [
        0.0,
        1.0,
        0.0,
        0.0
      ]
    ],
    "W_V": [
      [
        0.0,
        1.0,
        1.0,
        0.0
      ],
      [
        1.0,
        0.0,
        0.0,
        -1.0
      ]
    ]
  },
  "W_O": [
    [
      1.0,
      0.0,
      0.5,
      0.0
    ],
    [
      0.0,
      1.0,
      0.0,
      0.5
    ],
    [
      0.5,
      0.0,
      1.0,
      0.0
    ],
    [
      0.0,
      0.5,
      0.0,
      1.0
    ]
  ],
  "head1": {
    "Q": [
      [
        1.0,
        0.0
      ],
      [
        0.0,
        1.0
      ],
      [
        1.0,
        1.0
      ]
    ],
    "K": [
      [
        1.0,
        0.0
      ],
      [
        0.0,
        1.0
      ],
      [
        0.0,
        0.0
      ]
    ],
    "V": [
      [
        2.0,
        0.0
      ],
      [
        0.0,
        3.0
      ],
      [
        1.0,
        1.0
      ]
    ],
    "S": [
      [
        0.7071067811865475,
        0.0,
        0.0
      ],
      [
        0.0,
        0.7071067811865475,
        0.0
      ],
      [
        0.7071067811865475,
        0.7071067811865475,
        0.0
      ]
    ],
    "A": [
      [
        0.5034898434845538,
        0.2482550782577231,
        0.2482550782577231
      ],
      [
        0.2482550782577231,
        0.5034898434845538,
        0.2482550782577231
      ],
      [
        0.4011120926797859,
        0.4011120926797859,
        0.1977758146404282
      ]
    ],
    "O": [
      [
        1.2552347652268308,
        0.9930203130308924
      ],
      [
        0.7447652347731694,
        1.7587246087113846
      ],
      [
        1.0,
        1.401112092679786
      ]
    ]
  },
  "head2": {
    "Q": [
      [
        1.0,
        0.0
      ],
      [
        0.0,
        1.0
      ],
      [
        0.0,
        0.0
      ]
    ],
    "K": [
      [
        1.0,
        0.0
      ],
      [
        0.0,
        1.0
      ],
      [
        1.0,
        1.0
      ]
    ],
    "V": [
      [
        1.0,
        1.0
      ],
      [
        1.0,
        -1.0
      ],
      [
        1.0,
        1.0
      ]
    ],
    "S": [
      [
        0.7071067811865475,
        0.0,
        0.7071067811865475
      ],
      [
        0.0,
        0.7071067811865475,
        0.7071067811865475
      ],
      [
        0.0,
        0.0,
        0.0
      ]
    ],
    "A": [
      [
        0.4011120926797859,
        0.1977758146404282,
        0.4011120926797859
      ],
      [
        0.1977758146404282,
        0.4011120926797859,
        0.4011120926797859
      ],
      [
        0.3333333333333333,
        0.3333333333333333,
        0.3333333333333333
      ]
    ],
    "O": [
      [
        1.0,
        0.6044483707191437
      ],
      [
        1.0,
        0.1977758146404282
      ],
      [
        1.0,
        0.3333333333333333
      ]
    ]
  },
  "concatenated": [
    [
      1.2552347652268308,
      0.9930203130308924,
      1.0,
      0.6044483707191437
    ],
    [
      0.7447652347731694,
      1.7587246087113846,
      1.0,
      0.1977758146404282
    ],
    [
      1.0,
      1.401112092679786,
      1.0,
      0.3333333333333333
    ]
  ],
  "multihead_output": [
    [
      1.7552347652268308,
      1.2952444983904643,
      1.6276173826134155,
      1.1009585272345899
    ],
    [
      1.2447652347731695,
      1.8576125160315986,
      1.3723826173865847,
      1.0771381189961204
    ],
    [
      1.5,
      1.5677787593464527,
      1.5,
      1.0338893796732262
    ]
  ],
  "causal": {
    "Q": [
      [
        1.0,
        0.0
      ],
      [
        0.0,
        1.0
      ],
      [
        1.0,
        1.0
      ]
    ],
    "K": [
      [
        1.0,
        0.0
      ],
      [
        0.0,
        1.0
      ],
      [
        0.0,
        0.0
      ]
    ],
    "V": [
      [
        2.0,
        0.0
      ],
      [
        0.0,
        3.0
      ],
      [
        1.0,
        1.0
      ]
    ],
    "S": [
      [
        0.7071067811865475,
        "-Infinity",
        "-Infinity"
      ],
      [
        0.0,
        0.7071067811865475,
        "-Infinity"
      ],
      [
        0.7071067811865475,
        0.7071067811865475,
        0.0
      ]
    ],
    "A": [
      [
        1.0,
        0.0,
        0.0
      ],
      [
        0.3302384506733431,
        0.6697615493266569,
        0.0
      ],
      [
        0.4011120926797859,
        0.4011120926797859,
        0.1977758146404282
      ]
    ],
    "O": [
      [
        2.0,
        0.0
      ],
      [
        0.6604769013466862,
        2.0092846479799706
      ],
      [
        1.0,
        1.401112092679786
      ]
    ]
  },
  "permutation": [
    2,
    0,
    1
  ],
  "verification": {
    "sdpa_max_error": 1.1102230246251565e-16,
    "multihead_max_error": 2.220446049250313e-16,
    "causal_sdpa_max_error": 1.1102230246251565e-16,
    "projection_gradient_max_errors": {
      "W_Q": 1.69847186662353e-10,
      "W_K": 1.1172254787972236e-10,
      "W_V": 2.0961737901004085e-10,
      "W_O": 1.762088253087768e-10
    },
    "gru_max_error": 2.7755575615628914e-17,
    "lstm_h_max_error": 0.0
  },
  "gate_checks": {
    "GRU": {
      "parameters": {
        "weight_ih": [
          [
            -0.25,
            -0.2
          ],
          [
            -0.15,
            -0.1
          ],
          [
            -0.04999999999999999,
            0.0
          ],
          [
            0.04999999999999999,
            0.09999999999999998
          ],
          [
            0.15000000000000002,
            0.2
          ],
          [
            0.25,
            0.30000000000000004
          ]
        ],
        "weight_hh": [
          [
            -0.25,
            -0.2
          ],
          [
            -0.15,
            -0.1
          ],
          [
            -0.04999999999999999,
            0.0
          ],
          [
            0.04999999999999999,
            0.09999999999999998
          ],
          [
            0.15000000000000002,
            0.2
          ],
          [
            0.25,
            0.30000000000000004
          ]
        ],
        "bias_ih": [
          -0.25,
          -0.2,
          -0.15,
          -0.1,
          -0.04999999999999999,
          0.0
        ],
        "bias_hh": [
          -0.25,
          -0.2,
          -0.15,
          -0.1,
          -0.04999999999999999,
          0.0
        ]
      },
      "x": [
        0.2,
        0.8
      ],
      "h_prev": [
        0.4219,
        -0.1974
      ],
      "r": [
        0.3151836964743715,
        0.36504233366047995
      ],
      "u": [
        0.4179742166675213,
        0.4728654355203179
      ],
      "candidate": [
        0.13098681779174484,
        0.2976006495091789
      ],
      "h": [
        0.2525810272434962,
        0.06353195179618079
      ]
    },
    "LSTM": {
      "parameters": {
        "weight_ih": [
          [
            -0.25,
            -0.2
          ],
          [
            -0.15,
            -0.1
          ],
          [
            -0.04999999999999999,
            0.0
          ],
          [
            0.04999999999999999,
            0.09999999999999998
          ],
          [
            0.15000000000000002,
            0.2
          ],
          [
            0.25,
            0.30000000000000004
          ],
          [
            0.35,
            0.4
          ],
          [
            0.44999999999999996,
            0.5
          ]
        ],
        "weight_hh": [
          [
            -0.25,
            -0.2
          ],
          [
            -0.15,
            -0.1
          ],
          [
            -0.04999999999999999,
            0.0
          ],
          [
            0.04999999999999999,
            0.09999999999999998
          ],
          [
            0.15000000000000002,
            0.2
          ],
          [
            0.25,
            0.30000000000000004
          ],
          [
            0.35,
            0.4
          ],
          [
            0.44999999999999996,
            0.5
          ]
        ],
        "bias_ih": [
          -0.25,
          -0.2,
          -0.15,
          -0.1,
          -0.04999999999999999,
          0.0,
          0.04999999999999999,
          0.09999999999999998
        ],
        "bias_hh": [
          -0.25,
          -0.2,
          -0.15,
          -0.1,
          -0.04999999999999999,
          0.0,
          0.04999999999999999,
          0.09999999999999998
        ]
      },
      "x": [
        0.2,
        0.8
      ],
      "h_prev": [
        0.4219,
        -0.1974
      ],
      "c_prev": [
        0.6,
        -0.2
      ],
      "i": [
        0.3151836964743715,
        0.36504233366047995
      ],
      "f": [
        0.4179742166675213,
        0.4728654355203179
      ],
      "candidate": [
        0.11331621395179524,
        0.32412992431901044
      ],
      "o": [
        0.6361528493542472,
        0.6859289894806629
      ],
      "c": [
        0.2864999531843203,
        0.023748056878542728
      ],
      "h": [
        0.1774295176002329,
        0.016286419086895414
      ]
    }
  }
}