adapters.py

from __future__ import annotations

import builtins
import math
import os
from collections import Counter
from collections.abc import Iterable
from typing import IO, Any, BinaryIO

import numpy.typing as npt
import regex as re
import torch
from jaxtyping import Bool, Float, Int
from torch import Tensor
from torch.optim import Optimizer
from .common import gpt2_bytes_to_unicode

_original_open = builtins.open


def _open_utf8_for_merges(file, *args, **kwargs):
    if "encoding" not in kwargs and isinstance(file, (str, os.PathLike)):
        mode = args[0] if args else "r"
        if "b" not in mode:
            return _original_open(file, *args, encoding="utf-8", **kwargs)
    return _original_open(file, *args, **kwargs)


builtins.open = _open_utf8_for_merges


def run_linear(
    d_in: int,
    d_out: int,
    weights: Float[Tensor, " d_out d_in"],
    in_features: Float[Tensor, " ... d_in"],
) -> Float[Tensor, " ... d_out"]:
    """
    Given the weights of a Linear layer, compute the transformation of a batched input.

    Args:
        in_dim (int): The size of the input dimension
        out_dim (int): The size of the output dimension
        weights (Float[Tensor, "d_out d_in"]): The linear weights to use
        in_features (Float[Tensor, "... d_in"]): The output tensor to apply the function to

    Returns:
        Float[Tensor, "... d_out"]: The transformed output of your linear module.
    """

    return in_features @ weights.t()


def run_embedding(
    vocab_size: int,
    d_model: int,
    weights: Float[Tensor, " vocab_size d_model"],
    token_ids: Int[Tensor, " ..."],
) -> Float[Tensor, " ... d_model"]:
    """
    Given the weights of an Embedding layer, get the embeddings for a batch of token ids.

    Args:
        vocab_size (int): The number of embeddings in the vocabulary
        d_model (int): The size of the embedding dimension
        weights (Float[Tensor, "vocab_size d_model"]): The embedding vectors to fetch from
        token_ids (Int[Tensor, "..."]): The set of token ids to fetch from the Embedding layer

    Returns:
        Float[Tensor, "... d_model"]: Batch of embeddings returned by your Embedding layer.
    """

    return weights[token_ids]


