content: import torch
  import torch.nn as nn
  import math
  import math
  from typing import Optional
  from typing import Callable, Iterable, Sequence, Union
  
  import torch
  
  def checkpoint(
      func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]],
      inputs: Sequence[torch.Tensor],
      params: Iterable[torch.Tensor],
      flag: bool,
  ):
      """
      Evaluate a function without caching intermediate activations, allowing for
      reduced memory at the expense of extra compute in the backward pass.
      :param func: the function to evaluate.
      :param inputs: the argument sequence to pass to `func`.
      :param params: a sequence of parameters `func` depends on but does not
                     explicitly take as arguments.
      :param flag: if False, disable gradient checkpointing.
      """
      if flag:
          args = tuple(inputs) + tuple(params)
          return CheckpointFunction.apply(func, len(inputs), *args)
      else:
          return func(*inputs)
  
  
  class CheckpointFunction(torch.autograd.Function):
      @staticmethod
      def forward(ctx, run_function, length, *args):
          ctx.run_function = run_function
          ctx.input_tensors = list(args[:length])
          ctx.input_params = list(args[length:])
          with torch.no_grad():
              output_tensors = ctx.run_function(*ctx.input_tensors)
          return output_tensors
  
      @staticmethod
      def backward(ctx, *output_grads):
          ctx.input_tensors = [x.detach().requires_grad_(True)
                               for x in ctx.input_tensors]
          with torch.enable_grad():
              # Fixes a bug where the first op in run_function modifies the
              # Tensor storage in place, which is not allowed for detach()'d
              # Tensors.
              shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
              output_tensors = ctx.run_function(*shallow_copies)
          input_grads = torch.autograd.grad(
              output_tensors,
              ctx.input_tensors + ctx.input_params,
              output_grads,
              allow_unused=True,
          )
          del ctx.input_tensors
          del ctx.input_params
          del output_tensors
          return (None, None) + input_grads
  
  
  def init_linear(l, stddev):
      nn.init.normal_(l.weight, std=stddev)
      if l.bias is not None:
          nn.init.constant_(l.bias, 0.0)
  
  class MLP(nn.Module):
      def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int, init_scale: float):
          super().__init__()
          self.width = width
          self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
          self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
          self.gelu = nn.GELU()
          init_linear(self.c_fc, init_scale)
          init_linear(self.c_proj, init_scale)
  
      def forward(self, x):
          return self.c_proj(self.gelu(self.c_fc(x)))
  
  class QKVMultiheadCrossAttention(nn.Module):
      def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_data: int):
          super().__init__()
          self.device = device
          self.dtype = dtype
          self.heads = heads
          self.n_data = n_data
  
      def forward(self, q, kv):
          _, n_ctx, _ = q.shape
          bs, n_data, width = kv.shape
          attn_ch = width // self.heads // 2
          scale = 1 / math.sqrt(math.sqrt(attn_ch))
          q = q.view(bs, n_ctx, self.heads, -1)
          kv = kv.view(bs, n_data, self.heads, -1)
          k, v = torch.split(kv, attn_ch, dim=-1)
          weight = torch.einsum(
              "bthc,bshc->bhts", q * scale, k * scale
          )  # More stable with f16 than dividing afterwards
          wdtype = weight.dtype
          weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
          return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
  
  
  
  class QKVMultiheadAttention(nn.Module):
      def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int):
          super().__init__()
          self.device = device
          self.dtype = dtype
          self.heads = heads
          self.n_ctx = n_ctx
  
      def forward(self, qkv):
          bs, n_ctx, width = qkv.shape
          attn_ch = width // self.heads // 3
          scale = 1 / math.sqrt(math.sqrt(attn_ch))
          qkv = qkv.view(bs, n_ctx, self.heads, -1)
          q, k, v = torch.split(qkv, attn_ch, dim=-1)
          weight = torch.einsum(
              "bthc,bshc->bhts", q * scale, k * scale
          )  # More stable with f16 than dividing afterwards
          wdtype = weight.dtype
          weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
          return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
  
  
  
  class MultiheadCrossAttention(nn.Module):
      def __init__(
          self,
          *,
          device: torch.device,
          dtype: torch.dtype,
          n_data: int,
          width: int,
          heads: int,
          init_scale: float,
          data_width: Optional[int] = None,
      ):
          super().__init__()
          self.n_data = n_data
          self.width = width
          self.heads = heads
          self.data_width = width if data_width is None else data_width
          self.c_q = nn.Linear(width, width, device=device, dtype=dtype)
          self.c_kv = nn.Linear(self.data_width, width * 2,
                                device=device, dtype=dtype)
          self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
          self.attention = QKVMultiheadCrossAttention(
              device=device, dtype=dtype, heads=heads, n_data=n_data
          )
          init_linear(self.c_q, init_scale)
          init_linear(self.c_kv, init_scale)
          init_linear(self.c_proj, init_scale)
  
      def forward(self, x, data):
          x = self.c_q(x)
          data = self.c_kv(data)
          x = checkpoint(self.attention, (x, data), (), True)
          x = self.c_proj(x)
          return x
  
  
  class MultiheadAttention(nn.Module):
      def __init__(
          self,
          *,
          device: torch.device,
          dtype: torch.dtype,
          n_ctx: int,
          width: int,
          heads: int,
          init_scale: float,
      ):
          super().__init__()
          self.n_ctx = n_ctx
          self.width = width
          self.heads = heads
          self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
          self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
          self.attention = QKVMultiheadAttention(device=device, dtype=dtype, heads=heads, n_ctx=n_ctx)
          init_linear(self.c_qkv, init_scale)
          init_linear(self.c_proj, init_scale)
  
      def forward(self, x):
          x = self.c_qkv(x)
          x = checkpoint(self.attention, (x,), (), True)
          x = self.c_proj(x)
          return x
  
  
  class ResidualTransformerBlock(nn.Module):
      def __init__(
          self,
          *,
          device: torch.device,
          dtype: torch.dtype,
          n_data: int,
          width: int,
          heads: int,
          data_width: Optional[int] = None,
          init_scale: float = 1.0,
      ):
          super().__init__()
  
          if data_width is None:
              data_width = width
  
          self.attn_cross = MultiheadCrossAttention(
              device=device,
              dtype=dtype,
              n_data=n_data,
              width=width,
              heads=heads,
              data_width=data_width,
              init_scale=init_scale,
          )
          self.attn_self = MultiheadAttention(
              device=device,
              dtype=dtype,
              n_ctx=n_data,
              width=width,
              heads=heads,
              init_scale=init_scale,
          )
          self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
          self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
          self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
          self.mlp = MLP(device=device, dtype=dtype,
                         width=width, init_scale=init_scale)
          self.ln_4 = nn.LayerNorm(width, device=device, dtype=dtype)
  
      def forward(self, x: torch.Tensor, data: torch.Tensor):
          x = x + self.attn_cross(self.ln_1(x), self.ln_2(data))
          x = x + self.attn_self(self.ln_3(x))
          x = x + self.mlp(self.ln_4(x))
          return x