def run_swiglu(
    d_model: int,
    d_ff: int,
    w1_weight: Float[Tensor, " d_ff d_model"],
    w2_weight: Float[Tensor, " d_model d_ff"],
    w3_weight: Float[Tensor, " d_ff d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    """Given the weights of a SwiGLU network, return
    the output of your implementation with these weights.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        d_ff (int): Dimensionality of the up-project happening internally to your swiglu.
        w1_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W1
        w2_weight (Float[Tensor, "d_model d_ff"]): Stored weights for W2
        w3_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W3
        in_features (Float[Tensor, "... d_model"]): Input embeddings to the feed-forward layer.

    Returns:
        Float[Tensor, "... d_model"]: Output embeddings of the same shape as the input embeddings.
    """
    # Example:
    # If your state dict keys match, you can use `load_state_dict()`
    # swiglu.load_state_dict(weights)
    # You can also manually assign the weights
    # swiglu.w1.weight.data = w1_weight
    # swiglu.w2.weight.data = w2_weight
    # swiglu.w3.weight.data = w3_weight
    w1_out = in_features @ w1_weight.t()
    w3_out = in_features @ w3_weight.t()
    gated = run_silu(w1_out) * w3_out
    return gated @ w2_weight.t()


def run_scaled_dot_product_attention(
    Q: Float[Tensor, " ... queries d_k"],
    K: Float[Tensor, " ... keys d_k"],
    V: Float[Tensor, " ... values d_v"],
    mask: Bool[Tensor, " ... queries keys"] | None = None,
) -> Float[Tensor, " ... queries d_v"]:
    """
    Given key (K), query (Q), and value (V) tensors, return
    the output of your scaled dot product attention implementation.

    Args:
        Q (Float[Tensor, " ... queries d_k"]): Query tensor
        K (Float[Tensor, " ... keys d_k"]): Key tensor
        V (Float[Tensor, " ... values d_v"]): Values tensor
        mask (Bool[Tensor, " ... queries keys"] | None): Mask tensor
    Returns:
        Float[Tensor, " ... queries d_v"]: Output of SDPA
    """
    d_k = Q.shape[-1]
    scores = (Q @ K.transpose(-1, -2)) / math.sqrt(d_k)
    if mask is not None:
        scores = torch.where(mask, scores, torch.full_like(scores, -float("inf")))
    attn = run_softmax(scores, dim=-1)
    return attn @ V


def run_multihead_self_attention(
    d_model: int,
    num_heads: int,
    q_proj_weight: Float[Tensor, " d_k d_in"],
    k_proj_weight: Float[Tensor, " d_k d_in"],
    v_proj_weight: Float[Tensor, " d_v d_in"],
    o_proj_weight: Float[Tensor, " d_model d_v"],
    in_features: Float[Tensor, " ... sequence_length d_in"],
) -> Float[Tensor, " ... sequence_length d_out"]:
    """
    Given the key, query, and value projection weights of a naive unbatched
    implementation of multi-head attention, return the output of an optimized batched
    implementation. This implementation should handle the key, query, and value projections
    for all heads in a single matrix multiply.
    This function should not use RoPE.
    See section 3.2.2 of Vaswani et al., 2017.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        num_heads (int): Number of heads to use in multi-headed attention.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        q_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the Q projection
        k_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the K projection
        v_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the V projection
        o_proj_weight (Float[Tensor, "d_model d_v"]): Weights for the output projection
        in_features (Float[Tensor, "... sequence_length d_in"]): Tensor to run your implementation on.

    Returns:
        Float[Tensor, " ... sequence_length d_out"]: Tensor with the output of running your optimized, batched multi-headed attention
        implementation with the given QKV projection weights and input features.
    """
    d_k = d_model // num_heads
    q = in_features @ q_proj_weight.t()
    k = in_features @ k_proj_weight.t()
    v = in_features @ v_proj_weight.t()
    new_shape = q.shape[:-1] + (num_heads, d_k)
    q = q.reshape(new_shape).transpose(-3, -2)
    k = k.reshape(new_shape).transpose(-3, -2)
    v = v.reshape(new_shape).transpose(-3, -2)
    seq_len = in_features.shape[-2]
    mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=in_features.device))
    for _ in range(q.ndim - mask.ndim):
        mask = mask.unsqueeze(0)
    attn_out = run_scaled_dot_product_attention(Q=q, K=k, V=v, mask=mask)
    attn_out = attn_out.transpose(-3, -2).reshape(in_features.shape[:-1] + (d_model,))
    return attn_out @ o_proj_weight.t()


def run_multihead_self_attention_with_rope(
    d_model: int,
    num_heads: int,
    max_seq_len: int,
    theta: float,
    q_proj_weight: Float[Tensor, " d_k d_in"],
    k_proj_weight: Float[Tensor, " d_k d_in"],
    v_proj_weight: Float[Tensor, " d_v d_in"],
    o_proj_weight: Float[Tensor, " d_model d_v"],
    in_features: Float[Tensor, " ... sequence_length d_in"],
    token_positions: Int[Tensor, " ... sequence_length"] | None = None,
) -> Float[Tensor, " ... sequence_length d_out"]:
    """
    Given the key, query, and value projection weights of a naive unbatched
    implementation of multi-head attention, return the output of an optimized batched
    implementation. This implementation should handle the key, query, and value projections
    for all heads in a single matrix multiply.
    This version of MHA should include RoPE.
    In this case, the RoPE embedding dimension must be the head embedding dimension (d_model // num_heads).
    See section 3.2.2 of Vaswani et al., 2017.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        num_heads (int): Number of heads to use in multi-headed attention.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        theta (float): RoPE parameter.
        q_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the Q projection
        k_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the K projection
        v_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the V projection
        o_proj_weight (Float[Tensor, "d_model d_v"]): Weights for the output projection
        in_features (Float[Tensor, "... sequence_length d_in"]): Tensor to run your implementation on.
        token_positions (Int[Tensor, " ... sequence_length"] | None): Optional tensor with the positions of the tokens

    Returns:
        Float[Tensor, " ... sequence_length d_out"]: Tensor with the output of running your optimized, batched multi-headed attention
        implementation with the given QKV projection weights and input features.
    """
    d_k = d_model // num_heads
    q = in_features @ q_proj_weight.t()
    k = in_features @ k_proj_weight.t()
    v = in_features @ v_proj_weight.t()
    new_shape = q.shape[:-1] + (num_heads, d_k)
    q = q.reshape(new_shape).transpose(-3, -2)
    k = k.reshape(new_shape).transpose(-3, -2)
    v = v.reshape(new_shape).transpose(-3, -2)
    if token_positions is None:
        seq_len = in_features.shape[-2]
        token_positions = torch.arange(seq_len, device=in_features.device)
        token_positions = token_positions.expand(in_features.shape[:-1])
    q = run_rope(d_k=d_k, theta=theta, max_seq_len=max_seq_len, in_query_or_key=q, token_positions=token_positions)
    k = run_rope(d_k=d_k, theta=theta, max_seq_len=max_seq_len, in_query_or_key=k, token_positions=token_positions)
    seq_len = in_features.shape[-2]
    mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=in_features.device))
    for _ in range(q.ndim - mask.ndim):
        mask = mask.unsqueeze(0)
    attn_out = run_scaled_dot_product_attention(Q=q, K=k, V=v, mask=mask)
    attn_out = attn_out.transpose(-3, -2).reshape(in_features.shape[:-1] + (d_model,))
    return attn_out @ o_proj_weight.t()


def run_rope(
    d_k: int,
    theta: float,
    max_seq_len: int,
    in_query_or_key: Float[Tensor, " ... sequence_length d_k"],
    token_positions: Int[Tensor, " ... sequence_length"],
) -> Float[Tensor, " ... sequence_length d_k"]:
    """
    Run RoPE for a given input tensor.

    Args:
        d_k (int): Embedding dimension size for the query or key tensor.
        theta (float): RoPE parameter.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        in_query_or_key (Float[Tensor, "... sequence_length d_k"]): Input tensor to run RoPE on.
        token_positions (Int[Tensor, "... sequence_length"]): Tensor of shape (batch_size, sequence_length) with the token positions
    Returns:
        Float[Tensor, " ... sequence_length d_k"]: Tensor with RoPEd input.
    """
    d_half = d_k // 2
    inv_freq = theta ** (-2 * torch.arange(0, d_half, device=in_query_or_key.device) / d_k)
    while token_positions.ndim < in_query_or_key.ndim - 1:
        token_positions = token_positions.unsqueeze(-2)
    angles = token_positions[..., None].to(in_query_or_key.device) * inv_freq
    cos = torch.cos(angles)
    sin = torch.sin(angles)
    x_even = in_query_or_key[..., 0::2]
    x_odd = in_query_or_key[..., 1::2]
    rotated_even = x_even * cos - x_odd * sin
    rotated_odd = x_even * sin + x_odd * cos
    out = torch.empty_like(in_query_or_key)
    out[..., 0::2] = rotated_even
    out[..., 1::2] = rotated_odd
    return out


def run_transformer_block(
    d_model: int,
    num_heads: int,
    d_ff: int,
    max_seq_len: int,
    theta: float,
    weights: dict[str, Tensor],
    in_features: Float[Tensor, " batch sequence_length d_model"],
) -> Float[Tensor, " batch sequence_length d_model"]:
    """
    Given the weights of a pre-norm Transformer block and input features,
    return the output of running the Transformer block on the input features.

    This function should use RoPE.
    Depending on your implementation, you may simply need to pass the relevant args
    to your TransformerBlock constructor, or you may need to initialize your own RoPE
    class and pass that instead.

    Args:
        d_model (int): The dimensionality of the Transformer block input.
        num_heads (int): Number of heads to use in multi-headed attention. `d_model` must be
            evenly divisible by `num_heads`.
        d_ff (int): Dimensionality of the feed-forward inner layer.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        theta (float): RoPE parameter.
        weights (dict[str, Tensor]):
            State dict of our reference implementation.
            The keys of this dictionary are:
            - `attn.q_proj.weight`
                The query projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.q_proj.weight == torch.cat([q_heads.0.weight, ..., q_heads.N.weight], dim=0)`.
            - `attn.k_proj.weight`
                The key projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.k_proj.weight == torch.cat([k_heads.0.weight, ..., k_heads.N.weight], dim=0)`.
            - `attn.v_proj.weight`
                The value projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_v),
                so `attn.v_proj.weight == torch.cat([v_heads.0.weight, ..., v_heads.N.weight], dim=0)`.
            - `attn.output_proj.weight`
                Weight of the multi-head self-attention output projection
                Shape is (d_model, d_model).
            - `ln1.weight`
                Weights of affine transform for the first RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `ffn.w1.weight`
                Weight of the first linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `ffn.w2.weight`
                Weight of the second linear transformation in the FFN.
                Shape is (d_ff, d_model).
            - `ffn.w3.weight`
                Weight of the third linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `ln2.weight`
                Weights of affine transform for the second RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
        in_features (Float[Tensor, "batch sequence_length d_model"]):
            Tensor to run your implementation on.

    Returns:
        Float[Tensor, "batch sequence_length d_model"] Tensor with the output of
        running the Transformer block on the input features while using RoPE.
    """
    x = in_features
    ln1 = run_rmsnorm(d_model=d_model, eps=1e-5, weights=weights["ln1.weight"], in_features=x)
    seq_len = x.shape[-2]
    causal_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device))
    q_proj_weight = weights["attn.q_proj.weight"]
    k_proj_weight = weights["attn.k_proj.weight"]
    v_proj_weight = weights["attn.v_proj.weight"]
    o_proj_weight = weights["attn.output_proj.weight"]
    d_k = d_model // num_heads
    q = ln1 @ q_proj_weight.t()
    k = ln1 @ k_proj_weight.t()
    v = ln1 @ v_proj_weight.t()
    new_shape = q.shape[:-1] + (num_heads, d_k)
    q = q.reshape(new_shape).transpose(-3, -2)
    k = k.reshape(new_shape).transpose(-3, -2)
    v = v.reshape(new_shape).transpose(-3, -2)
    token_positions = torch.arange(seq_len, device=x.device).expand(x.shape[:-1])
    q = run_rope(d_k=d_k, theta=theta, max_seq_len=max_seq_len, in_query_or_key=q, token_positions=token_positions)
    k = run_rope(d_k=d_k, theta=theta, max_seq_len=max_seq_len, in_query_or_key=k, token_positions=token_positions)
    mask = causal_mask
    for _ in range(q.ndim - mask.ndim):
        mask = mask.unsqueeze(0)
    attn_out = run_scaled_dot_product_attention(Q=q, K=k, V=v, mask=mask)
    attn_out = attn_out.transpose(-3, -2).reshape(x.shape[:-1] + (d_model,))
    x = x + (attn_out @ o_proj_weight.t())
    ln2 = run_rmsnorm(d_model=d_model, eps=1e-5, weights=weights["ln2.weight"], in_features=x)
    w1_weight = weights["ffn.w1.weight"]
    w2_weight = weights["ffn.w2.weight"]
    w3_weight = weights["ffn.w3.weight"]
    ffn_out = run_swiglu(d_model=d_model, d_ff=d_ff, w1_weight=w1_weight, w2_weight=w2_weight, w3_weight=w3_weight, in_features=ln2)
    return x + ffn_out


def run_transformer_lm(
    vocab_size: int,
    context_length: int,
    d_model: int,
    num_layers: int,
    num_heads: int,
    d_ff: int,
    rope_theta: float,
    weights: dict[str, Tensor],
    in_indices: Int[Tensor, " batch_size sequence_length"],
) -> Float[Tensor, " batch_size sequence_length vocab_size"]:
    """Given the weights of a Transformer language model and input indices,
    return the output of running a forward pass on the input indices.

    This function should use RoPE.

    Args:
        vocab_size (int): The number of unique items in the output vocabulary to be predicted.
        context_length (int): The maximum number of tokens to process at once.
        d_model (int): The dimensionality of the model embeddings and sublayer outputs.
        num_layers (int): The number of Transformer layers to use.
        num_heads (int): Number of heads to use in multi-headed attention. `d_model` must be
            evenly divisible by `num_heads`.
        d_ff (int): Dimensionality of the feed-forward inner layer (section 3.3).
        rope_theta (float): The RoPE Theta parameter.
        weights (dict[str, Tensor]):
            State dict of our reference implementation. {num_layers} refers to an
            integer between `0` and `num_layers - 1` (the layer index).
            The keys of this dictionary are:
            - `token_embeddings.weight`
                Token embedding matrix. Shape is (vocab_size, d_model).
            - `layers.{num_layers}.attn.q_proj.weight`
                The query projections for all `num_heads` attention heads.
                Shape is (num_heads * (d_model / num_heads), d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.q_proj.weight == torch.cat([q_heads.0.weight, ..., q_heads.N.weight], dim=0)`.
            - `layers.{num_layers}.attn.k_proj.weight`
                The key projections for all `num_heads` attention heads.
                Shape is (num_heads * (d_model / num_heads), d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.k_proj.weight == torch.cat([k_heads.0.weight, ..., k_heads.N.weight], dim=0)`.
            - `layers.{num_layers}.attn.v_proj.weight`
                The value projections for all `num_heads` attention heads.
                Shape is (num_heads * (d_model / num_heads), d_model).
                The rows are ordered by matrices of shape (num_heads, d_v),
                so `attn.v_proj.weight == torch.cat([v_heads.0.weight, ..., v_heads.N.weight], dim=0)`.
            - `layers.{num_layers}.attn.output_proj.weight`
                Weight of the multi-head self-attention output projection
                Shape is ((d_model / num_heads) * num_heads, d_model).
            - `layers.{num_layers}.ln1.weight`
                Weights of affine transform for the first RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `layers.{num_layers}.ffn.w1.weight`
                Weight of the first linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `layers.{num_layers}.ffn.w2.weight`
                Weight of the second linear transformation in the FFN.
                Shape is (d_ff, d_model).
            - `layers.{num_layers}.ffn.w3.weight`
                Weight of the third linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `layers.{num_layers}.ln2.weight`
                Weights of affine transform for the second RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `ln_final.weight`
                Weights of affine transform for RMSNorm applied to the output of the final transformer block.
                Shape is (d_model, ).
            - `lm_head.weight`
                Weights of the language model output embedding.
                Shape is (vocab_size, d_model).
        in_indices (Int[Tensor, "batch_size sequence_length"]) Tensor with input indices to run the language model on. Shape is (batch_size, sequence_length), where
            `sequence_length` is at most `context_length`.

    Returns:
        Float[Tensor, "batch_size sequence_length vocab_size"]: Tensor with the predicted unnormalized
        next-word distribution for each token.
    """
    x = run_embedding(vocab_size=vocab_size, d_model=d_model, weights=weights["token_embeddings.weight"], token_ids=in_indices)
    for layer in range(num_layers):
        prefix = f"layers.{layer}."
        block_weights = {k[len(prefix) :]: v for k, v in weights.items() if k.startswith(prefix)}
        x = run_transformer_block(
            d_model=d_model,
            num_heads=num_heads,
            d_ff=d_ff,
            max_seq_len=context_length,
            theta=rope_theta,
            weights=block_weights,
            in_features=x,
        )
    x = run_rmsnorm(d_model=d_model, eps=1e-5, weights=weights["ln_final.weight"], in_features=x)
    return x @ weights["lm_head.weight"].t()


def run_rmsnorm(
    d_model: int,
    eps: float,
    weights: Float[Tensor, " d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    """Given the weights of a RMSNorm affine transform,
    return the output of running RMSNorm on the input features.

    Args:
        d_model (int): The dimensionality of the RMSNorm input.
        eps: (float): A value added to the denominator for numerical stability.
        weights (Float[Tensor, "d_model"]): RMSNorm weights.
        in_features (Float[Tensor, "... d_model"]): Input features to run RMSNorm on. Can have arbitrary leading
            dimensions.

    Returns:
        Float[Tensor,"... d_model"]: Tensor of with the same shape as `in_features` with the output of running
        RMSNorm of the `in_features`.
    """
    rms = torch.sqrt(torch.mean(in_features * in_features, dim=-1, keepdim=True) + eps)
    return in_features / rms * weights


def run_silu(in_features: Float[Tensor, " ..."]) -> Float[Tensor, " ..."]:
    """Given a tensor of inputs, return the output of applying SiLU
    to each element.

    Args:
        in_features(Float[Tensor, "..."]): Input features to run SiLU on. Shape is arbitrary.

    Returns:
        Float[Tensor,"..."]: of with the same shape as `in_features` with the output of applying
        SiLU to each element.
    """
    return in_features * torch.sigmoid(in_features)


def run_get_batch(
    dataset: npt.NDArray, batch_size: int, context_length: int, device: str
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Given a dataset (a 1D numpy array of integers) and a desired batch size and
    context length, sample language modeling input sequences and their corresponding
    labels from the dataset.

    Args:
        dataset (np.array): 1D numpy array of integer token IDs in the dataset.
        batch_size (int): Desired batch size to sample.
        context_length (int): Desired context length of each sampled example.
        device (str): PyTorch device string (e.g., 'cpu' or 'cuda:0') indicating the device
            to place the sampled input sequences and labels on.

    Returns:
        Tuple of torch.LongTensors of shape (batch_size, context_length). The first tuple item
        is the sampled input sequences, and the second tuple item is the corresponding
        language modeling labels.
    """
    max_start = len(dataset) - context_length
    starts = torch.randint(0, max_start, (batch_size,))
    x = torch.stack([torch.from_numpy(dataset[s : s + context_length]) for s in starts.tolist()])
    y = torch.stack([torch.from_numpy(dataset[s + 1 : s + context_length + 1]) for s in starts.tolist()])
    return x.to(device=device), y.to(device=device)


def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]:
    """
    Given a tensor of inputs, return the output of softmaxing the given `dim`
    of the input.

    Args:
        in_features (Float[Tensor, "..."]): Input features to softmax. Shape is arbitrary.
        dim (int): Dimension of the `in_features` to apply softmax to.

    Returns:
        Float[Tensor, "..."]: Tensor of with the same shape as `in_features` with the output of
        softmax normalizing the specified `dim`.
    """
    shifted = in_features - torch.max(in_features, dim=dim, keepdim=True).values
    exp = torch.exp(shifted)
    return exp / torch.sum(exp, dim=dim, keepdim=True)


def run_cross_entropy(
    inputs: Float[Tensor, " batch_size vocab_size"], targets: Int[Tensor, " batch_size"]
) -> Float[Tensor, ""]:
    """Given a tensor of inputs and targets, compute the average cross-entropy
    loss across examples.

    Args:
        inputs (Float[Tensor, "batch_size vocab_size"]): inputs[i][j] is the
            unnormalized logit of jth class for the ith example.
        targets (Int[Tensor, "batch_size"]): Tensor of shape (batch_size,) with the index of the correct class.
            Each value must be between 0 and `num_classes - 1`.

    Returns:
        Float[Tensor, ""]: The average cross-entropy loss across examples.
    """
    max_logits = torch.max(inputs, dim=-1, keepdim=True).values
    logsumexp = max_logits + torch.log(torch.sum(torch.exp(inputs - max_logits), dim=-1, keepdim=True))
    log_probs = inputs - logsumexp
    losses = -log_probs[torch.arange(targets.shape[0]), targets]
    return losses.mean()


def run_gradient_clipping(parameters: Iterable[torch.nn.Parameter], max_l2_norm: float) -> None:
    """Given a set of parameters, clip their combined gradients to have l2 norm at most max_l2_norm.

    Args:
        parameters (Iterable[torch.nn.Parameter]): collection of trainable parameters.
        max_l2_norm (float): a positive value containing the maximum l2-norm.

    The gradients of the parameters (parameter.grad) should be modified in-place.
    """
    grads = [p.grad for p in parameters if p.grad is not None]
    if not grads:
        return
    total_norm = torch.sqrt(sum(torch.sum(g * g) for g in grads))
    if total_norm <= max_l2_norm:
        return
    scale = max_l2_norm / (total_norm + 1e-6)
    for g in grads:
        g.mul_(scale)


def get_adamw_cls() -> Any:
    """
    Returns a torch.optim.Optimizer that implements AdamW.
    """
    class AdamW(Optimizer):
        def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
            defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
            super().__init__(params, defaults)

        def step(self, closure=None):
            loss = None
            if closure is not None:
                loss = closure()
            with torch.no_grad():
                for group in self.param_groups:
                    lr = group["lr"]
                    beta1, beta2 = group["betas"]
                    eps = group["eps"]
                    weight_decay = group["weight_decay"]
                    for p in group["params"]:
                        if p.grad is None:
                            continue
                        grad = p.grad
                        state = self.state[p]
                        if len(state) == 0:
                            state["step"] = 0
                            state["exp_avg"] = torch.zeros_like(p)
                            state["exp_avg_sq"] = torch.zeros_like(p)
                        exp_avg = state["exp_avg"]
                        exp_avg_sq = state["exp_avg_sq"]
                        state["step"] += 1
                        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                        bias_correction1 = 1 - beta1 ** state["step"]
                        bias_correction2 = 1 - beta2 ** state["step"]
                        step_size = lr * math.sqrt(bias_correction2) / bias_correction1
                        if weight_decay != 0:
                            p.add_(p, alpha=-lr * weight_decay)
                        denom = exp_avg_sq.sqrt().add_(eps)
                        p.addcdiv_(exp_avg, denom, value=-step_size)
            return loss

    return AdamW


def run_get_lr_cosine_schedule(
    it: int,
    max_learning_rate: float,
    min_learning_rate: float,
    warmup_iters: int,
    cosine_cycle_iters: int,
):
    """
    Given the parameters of a cosine learning rate decay schedule (with linear
    warmup) and an iteration number, return the learning rate at the given
    iteration under the specified schedule.

    Args:
        it (int): Iteration number to get learning rate for.
        max_learning_rate (float): alpha_max, the maximum learning rate for
            cosine learning rate schedule (with warmup).
        min_learning_rate (float): alpha_min, the minimum / final learning rate for
            the cosine learning rate schedule (with warmup).
        warmup_iters (int): T_w, the number of iterations to linearly warm-up
            the learning rate.
        cosine_cycle_iters (int): T_c, the number of cosine annealing iterations.

    Returns:
        Learning rate at the given iteration under the specified schedule.
    """
    if it < warmup_iters:
        return max_learning_rate * it / warmup_iters
    if it >= cosine_cycle_iters:
        return min_learning_rate
    decay_iters = cosine_cycle_iters - warmup_iters
    if decay_iters <= 0:
        return min_learning_rate
    t = (it - warmup_iters) / decay_iters
    cosine = 0.5 * (1 + math.cos(math.pi * t))
    return min_learning_rate + cosine * (max_learning_rate - min_learning_rate)


def run_save_checkpoint(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    iteration: int,
    out: str | os.PathLike | BinaryIO | IO[bytes],
):
    """
    Given a model, optimizer, and an iteration number, serialize them to disk.

    Args:
        model (torch.nn.Module): Serialize the state of this model.
        optimizer (torch.optim.Optimizer): Serialize the state of this optimizer.
        iteration (int): Serialize this value, which represents the number of training iterations
            we've completed.
        out (str | os.PathLike | BinaryIO | IO[bytes]): Path or file-like object to serialize the model, optimizer, and iteration to.
    """
    payload = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "iteration": iteration,
    }
    torch.save(payload, out)


def run_load_checkpoint(
    src: str | os.PathLike | BinaryIO | IO[bytes],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
) -> int:
    """
    Given a serialized checkpoint (path or file-like object), restore the
    serialized state to the given model and optimizer.
    Return the number of iterations that we previously serialized in
    the checkpoint.

    Args:
        src (str | os.PathLike | BinaryIO | IO[bytes]): Path or file-like object to serialized checkpoint.
        model (torch.nn.Module): Restore the state of this model.
        optimizer (torch.optim.Optimizer): Restore the state of this optimizer.
    Returns:
        int: the previously-serialized number of iterations.
    """
    payload = torch.load(src, map_location="cpu")
    model.load_state_dict(payload["model_state_dict"])
    optimizer.load_state_dict(payload["optimizer_state_dict"])
    return int(payload["iteration"])


def get_tokenizer(
    vocab: dict[int, bytes],
    merges: list[tuple[bytes, bytes]],
    special_tokens: list[str] | None = None,
) -> Any:
    """Given a vocabulary, a list of merges, and a list of special tokens,
    return a BPE tokenizer that uses the provided vocab, merges, and special tokens.

    Args:
        vocab (dict[int, bytes]): The tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
            to bytes (token bytes)
        merges (list[tuple[bytes, bytes]]): BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
            representing that <token1> was merged with <token2>.
            Merges are ordered by order of creation.
        special_tokens (list[str] | None): A list of string special tokens for the tokenizer. These strings will never
            be split into multiple tokens, and will always be kept as a single token.

    Returns:
        A BPE tokenizer that uses the provided vocab, merges, and special tokens.
    """
    class BPETokenizer:
        def __init__(self, vocab, merges, special_tokens):
            self.byte_encoder = gpt2_bytes_to_unicode()
            self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
            self.id_to_token_bytes = dict(vocab)
            self.bytes_to_id = {v: k for k, v in self.id_to_token_bytes.items()}
            self.special_tokens = special_tokens or []
            self.special_token_set = set(self.special_tokens)
            self.special_token_bytes = {s: s.encode("utf-8") for s in self.special_tokens}
            self.special_token_ids = {
                s: self.bytes_to_id[self.special_token_bytes[s]]
                for s in self.special_tokens
                if self.special_token_bytes[s] in self.bytes_to_id
            }
            self.bpe_ranks = {
                (self._bytes_to_unicode_str(a), self._bytes_to_unicode_str(b)): i for i, (a, b) in enumerate(merges)
            }
            self.cache = {}
            self.pattern = re.compile(
                r"'(?:[sdmt]|ll|ve|re)| ?\p{L}++| ?\p{N}++| ?[^\s\p{L}\p{N}]++|\s++$|\s+(?!\S)|\s"
            )
            if self.special_tokens:
                sorted_tokens = sorted(self.special_tokens, key=len, reverse=True)
                self.special_regex = re.compile("|".join(re.escape(tok) for tok in sorted_tokens))
            else:
                self.special_regex = None
            self.token_to_id = {
                self._bytes_to_unicode_str(token_bytes): token_id for token_id, token_bytes in self.id_to_token_bytes.items()
            }

        def _bytes_to_unicode_str(self, token_bytes):
            if isinstance(token_bytes, str):
                return token_bytes
            return "".join(self.byte_encoder[b] for b in token_bytes)

        def _get_pairs(self, word):
            return {(word[i], word[i + 1]) for i in range(len(word) - 1)}

        def _bpe(self, token):
            if token in self.cache:
                return self.cache[token]
            word = tuple(token)
            pairs = self._get_pairs(word)
            if not pairs:
                self.cache[token] = word
                return word
            while True:
                min_pair = None
                min_rank = None
                for pair in pairs:
                    rank = self.bpe_ranks.get(pair)
                    if rank is None:
                        continue
                    if min_rank is None or rank < min_rank:
                        min_rank = rank
                        min_pair = pair
                if min_pair is None:
                    break
                first, second = min_pair
                new_word = []
                i = 0
                while i < len(word):
                    if i < len(word) - 1 and word[i] == first and word[i + 1] == second:
                        new_word.append(first + second)
                        i += 2
                    else:
                        new_word.append(word[i])
                        i += 1
                word = tuple(new_word)
                if len(word) == 1:
                    break
                pairs = self._get_pairs(word)
            self.cache[token] = word
            return word

        def _encode_piece(self, piece):
            ids = []
            for token in self.pattern.findall(piece):
                token_bytes = token.encode("utf-8")
                token_unicode = self._bytes_to_unicode_str(token_bytes)
                for bpe_token in self._bpe(token_unicode):
                    ids.append(self.token_to_id[bpe_token])
            return ids

        def encode(self, text):
            if text == "":
                return []
            if not self.special_tokens:
                return self._encode_piece(text)
            ids = []
            last = 0
            for match in self.special_regex.finditer(text):
                if match.start() > last:
                    ids.extend(self._encode_piece(text[last : match.start()]))
                special = match.group(0)
                ids.append(self.special_token_ids[special])
                last = match.end()
            if last < len(text):
                ids.extend(self._encode_piece(text[last:]))
            return ids

        def encode_iterable(self, iterable):
            for piece in iterable:
                for _id in self.encode(piece):
                    yield _id

        def decode(self, ids):
            if not ids:
                return ""
            byte_chunks = [self.id_to_token_bytes[_id] for _id in ids]
            return b"".join(byte_chunks).decode("utf-8", errors="replace")

    return BPETokenizer(vocab, merges, special_tokens)


def run_train_bpe(
    input_path: str | os.PathLike,
    vocab_size: int,
    special_tokens: list[str],
    **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """Given the path to an input corpus, run train a BPE tokenizer and
    output its vocabulary and merges.

    Args:
        input_path (str | os.PathLike): Path to BPE tokenizer training data.
        vocab_size (int): Total number of items in the tokenizer's vocabulary (including special tokens).
        special_tokens (list[str]): A list of string special tokens to be added to the tokenizer vocabulary.
            These strings will never be split into multiple tokens, and will always be
            kept as a single token. If these special tokens occur in the `input_path`,
            they are treated as any other string.

    Returns:
        tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
            vocab:
                The trained tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
                to bytes (token bytes)
            merges:
                BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
                representing that <token1> was merged with <token2>.
                Merges are ordered by order of creation.
    """
    with open(input_path, "r", encoding="utf-8") as f:
        corpus = f.read()
    special_tokens = special_tokens or []
    pattern = re.compile(
        r"'(?:[sdmt]|ll|ve|re)| ?\p{L}++| ?\p{N}++| ?[^\s\p{L}\p{N}]++|\s++$|\s+(?!\S)|\s"
    )
    if special_tokens:
        sorted_tokens = sorted(special_tokens, key=len, reverse=True)
        special_regex = re.compile("|".join(re.escape(tok) for tok in sorted_tokens))
    else:
        special_regex = None

    def split_special(text):
        if not special_regex:
            return [(text, False)]
        parts = []
        last = 0
        for match in special_regex.finditer(text):
            if match.start() > last:
                parts.append((text[last : match.start()], False))
            parts.append((match.group(0), True))
            last = match.end()
        if last < len(text):
            parts.append((text[last:], False))
        return parts

    word_counts = {}
    word_first_index = {}
    word_pairs_cache = {}
    word_index = 0

    def get_pairs_with_pos(word):
        pairs = []
        for i in range(len(word) - 1):
            pairs.append(((word[i], word[i + 1]), i))
        return pairs

    for segment, is_special in split_special(corpus):
        if is_special:
            continue
        for token in pattern.findall(segment):
            token_bytes = token.encode("utf-8")
            if not token_bytes:
                continue
            word = tuple(bytes([b]) for b in token_bytes)
            if word not in word_counts:
                word_counts[word] = 1
                word_first_index[word] = word_index
                word_pairs_cache[word] = get_pairs_with_pos(word)
            else:
                word_counts[word] += 1
            word_index += 1

    merges = []
    vocab_symbols = [token.encode("utf-8") for token in special_tokens]
    vocab_set = set(vocab_symbols)
    for b in range(256):
        sym = bytes([b])
        if sym not in vocab_set:
            vocab_symbols.append(sym)
            vocab_set.add(sym)

    def merge_word(word, pair):
        merged = []
        i = 0
        while i < len(word):
            if i < len(word) - 1 and word[i] == pair[0] and word[i + 1] == pair[1]:
                merged.append(word[i] + word[i + 1])
                i += 2
            else:
                merged.append(word[i])
                i += 1
        return tuple(merged)

    pair_counts = {}
    pair_first_seen = {}
    pair_to_words = {}
    for word, count in word_counts.items():
        pairs = word_pairs_cache[word]
        for pair, pos in pairs:
            pair_counts[pair] = pair_counts.get(pair, 0) + count
            pair_to_words.setdefault(pair, set()).add(word)
            candidate = (word_first_index[word], pos)
            existing = pair_first_seen.get(pair)
            if existing is None or candidate < existing:
                pair_first_seen[pair] = candidate

    while len(vocab_symbols) < vocab_size and pair_counts:
        best_pair = None
        best_count = -1
        for pair, count in pair_counts.items():
            if count > best_count or (count == best_count and (best_pair is None or pair > best_pair)):
                best_pair = pair
                best_count = count
        if best_pair is None:
            break
        merges.append(best_pair)
        merged_symbol = best_pair[0] + best_pair[1]
        if merged_symbol not in vocab_set:
            vocab_symbols.append(merged_symbol)
            vocab_set.add(merged_symbol)

        affected_words = list(pair_to_words.get(best_pair, set()))
        pairs_to_recompute = set()
        for word in affected_words:
            count = word_counts.pop(word)
            word_index = word_first_index.pop(word)
            pairs = word_pairs_cache.pop(word)
            for pair, pos in pairs:
                new_count = pair_counts[pair] - count
                if new_count <= 0:
                    pair_counts.pop(pair, None)
                    pair_first_seen.pop(pair, None)
                    pair_to_words.pop(pair, None)
                else:
                    pair_counts[pair] = new_count
                    pair_to_words[pair].discard(word)
                    if pair_first_seen.get(pair) and pair_first_seen[pair][0] == word_index:
                        pairs_to_recompute.add(pair)

            new_word = merge_word(word, best_pair)
            if new_word in word_counts:
                word_counts[new_word] += count
                if word_first_index[new_word] > word_index:
                    word_first_index[new_word] = word_index
                    pairs_to_recompute.update({pair for pair, _ in word_pairs_cache[new_word]})
            else:
                word_counts[new_word] = count
                word_first_index[new_word] = word_index
                word_pairs_cache[new_word] = get_pairs_with_pos(new_word)

            pairs_new = word_pairs_cache[new_word]
            for pair, pos in pairs_new:
                pair_counts[pair] = pair_counts.get(pair, 0) + count
                pair_to_words.setdefault(pair, set()).add(new_word)
                candidate = (word_first_index[new_word], pos)
                existing = pair_first_seen.get(pair)
                if existing is None or candidate < existing:
                    pair_first_seen[pair] = candidate

        for pair in pairs_to_recompute:
            if pair not in pair_counts:
                continue
            words_with_pair = pair_to_words.get(pair)
            if not words_with_pair:
                pair_first_seen.pop(pair, None)
                continue
            best_candidate = None
            for other_word in words_with_pair:
                other_index = word_first_index[other_word]
                for other_pair, pos in word_pairs_cache[other_word]:
                    if other_pair == pair:
                        candidate = (other_index, pos)
                        if best_candidate is None or candidate < best_candidate:
                            best_candidate = candidate
                        break
            pair_first_seen[pair] = best_candidate

    vocab = {i: sym for i, sym in enumerate(vocab_symbols[:vocab_size])}
    return vocab, merges