vq-bet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c5f6a94a-b1cd-44ff-826c-f8ca67f94d5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tinygrad\n",
    "from tinygrad import Tensor, nn, dtypes\n",
    "from tinygrad.dtype import ConstType\n",
    "from tinygrad.nn.optim import Optimizer\n",
    "import einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2443d738-61f9-4ea3-a7f6-209bd591a0be",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import warnings\n",
    "from collections import deque\n",
    "from typing import Callable, List, Union, Literal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911d848b-c5f5-42e0-9835-7f9c051a198c",
   "metadata": {},
   "source": [
    "# Vector-Quantized Behavior Transformers (VQ-BeTs)\n",
    "\n",
    "Vector-Quantized Behavior Transformers (VQ-BeTs) represent a significant advancement in behavior modeling and generation. They address key limitations of previous approaches like Behavior Transformers (BeT) by introducing a hierarchical vector quantization module for tokenizing continuous actions.\n",
    "\n",
    "Why VQ-BeTs?\n",
    "\n",
    "1. Handling Multimodal Action Distributions: Real-world behaviors often have multiple valid actions for a given state. VQ-BeTs can capture this multimodality more effectively than methods using simple discretization like k-means clustering.\n",
    "\n",
    "2. Scalability: Unlike k-means, the hierarchical vector quantization in VQ-BeTs scales well to high-dimensional action spaces and long sequences, making it suitable for complex, long-horizon tasks.\n",
    "\n",
    "3. Gradient Information: The vector quantization process allows for gradient flow, enabling more effective training compared to non-differentiable discretization methods.\n",
    "\n",
    "4. Versatility: VQ-BeTs can handle both conditional and unconditional behavior generation, as well as partial observations, making them applicable to a wide range of tasks.\n",
    "\n",
    "5. Improved Performance: Across various environments, including simulated manipulation, autonomous driving, and robotics, VQ-BeTs have shown superior performance compared to state-of-the-art models like BeT and Diffusion Policies.\n",
    "\n",
    "6. Faster Inference: VQ-BeTs demonstrate significantly faster inference times compared to some alternatives, such as Diffusion Policies."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d8aba5-c811-4af9-9ffa-7d2845cef2c6",
   "metadata": {},
   "source": [
    "# Implementing a Multi-Layer Perceptron (MLP) in tinygrad\n",
    "\n",
    "In this cell, we'll define a simple Multi-Layer Perceptron (MLP) class using tinygrad. An MLP is a type of feedforward neural network consisting of multiple layers of neurons. \n",
    "\n",
    "Our MLP class will:\n",
    "1. Accept an input dimension and a list of hidden layer dimensions\n",
    "2. Construct a sequence of linear layers with ReLU activations between them\n",
    "3. Use the last dimension in the list as the output dimension\n",
    "\n",
    "Key points:\n",
    "- We use `nn.Linear` for the fully connected layers\n",
    "- ReLU activation is applied after each hidden layer\n",
    "- The final layer doesn't have an activation function, allowing for flexible use in various tasks\n",
    "\n",
    "Let's implement the MLP class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b85e8ddf-6c2c-4dd6-98d8-1c7f6febe033",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP:\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels: int,\n",
    "        hidden_channels: List[int],\n",
    "    ):\n",
    "        layers = []\n",
    "        in_dim = in_channels\n",
    "        for hidden_dim in hidden_channels[:-1]:\n",
    "            layers.append(nn.Linear(in_dim, hidden_dim))\n",
    "            layers.append(Tensor.relu)\n",
    "            in_dim = hidden_dim\n",
    "        self.fc = nn.Linear(in_dim, hidden_channels[-1])\n",
    "        self.layers = layers\n",
    "    def __call__(self, x:Tensor) -> Tensor:\n",
    "        return self.fc(x.sequential(self.layers))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70806f08-fc58-442e-9a76-ecc4ece322fc",
   "metadata": {},
   "source": [
    "# Implementing Focal Loss in tinygrad\n",
    "\n",
    "Next, we'll implement the Focal Loss function, which is particularly useful for dealing with class imbalance in classification tasks. This implementation is adapted from the miniBET project.\n",
    "\n",
    "Focal Loss modifies standard cross-entropy loss to down-weight easy examples and focus training on hard negatives. It introduces a focusing parameter γ (gamma) that adjusts the rate at which easy examples are down-weighted.\n",
    "\n",
    "Key features of this implementation:\n",
    "1. Supports both 2D and 3D input tensors\n",
    "2. Uses log_softmax for numerical stability\n",
    "3. Allows for mean or sum reduction of the loss\n",
    "\n",
    "Let's examine the FocalLoss class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48881053-e9ec-41d6-a472-77249e7b662c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FocalLoss:\n",
    "    \"\"\"\n",
    "    From https://github.com/notmahi/miniBET/blob/main/behavior_transformer/bet.py\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, gamma: float = 0, size_average: bool = True):\n",
    "        super().__init__()\n",
    "        self.gamma = gamma\n",
    "        self.size_average = size_average\n",
    "\n",
    "    def __call__(self, x: Tensor, target: Tensor) -> Tensor:\n",
    "        if len(x.shape) == 3:\n",
    "            N, T, _ = x.shape\n",
    "            logpt = x.log_softmax(axis=-1)\n",
    "            logpt = logpt.gather(-1, target.view(N, T, 1)).view(N, T)\n",
    "        elif len(x.shape) == 2:\n",
    "            logpt = x.log_softmax(axis=-1)\n",
    "            logpt = logpt.gather(-1, target.view(-1, 1)).view(-1)\n",
    "        pt = logpt.exp()\n",
    "\n",
    "        loss = -1 * (1 - pt).pow(self.gamma) * logpt\n",
    "        return loss.mean() if self.size_average else loss.sum()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6cd46a9-6f6c-48c9-9b34-d5bc3bd4fa11",
   "metadata": {},
   "source": [
    "# Implementing the cdist Function for VQ-BeTs\n",
    "\n",
    "In the context of Vector-Quantized Behavior Transformers (VQ-BeTs), calculating distances between vectors is a crucial operation, particularly in the vector quantization process. The `cdist` function computes the pairwise distance between two sets of vectors, which is essential for finding the nearest codebook vectors during quantization.\n",
    "\n",
    "Why is this important for VQ-BeTs?\n",
    "\n",
    "1. Vector Quantization: VQ-BeTs use vector quantization to discretize continuous action spaces. This process involves finding the nearest codebook vector for each input vector, which requires computing distances.\n",
    "\n",
    "2. Codebook Updates: During training, the codebook vectors are updated based on their proximity to input vectors. The `cdist` function helps in determining these proximities efficiently.\n",
    "\n",
    "3. Efficiency: By implementing `cdist` using matrix operations, we can compute all pairwise distances in parallel, which is much faster than iterating over vectors individually. This is crucial for the performance of VQ-BeTs, especially when dealing with large action spaces or long sequences.\n",
    "\n",
    "4. Flexibility: While we're primarily using Euclidean distance (p=2), the function is designed to potentially accommodate other distance metrics, which could be useful for experimenting with different quantization schemes.\n",
    "\n",
    "Let's examine the implementation of the `cdist` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a7d5f72-8048-4ad6-8790-6dde86260fe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cdist(x1, x2, p=2.0):\n",
    "    # Ensure inputs are Tensors\n",
    "    x1 = x1 if isinstance(x1, Tensor) else Tensor(x1)\n",
    "    x2 = x2 if isinstance(x2, Tensor) else Tensor(x2)\n",
    "    \n",
    "    # Compute squared Euclidean distance\n",
    "    x1_norm = x1.square().sum(axis=-1, keepdim=True)\n",
    "    x2_norm = x2.square().sum(axis=-1, keepdim=True)\n",
    "    \n",
    "    cross_term = x1 @ x2.transpose()\n",
    "    \n",
    "    dist = x1_norm + x2_norm.transpose() - 2 * cross_term\n",
    "    \n",
    "    # Handle floating point errors (negative distances)\n",
    "    dist = dist.maximum(Tensor.zeros_like(dist))\n",
    "    \n",
    "    # For p=2 (Euclidean distance), take the square root\n",
    "    if p == 2:\n",
    "        return dist.sqrt()\n",
    "    else:\n",
    "        # For other p-norms, we'd need to implement a different calculation\n",
    "        raise NotImplementedError(\"Only p=2 (Euclidean distance) is currently implemented\")\n",
    "\n",
    "Tensor.cdist = cdist\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2be079d2-962b-45b3-a231-721b855b1385",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opened device METAL from pid:8264\n",
      "opened device NPY from pid:8264\n",
      "\u001b[32m*** METAL      1\u001b[0m \u001b[33mcopy       24,   METAL <- NPY    \u001b[0m        mem  0.00 GB tm     27.96us/     0.03ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL      2\u001b[0m \u001b[33mcopy       16,   METAL <- NPY    \u001b[0m        mem  0.00 GB tm     12.42us/     0.04ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "hc     : \u001b[36m3\u001b[0m \u001b[35m2\u001b[0m                            :     3.54 us  <  beam2  : \u001b[35m2\u001b[0m \u001b[33m3\u001b[0m                            :     7.00 us\n",
      "\u001b[32m*** METAL      3\u001b[0m r_\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                    mem  0.00 GB tm      7.71us/     0.05ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['square', 'sum']\n",
      "hc     : \u001b[36m2\u001b[0m \u001b[35m2\u001b[0m                            :     6.46 us  <  beam2  : \u001b[36m2\u001b[0m \u001b[31m2\u001b[0m                            :     6.75 us\n",
      "\u001b[32m*** METAL      4\u001b[0m r_\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                    mem  0.00 GB tm      7.38us/     0.06ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['square', 'sum']\n",
      "beam2  : \u001b[34m3\u001b[0m \u001b[36m2\u001b[0m \u001b[31m2\u001b[0m                          :     7.12 us  <  hc     : \u001b[36m3\u001b[0m \u001b[36m2\u001b[0m \u001b[35m2\u001b[0m                          :     7.12 us\n",
      "\u001b[32m*** METAL      5\u001b[0m r_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m\u001b[0m                                  mem  0.00 GB tm      7.87us/     0.06ms (     0.01 GFLOPS    0.0|0.0     GB/s) ['__add__', '__matmul__', '__rmul__', '__sub__', 'maximum', 'sqrt']\n",
      "opened device CLANG from pid:8264\n",
      "\u001b[32m*** CLANG      6\u001b[0m \u001b[33mcopy       24,   CLANG <- METAL  \u001b[0m        mem  0.00 GB tm     38.96us/     0.10ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "[[3.1192703  2.0958931 ]\n",
      " [2.713841   3.8321726 ]\n",
      " [2.2830095  0.37910157]]\n"
     ]
    }
   ],
   "source": [
    "x1 = Tensor([[0.9041, 0.0196], [-0.3108, -2.4423], [-0.4821, 1.059]])\n",
    "x2 = Tensor([[-2.1763, -0.4713], [-0.6986, 1.3702]])\n",
    "\n",
    "result = cdist(x1, x2)\n",
    "print(result.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed25c6e-196e-4430-a249-cdd515806eca",
   "metadata": {},
   "source": [
    "# Utility Functions for VQ-BeTs Implementation\n",
    "\n",
    "In implementing Vector-Quantized Behavior Transformers (VQ-BeTs), we rely on a set of utility functions that handle various operations crucial to the model's functionality. These functions cover a range of tasks from basic tensor operations to more complex procedures like k-means clustering and Gumbel-Softmax sampling.\n",
    "\n",
    "Why are these utility functions important for VQ-BeTs?\n",
    "\n",
    "1. Tensor Operations: Functions like `unbind`, `cdist`, and `batched_embedding` provide efficient ways to manipulate tensors, which is essential for handling the high-dimensional data in VQ-BeTs.\n",
    "\n",
    "2. Vector Quantization: The `kmeans` function is central to the vector quantization process, allowing us to create and update codebooks.\n",
    "\n",
    "3. Sampling and Noise: Functions like `gumbel_noise` and `gumbel_sample` are crucial for the stochastic aspects of VQ-BeTs, enabling the model to handle uncertainty and explore the action space effectively.\n",
    "\n",
    "4. Initialization and Normalization: Functions like `normal_`, `uniform_init`, and `laplace_smoothing` help in properly initializing and normalizing model parameters, which is critical for stable training of VQ-BeTs.\n",
    "\n",
    "5. Loss Computation: The `orthogonal_loss_fn` is used to encourage diversity in the codebook, which is important for capturing a wide range of behaviors.\n",
    "\n",
    "6. Efficient Computation: Many of these functions (e.g., `batched_bincount`, `batched_sample_vectors`) are designed to operate efficiently on batches of data, which is crucial for training VQ-BeTs on large datasets.\n",
    "\n",
    "These utility functions form the backbone of our VQ-BeTs implementation, enabling efficient computation and providing the necessary tools for the model's core operations. Let's examine each of these functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ac0f3c2d-1982-4224-be15-8198b3f31271",
   "metadata": {},
   "outputs": [],
   "source": [
    "# General utils\n",
    "\n",
    "def noop(*args, **kwargs):\n",
    "    pass\n",
    "\n",
    "\n",
    "def identity(t):\n",
    "    return t\n",
    "\n",
    "def unbind(x: Tensor):\n",
    "    return tuple(x[i] for i in range(x.shape[0]))\n",
    "\n",
    "def cdist(x, y):\n",
    "    x2 = reduce(x.square(), \"b n d -> b n\", \"sum\")\n",
    "    y2 = reduce(y.square(), \"b n d -> b n\", \"sum\")\n",
    "    xy = Tensor.einsum(\"b i d, b j d -> b i j\", x, y) * -2\n",
    "    return (rearrange(x2, \"b i -> b i 1\") + rearrange(y2, \"b j -> b 1 j\") + xy).sqrt()\n",
    "\n",
    "\n",
    "def log(t, eps=1e-20):\n",
    "    return t.clamp(min_=eps).log()\n",
    "\n",
    "\n",
    "def ema_inplace(old, new, decay):\n",
    "    is_mps = str(old.device).startswith('METAL')\n",
    "\n",
    "    if not is_mps:\n",
    "        old.assign(old.detach().lerp(new.detach(), 1 - decay))\n",
    "    else:\n",
    "        old.assign(old.detach().mul(decay).add(new.detach() * (1 - decay)))\n",
    "\n",
    "from einops import pack, rearrange, reduce, repeat, unpack\n",
    "\n",
    "def pack_one(t, pattern):\n",
    "    return pack([t], pattern)\n",
    "\n",
    "\n",
    "def unpack_one(t, ps, pattern):\n",
    "    return unpack(t, ps, pattern)[0]\n",
    "\n",
    "def normal_(tensor: Tensor, mean: float = 0.0, std: float = 1.0) -> Tensor:\n",
    "    \"\"\"\n",
    "    Fills the input Tensor with values drawn from the normal distribution N(mean, std^2).\n",
    "    \n",
    "    Args:\n",
    "    tensor (Tensor): an n-dimensional `Tensor`\n",
    "    mean (float): the mean of the normal distribution\n",
    "    std (float): the standard deviation of the normal distribution\n",
    "    \n",
    "    Returns:\n",
    "    Tensor: the modified input tensor\n",
    "    \"\"\"\n",
    "    # Create a new tensor with values drawn from a normal distribution\n",
    "    normal_tensor = Tensor.randn(*tensor.shape) * std + mean\n",
    "    \n",
    "    # In-place update of the input tensor\n",
    "    tensor.assign(normal_tensor)\n",
    "    \n",
    "    return tensor\n",
    "\n",
    "def uniform_init(*shape):\n",
    "    return Tensor.kaiming_uniform(shape)\n",
    "\n",
    "\n",
    "def gumbel_noise(t):\n",
    "    noise = Tensor.uniform(t.shape, low=0, high=1)\n",
    "    return -(-noise.log()).log()\n",
    "\n",
    "\n",
    "def gumbel_sample(\n",
    "    logits,\n",
    "    temperature=1.0,\n",
    "    stochastic=False,\n",
    "    straight_through=False,\n",
    "    reinmax=False,\n",
    "    dim=-1,\n",
    "    training=True,\n",
    "):\n",
    "    dtype, size = logits.dtype, logits.shape[dim]\n",
    "\n",
    "    if training and stochastic and temperature > 0:\n",
    "        sampling_logits = (logits / temperature) + gumbel_noise(logits)\n",
    "    else:\n",
    "        sampling_logits = logits\n",
    "\n",
    "    ind = sampling_logits.argmax(axis=dim)\n",
    "    one_hot = ind.one_hot(size).cast(dtype=dtype)\n",
    "\n",
    "    assert not (\n",
    "        reinmax and not straight_through\n",
    "    ), \"reinmax can only be turned on if using straight through gumbel softmax\"\n",
    "\n",
    "    if not straight_through or temperature <= 0.0 or not training:\n",
    "        return ind, one_hot\n",
    "\n",
    "    # use reinmax for better second-order accuracy - https://arxiv.org/abs/2304.08612\n",
    "    # algorithm 2\n",
    "\n",
    "    if reinmax:\n",
    "        π0 = logits.softmax(axis=dim)\n",
    "        π1 = (one_hot + (logits / temperature).softmax(axis=dim)) / 2\n",
    "        π1 = ((log(π1) - logits).detach() + logits).softmax(axis=1)\n",
    "        π2 = 2 * π1 - 0.5 * π0\n",
    "        one_hot = π2 - π2.detach() + one_hot\n",
    "    else:\n",
    "        π1 = (logits / temperature).softmax(axis=dim)\n",
    "        one_hot = one_hot + π1 - π1.detach()\n",
    "\n",
    "    return ind, one_hot\n",
    "\n",
    "\n",
    "def laplace_smoothing(x, n_categories, eps=1e-5, dim=-1):\n",
    "    denom = x.sum(axis=dim, keepdim=True)\n",
    "    return (x + eps) / (denom + n_categories * eps)\n",
    "\n",
    "\n",
    "def sample_vectors(samples, num): return samples[Tensor.randint((num,), low=0, high=samples.shape[0])]\n",
    "\n",
    "def batched_sample_vectors(samples, num):\n",
    "    return Tensor.stack(*[sample_vectors(sample, num) for sample in unbind(samples)], dim=0)\n",
    "\n",
    "def pad_shape(shape, size, dim=0):\n",
    "    return [size if i == dim else s for i, s in enumerate(shape)]\n",
    "\n",
    "\n",
    "def sample_multinomial(total_count, probs):\n",
    "    remainder = Tensor(1.0)\n",
    "    sample = Tensor.empty_like(probs, dtype=dtypes.long)\n",
    "\n",
    "    for i, p in enumerate(probs):\n",
    "        # my hacky way of getting binomial up in tinygrad\n",
    "        u = Tensor.rand(total_count)\n",
    "        s = (u <= (p / remainder)).sum(axis=-1)\n",
    "        sample[i] = s\n",
    "        total_count -= s\n",
    "        remainder -= p\n",
    "\n",
    "    return sample\n",
    "\n",
    "def batched_bincount(x, *, minlength):\n",
    "    batch = x.shape[0]\n",
    "    \n",
    "    # Create a zero tensor for the target\n",
    "    target = Tensor.zeros(batch, minlength)\n",
    "    \n",
    "    # Iterate over each batch\n",
    "    for i in range(batch):\n",
    "        # Get the current batch\n",
    "        x_batch = x[i]\n",
    "        \n",
    "        # Iterate over each element in the batch\n",
    "        for j in range(x_batch.shape[0]):\n",
    "            # Get the index (bin) for the current element\n",
    "            index = x_batch[j]\n",
    "            \n",
    "            # Increment the count for this index in the target tensor\n",
    "            target[i, index] += 1\n",
    "    \n",
    "    return target\n",
    "\n",
    "def scatter_tensor(self, dim:int, index:Tensor, src:Union[Tensor, ConstType], reduce:Union[None, Literal['multiply'], Literal['add']] = None):\n",
    "    \"\"\"\n",
    "    Scatters `src` values along an axis specified by `dim`.\n",
    "    apply `add` or `multiply` reduction operation with `reduce`.\n",
    "    \"\"\"\n",
    "    index, dim  = index.to(self.device), self._resolve_dim(dim)\n",
    "    if not isinstance(src, Tensor): src = Tensor(src, device=self.device, dtype=self.dtype)._broadcast_to(index.shape)\n",
    "    assert index.ndim == self.ndim == src.ndim, f\"self.ndim must equal index.ndim, {self.ndim=}, {index.ndim=}\"\n",
    "    assert all((s >= i if d != dim else True) and srcs >= i for d,(s,i,srcs) in enumerate(zip(self.shape, index.shape, src.shape))), \\\n",
    "      f\"Expected {index.shape=} to be <= {self.shape=} apart from dimension {dim} and to be <= {src.shape=}\"\n",
    "    mask = (index.unsqueeze(-1) == Tensor.arange(self.shape[dim], requires_grad=False, device=self.device)).transpose(-1, dim)\n",
    "    src = src.shrink(tuple((0, sh) for sh in index.shape)).unsqueeze(-1).transpose(-1, dim)\n",
    "    src = src.expand(tuple(self.size(i) if i == dim else None for i in range(src.ndim)))\n",
    "    src = src.pad(tuple((0, max(xs-ss, 0)) for ss, xs in zip(src.shape[:-1], self.shape)) + (None,))\n",
    "    mask = mask.pad(tuple((0, max(xs-ms, 0)) for ms, xs in zip(mask.shape[:-1], self.shape)) + (None,))\n",
    "    if reduce is None:\n",
    "      nan_masked = mask.where(src*mask, float(\"nan\"))\n",
    "      masked_src = functools.reduce(lambda x,y: y.isnan().where(x, y), nan_masked.split(1, -1))\n",
    "      return (mask.any(-1).where(masked_src.squeeze(), self)).cast(self.dtype)\n",
    "    if reduce == \"add\": return ((mask*src).sum(-1) + self).cast(self.dtype)\n",
    "    if reduce == \"multiply\": return (mask.where(mask*src, 1).prod(-1) * self).cast(self.dtype)\n",
    "\n",
    "Tensor.scatter = scatter_tensor\n",
    "\n",
    "def gpu_radix_sort(self):\n",
    "    n = self.size(0)  # Number of elements\n",
    "    if n == 0: return self\n",
    "    num_bits = 4  # Number of bits to process per pass (chunk size)\n",
    "    max_digit = 1 << num_bits  # Base = 2^num_bits\n",
    "    \n",
    "    sorted_output = self  # Initial unsorted tensor\n",
    "\n",
    "    for pass_idx in range(32 // num_bits):  # Assuming 32-bit integers\n",
    "        # Extract digits for the current pass\n",
    "        digits = (sorted_output >> (num_bits * pass_idx)) & (max_digit - 1)\n",
    "        print(f\"Pass {pass_idx}, Digits: {digits.numpy()}\")\n",
    "\n",
    "        # Create histogram (count occurrences of each digit)\n",
    "        histogram = Tensor.zeros(max_digit)\n",
    "        histogram = histogram.scatter(0, digits, Tensor.ones_like(digits), 'add')\n",
    "        print(f\"Pass {pass_idx}, Histogram: {histogram.numpy()}\")\n",
    "\n",
    "        # Compute prefix sum (cumulative count)\n",
    "        prefix_sum = histogram.cumsum(0)\n",
    "        print(f\"Pass {pass_idx}, Prefix Sum: {prefix_sum.numpy()}\")\n",
    "\n",
    "        # Initialize positions using prefix sum\n",
    "        positions = (prefix_sum - histogram).int()  # Corrected positions\n",
    "        print(f\"Pass {pass_idx}, Initial Positions: {positions.numpy()}\")\n",
    "\n",
    "        # Scatter into sorted output\n",
    "        new_output = Tensor.zeros_like(sorted_output).contiguous()\n",
    "        for i in range(n):\n",
    "            digit = digits[i]\n",
    "            position = positions[digit]\n",
    "            new_output[position] = sorted_output[i]\n",
    "            positions[digit] += 1  # Increment position after assignment\n",
    "\n",
    "        sorted_output = new_output\n",
    "        print(f\"Pass {pass_idx}, Sorted Output: {sorted_output.numpy()}\")\n",
    "\n",
    "    return sorted_output\n",
    "\n",
    "Tensor.sort = gpu_radix_sort\n",
    "\n",
    "def unique(self):\n",
    "    # Step 1: Sort the tensor\n",
    "    sorted_tensor = self.sort()\n",
    "    \n",
    "    # Step 2: Find indices where values change\n",
    "    # Prepend a value outside the range to ensure the first element is always included\n",
    "    diffs = Tensor.cat(Tensor([True], requires_grad=False), sorted_tensor[1:] != sorted_tensor[:-1])\n",
    "    \n",
    "    # Step 3: Index the unique elements\n",
    "    unique_with_zero = diffs.where(sorted_tensor, Tensor([0], requires_grad=False))\n",
    "    result = Tensor.empty(unique_with_zero.shape, dtype=sorted_tensor.dtype)\n",
    "    idx = 0\n",
    "    for i in range(unique_with_zero.shape[0]):\n",
    "        if unique_with_zero[i] != 0:\n",
    "            result[idx] = unique_with_zero[i]\n",
    "            idx += 1\n",
    "    \n",
    "    return result[:idx]\n",
    "\n",
    "Tensor.unique = unique\n",
    "\n",
    "def kmeans(\n",
    "    samples,\n",
    "    num_clusters,\n",
    "    num_iters=10,\n",
    "    sample_fn=batched_sample_vectors,\n",
    "    all_reduce_fn=noop,\n",
    "):\n",
    "    num_codebooks, dim, dtype = (\n",
    "        samples.shape[0],\n",
    "        samples.shape[-1],\n",
    "        samples.dtype,\n",
    "    )\n",
    "\n",
    "    means = sample_fn(samples, num_clusters)\n",
    "\n",
    "    for _ in range(num_iters):\n",
    "        dists = -Tensor.cdist(samples, means, p=2)\n",
    "        \n",
    "        # Replace argmax\n",
    "        buckets = dists.argmax(axis=-1)\n",
    "\n",
    "        bins = batched_bincount(buckets, minlength=num_clusters)\n",
    "        all_reduce_fn(bins)\n",
    "\n",
    "        zero_mask = (bins == 0)\n",
    "        bins_min_clamped = bins.masked_fill(zero_mask, 1)\n",
    "\n",
    "        new_means = Tensor.zeros((num_codebooks, num_clusters, dim), dtype=dtype)\n",
    "        new_means.scatter(1, repeat(buckets, \"h n -> h n d\", d=dim), samples, 'add')\n",
    "        new_means = new_means / rearrange(bins_min_clamped, \"... -> ... 1\")\n",
    "        all_reduce_fn(new_means)\n",
    "\n",
    "        means = rearrange(zero_mask, \"... -> ... 1\").where(means, new_means)\n",
    "\n",
    "    return means, bins\n",
    "\n",
    "\n",
    "def batched_embedding(indices, embeds):\n",
    "    batch, dim = indices.shape[1], embeds.shape[-1]\n",
    "    indices = repeat(indices, \"h b n -> h b n d\", d=dim)\n",
    "    embeds = repeat(embeds, \"h c d -> h b c d\", b=batch)\n",
    "    return embeds.gather(2, indices)\n",
    "\n",
    "\n",
    "def orthogonal_loss_fn(t):\n",
    "    # eq (2) from https://arxiv.org/abs/2112.00384\n",
    "    h, n = t.shape[:2]\n",
    "    normed_codes = F.normalize(t, p=2, dim=-1)\n",
    "    cosine_sim = Tensor.einsum(\"h i d, h j d -> h i j\", normed_codes, normed_codes)\n",
    "    return cosine_sim.square().sum() / (h * n.square()) - (1 / n)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e1814ae-e5af-4a6c-ab17-835cfc015676",
   "metadata": {},
   "source": [
    "# EuclideanCodebook: Vector Quantization for VQ-BeTs\n",
    "\n",
    "The `EuclideanCodebook` class is a sophisticated implementation of vector quantization, a key component in Vector-Quantized Behavior Transformers (VQ-BeTs). This class provides a flexible and powerful mechanism for discretizing continuous action spaces.\n",
    "\n",
    "Key Features of the EuclideanCodebook:\n",
    "\n",
    "1. **Multi-Codebook Support**: Can handle multiple codebooks, allowing for more expressive representation of complex behaviors.\n",
    "\n",
    "2. **Initialization Strategies**:\n",
    "   - Supports both uniform and k-means initialization of codebooks\n",
    "   - Flexible initialization options for different use cases\n",
    "\n",
    "3. **Adaptive Codebook Management**:\n",
    "   - Exponential Moving Average (EMA) updates for codebook vectors\n",
    "   - Mechanism to detect and replace \"dead\" (unused) codebook entries\n",
    "   - Optional learnable codebook\n",
    "\n",
    "4. **Advanced Sampling Techniques**:\n",
    "   - Gumbel-Softmax sampling for stochastic vector selection\n",
    "   - Temperature-controlled sampling to balance exploration and exploitation\n",
    "\n",
    "5. **Affine Transformation Support**:\n",
    "   - Optional affine parameter tracking\n",
    "   - Batch and codebook-level mean and variance normalization\n",
    "\n",
    "Why is this important for VQ-BeTs?\n",
    "\n",
    "- **Representation Learning**: Enables learning discrete representations of continuous action spaces\n",
    "- **Computational Efficiency**: Reduces the dimensionality of action representations\n",
    "- **Behavior Modeling**: Captures complex, multi-modal action distributions\n",
    "\n",
    "Let's dive into the implementation of the EuclideanCodebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "768b1315-9277-4521-88b4-d6226be84c52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "class EuclideanCodebook:\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim,\n",
    "        codebook_size,\n",
    "        num_codebooks=1,\n",
    "        kmeans_init=False,\n",
    "        kmeans_iters=10,\n",
    "        sync_kmeans=True,\n",
    "        decay=0.8,\n",
    "        eps=1e-5,\n",
    "        threshold_ema_dead_code=2,\n",
    "        reset_cluster_size=None,\n",
    "        use_ddp=False,\n",
    "        learnable_codebook=False,\n",
    "        gumbel_sample=gumbel_sample,\n",
    "        sample_codebook_temp=1.0,\n",
    "        ema_update=True,\n",
    "        affine_param=False,\n",
    "        sync_affine_param=False,\n",
    "        affine_param_batch_decay=0.99,\n",
    "        affine_param_codebook_decay=0.9,\n",
    "    ):\n",
    "        self.transform_input = identity\n",
    "\n",
    "        self.decay = decay\n",
    "        self.ema_update = ema_update\n",
    "\n",
    "        init_fn = Tensor.kaiming_uniform if not kmeans_init else Tensor.zeros\n",
    "        embed = init_fn(num_codebooks, codebook_size, dim, requires_grad=learnable_codebook)\n",
    "\n",
    "        self.codebook_size = codebook_size\n",
    "        self.num_codebooks = num_codebooks\n",
    "\n",
    "        self.kmeans_iters = kmeans_iters\n",
    "        self.eps = eps\n",
    "        self.threshold_ema_dead_code = threshold_ema_dead_code\n",
    "        self.reset_cluster_size = (\n",
    "            reset_cluster_size if (reset_cluster_size is not None) else threshold_ema_dead_code\n",
    "        )\n",
    "\n",
    "        assert callable(gumbel_sample)\n",
    "        self.gumbel_sample = gumbel_sample\n",
    "        self.sample_codebook_temp = sample_codebook_temp\n",
    "\n",
    "        assert not (\n",
    "            use_ddp and num_codebooks > 1 and kmeans_init\n",
    "        ), \"kmeans init is not compatible with multiple codebooks in distributed environment for now\"\n",
    "\n",
    "        # we aren't dealing with distributed work for now. Dear god no no no no\n",
    "        self.sample_fn = batched_sample_vectors\n",
    "        self.kmeans_all_reduce_fn = noop\n",
    "        self.all_reduce_fn = noop\n",
    "\n",
    "        self.initted = not kmeans_init\n",
    "        self.cluster_size = Tensor.zeros((num_codebooks, codebook_size), requires_grad=False)\n",
    "        self.embed_avg = embed.detach()\n",
    "\n",
    "        self.learnable_codebook = learnable_codebook\n",
    "        self.embed = embed\n",
    "        self.embed.requires_grad = learnable_codebook\n",
    "\n",
    "        # affine related params\n",
    "\n",
    "        self.affine_param = affine_param\n",
    "        self.sync_affine_param = sync_affine_param\n",
    "\n",
    "        if not affine_param:\n",
    "            return\n",
    "\n",
    "        self.affine_param_batch_decay = affine_param_batch_decay\n",
    "        self.affine_param_codebook_decay = affine_param_codebook_decay\n",
    "\n",
    "        self.batch_mean = Tensor(None, requires_grad=False)\n",
    "        self.batch_variance = Tensor(None, requires_grad=False)\n",
    "\n",
    "        self.codebook_mean_needs_init = Tensor([True], requires_grad=False)\n",
    "        self.codebook_mean = Tensor.empty(num_codebooks, 1, dim, requires_grad=False)\n",
    "        self.codebook_variance_needs_init = Tensor([True], requires_grad=False)\n",
    "        self.codebook_variance = Tensor.empty(num_codebooks, 1, dim, requires_grad=False)\n",
    "\n",
    "    # Don't jit this\n",
    "    def init_embed_(self, data, mask=None):\n",
    "        if self.initted:\n",
    "            return\n",
    "\n",
    "        if mask is not None:\n",
    "            c = data.shape[0]\n",
    "            data = rearrange(data[mask], \"(c n) d -> c n d\", c=c)\n",
    "\n",
    "        embed, cluster_size = kmeans(\n",
    "            data,\n",
    "            self.codebook_size,\n",
    "            self.kmeans_iters,\n",
    "            sample_fn=self.sample_fn,\n",
    "            all_reduce_fn=self.kmeans_all_reduce_fn,\n",
    "        )\n",
    "\n",
    "        embed_sum = embed * rearrange(cluster_size, \"... -> ... 1\")\n",
    "\n",
    "        self.embed = self.embed.replace(embed)\n",
    "        self.embed_avg = self.embed_avg.replace(embed_sum)\n",
    "        self.cluster_size = self.cluster_size.replace(cluster_size)\n",
    "        self.initted = True\n",
    "\n",
    "    # don't jit\n",
    "    def update_with_decay(self, buffer_name, new_value, decay):\n",
    "        old_value = getattr(self, buffer_name)\n",
    "\n",
    "        needs_init = getattr(self, buffer_name + \"_needs_init\", False)\n",
    "\n",
    "        if needs_init:\n",
    "            setattr(self, buffer_name + \"_needs_init\", Tensor([False], requires_grad=False))\n",
    "\n",
    "        if not (old_value is not None) or needs_init:\n",
    "            setattr(self, buffer_name, new_value.detach())\n",
    "\n",
    "            return\n",
    "\n",
    "        value = old_value * decay + new_value.detach() * (1 - decay)\n",
    "        setattr(self, buffer_name, value)\n",
    "\n",
    "    # don't jit\n",
    "    def update_affine(self, data, embed, mask=None):\n",
    "        assert self.affine_param\n",
    "\n",
    "        # don't use bessel correction\n",
    "        var_fn = partial(Tensor.var, correction=0)\n",
    "\n",
    "        # calculate codebook mean and variance\n",
    "\n",
    "        embed = rearrange(embed, \"h ... d -> h (...) d\")\n",
    "\n",
    "        if Tensor.training:\n",
    "            self.update_with_decay(\n",
    "                \"codebook_mean\",\n",
    "                reduce(embed, \"h n d -> h 1 d\", \"mean\"),\n",
    "                self.affine_param_codebook_decay,\n",
    "            )\n",
    "            self.update_with_decay(\n",
    "                \"codebook_variance\",\n",
    "                reduce(embed, \"h n d -> h 1 d\", var_fn),\n",
    "                self.affine_param_codebook_decay,\n",
    "            )\n",
    "\n",
    "        # prepare batch data, which depends on whether it has masking\n",
    "\n",
    "        data = rearrange(data, \"h ... d -> h (...) d\")\n",
    "\n",
    "        if mask is not None:\n",
    "            c = data.shape[0]\n",
    "            data = rearrange(data[mask], \"(c n) d -> c n d\", c=c)\n",
    "\n",
    "        # calculate batch mean and variance\n",
    "\n",
    "        if not self.sync_affine_param:\n",
    "            self.update_with_decay(\n",
    "                \"batch_mean\",\n",
    "                reduce(data, \"h n d -> h 1 d\", \"mean\"),\n",
    "                self.affine_param_batch_decay,\n",
    "            )\n",
    "            self.update_with_decay(\n",
    "                \"batch_variance\",\n",
    "                reduce(data, \"h n d -> h 1 d\", var_fn),\n",
    "                self.affine_param_batch_decay,\n",
    "            )\n",
    "            return\n",
    "\n",
    "        num_vectors = data.shape[-2]\n",
    "\n",
    "        # number of vectors, for denominator\n",
    "\n",
    "        num_vectors = Tensor([num_vectors], dtype=data.dtype)\n",
    "        # calculate distributed mean\n",
    "\n",
    "        batch_sum = reduce(data, \"h n d -> h 1 d\", \"sum\")\n",
    "        batch_mean = batch_sum / num_vectors\n",
    "\n",
    "        self.update_with_decay(\"batch_mean\", batch_mean, self.affine_param_batch_decay)\n",
    "\n",
    "        # calculate distributed variance\n",
    "\n",
    "        variance_numer = reduce((data - batch_mean).square(), \"h n d -> h 1 d\", \"sum\")\n",
    "        batch_variance = variance_numer / num_vectors\n",
    "\n",
    "        self.update_with_decay(\"batch_variance\", batch_variance, self.affine_param_batch_decay)\n",
    "\n",
    "    def replace(self, batch_samples, batch_mask):\n",
    "        for ind, (samples, mask) in enumerate(\n",
    "            zip(unbind(batch_samples), unbind(batch_mask), strict=False)\n",
    "        ):\n",
    "            if not mask.any().item():\n",
    "                continue\n",
    "\n",
    "            sampled = self.sample_fn(rearrange(samples, \"... -> 1 ...\"), mask.sum().item())\n",
    "            sampled = rearrange(sampled, \"1 ... -> ...\")\n",
    "\n",
    "            self.embed[ind][mask] = sampled\n",
    "\n",
    "            self.cluster_size[ind][mask] = self.reset_cluster_size\n",
    "            self.embed_avg[ind][mask] = sampled * self.reset_cluster_size\n",
    "\n",
    "    def expire_codes_(self, batch_samples):\n",
    "        if self.threshold_ema_dead_code == 0:\n",
    "            return\n",
    "\n",
    "        expired_codes = self.cluster_size < self.threshold_ema_dead_code\n",
    "\n",
    "        if not expired_codes.any().item():\n",
    "            return\n",
    "\n",
    "        batch_samples = rearrange(batch_samples, \"h ... d -> h (...) d\")\n",
    "        self.replace(batch_samples, batch_mask=expired_codes)\n",
    "        \n",
    "    def __call__(self, x, sample_codebook_temp=None, mask=None, freeze_codebook=False):\n",
    "        needs_codebook_dim = x.ndim < 4\n",
    "        sample_codebook_temp = (\n",
    "            sample_codebook_temp if (sample_codebook_temp is not None) else self.sample_codebook_temp\n",
    "        )\n",
    "\n",
    "        x = x.float()\n",
    "\n",
    "        if needs_codebook_dim:\n",
    "            x = rearrange(x, \"... -> 1 ...\")\n",
    "\n",
    "        flatten, ps = pack_one(x, \"h * d\")\n",
    "\n",
    "        if mask is not None:\n",
    "            mask = repeat(\n",
    "                mask,\n",
    "                \"b n -> c (b h n)\",\n",
    "                c=flatten.shape[0],\n",
    "                h=flatten.shape[-2] // (mask.shape[0] * mask.shape[1]),\n",
    "            )\n",
    "\n",
    "        self.init_embed_(flatten, mask=mask)\n",
    "\n",
    "        if self.affine_param:\n",
    "            self.update_affine(flatten, self.embed, mask=mask)\n",
    "\n",
    "        embed = self.embed if self.learnable_codebook else self.embed.detach()\n",
    "\n",
    "        if self.affine_param:\n",
    "            codebook_std = self.codebook_variance.clamp(min_=1e-5).sqrt()\n",
    "            batch_std = self.batch_variance.clamp(min_=1e-5).sqrt()\n",
    "            embed = (embed - self.codebook_mean) * (batch_std / codebook_std) + self.batch_mean\n",
    "\n",
    "        dist = -cdist(flatten, embed)\n",
    "\n",
    "        embed_ind, embed_onehot = self.gumbel_sample(\n",
    "            dist, dim=-1, temperature=sample_codebook_temp, training=Tensor.training\n",
    "        )\n",
    "\n",
    "        embed_ind = unpack_one(embed_ind, ps, \"h *\")\n",
    "\n",
    "        if Tensor.training:\n",
    "            unpacked_onehot = unpack_one(embed_onehot, ps, \"h * c\")\n",
    "            quantize = Tensor.einsum(\"h b n c, h c d -> h b n d\", unpacked_onehot, embed)\n",
    "        else:\n",
    "            quantize = batched_embedding(embed_ind, embed)\n",
    "\n",
    "        if Tensor.training and self.ema_update and not freeze_codebook:\n",
    "            if self.affine_param:\n",
    "                flatten = (flatten - self.batch_mean) * (codebook_std / batch_std) + self.codebook_mean\n",
    "\n",
    "            if mask is not None:\n",
    "                embed_onehot[mask.logical_not().int()] = 0.0\n",
    "\n",
    "            cluster_size = embed_onehot.sum(axis=1)\n",
    "\n",
    "            self.all_reduce_fn(cluster_size)\n",
    "            ema_inplace(self.cluster_size, cluster_size, self.decay)\n",
    "\n",
    "            embed_sum = Tensor.einsum(\"h n d, h n c -> h c d\", flatten, embed_onehot)\n",
    "            self.all_reduce_fn(embed_sum.contiguous())\n",
    "            ema_inplace(self.embed_avg, embed_sum, self.decay)\n",
    "\n",
    "            cluster_size = laplace_smoothing(\n",
    "                self.cluster_size, self.codebook_size, self.eps\n",
    "            ) * self.cluster_size.sum(axis=-1, keepdim=True)\n",
    "\n",
    "            embed_normalized = self.embed_avg / rearrange(cluster_size, \"... -> ... 1\")\n",
    "            self.embed.replace(embed_normalized)\n",
    "            self.expire_codes_(x)\n",
    "\n",
    "        if needs_codebook_dim:\n",
    "            quantize, embed_ind = tuple(rearrange(t, \"1 ... -> ...\") for t in (quantize, embed_ind))\n",
    "\n",
    "        dist = unpack_one(dist, ps, \"h * d\")\n",
    "\n",
    "        return quantize, embed_ind, dist"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a30471de-c429-4ad7-a258-e56f49a97dcf",
   "metadata": {},
   "source": [
    "# VectorQuantize: The Core of VQ-BeTs\n",
    "\n",
    "The `VectorQuantize` class is the cornerstone of Vector-Quantized Behavior Transformers (VQ-BeTs). This class implements the vector quantization process, which is crucial for transforming continuous action spaces into discrete representations that can be effectively processed by transformer models.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **Multi-Head Support**:\n",
    "   - Allows for multiple independent codebooks, enabling more nuanced representation of complex behaviors.\n",
    "   - Essential for capturing different aspects of behavior in parallel.\n",
    "\n",
    "2. **Flexible Codebook Initialization**:\n",
    "   - Supports both uniform and k-means initialization.\n",
    "   - K-means initialization can lead to better initial representations of the action space.\n",
    "\n",
    "3. **Commitment Loss**:\n",
    "   - Encourages the encoder to produce embeddings close to the codebook vectors.\n",
    "   - Crucial for learning meaningful discrete representations of continuous actions.\n",
    "\n",
    "4. **Straight-Through Estimator**:\n",
    "   - Allows gradient flow through the discretization process.\n",
    "   - Essential for end-to-end training of the VQ-BeT model.\n",
    "\n",
    "5. **Gumbel-Softmax Sampling**:\n",
    "   - Enables stochastic selection of codebook vectors.\n",
    "   - Facilitates exploration in the action space during training.\n",
    "\n",
    "6. **EMA Codebook Updates**:\n",
    "   - Provides stable updates to the codebook vectors.\n",
    "   - Important for consistent learning and adaptation to evolving behavior patterns.\n",
    "\n",
    "7. **Orthogonal Regularization**:\n",
    "   - Encourages diversity in the codebook vectors.\n",
    "   - Helps in capturing a wide range of distinct behaviors.\n",
    "\n",
    "8. **Affine Transformation Support**:\n",
    "   - Allows for normalization of inputs and codebook vectors.\n",
    "   - Aids in handling varying scales and distributions in the action space.\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Discrete Representation**: Enables the transformation of continuous actions into discrete tokens, allowing the use of transformer architectures for behavior modeling.\n",
    "- **Efficient Learning**: By clustering similar actions, it helps in identifying and learning common behavior patterns.\n",
    "- **Scalability**: Allows handling of high-dimensional and complex action spaces typical in robotics and other advanced control tasks.\n",
    "- **Exploration-Exploitation Balance**: The stochastic nature of the quantization process aids in balancing exploration and exploitation during learning.\n",
    "\n",
    "Let's examine the implementation of the VectorQuantize class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "542d1f55-7615-44b9-83e8-f0fb91abc5ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "class VectorQuantize:\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim,\n",
    "        codebook_size,\n",
    "        codebook_dim=None,\n",
    "        heads=1,\n",
    "        separate_codebook_per_head=False,\n",
    "        decay=0.8,\n",
    "        eps=1e-5,\n",
    "        kmeans_init=False,\n",
    "        kmeans_iters=10,\n",
    "        sync_kmeans=True,\n",
    "        threshold_ema_dead_code=0,\n",
    "        channel_last=True,\n",
    "        accept_image_fmap=False,\n",
    "        commitment_weight=1.0,\n",
    "        commitment_use_cross_entropy_loss=False,\n",
    "        orthogonal_reg_weight=0.0,\n",
    "        orthogonal_reg_active_codes_only=False,\n",
    "        orthogonal_reg_max_codes=None,\n",
    "        stochastic_sample_codes=False,\n",
    "        sample_codebook_temp=1.0,\n",
    "        straight_through=False,\n",
    "        reinmax=False,  # using reinmax for improved straight-through, assuming straight through helps at all\n",
    "        sync_codebook=None,\n",
    "        sync_affine_param=False,\n",
    "        ema_update=True,\n",
    "        learnable_codebook=False,\n",
    "        in_place_codebook_optimizer: Callable[\n",
    "            ..., Optimizer\n",
    "        ] = None,  # Optimizer used to update the codebook embedding if using learnable_codebook\n",
    "        affine_param=False,\n",
    "        affine_param_batch_decay=0.99,\n",
    "        affine_param_codebook_decay=0.9,\n",
    "        sync_update_v=0.0,  # the v that controls optimistic vs pessimistic update for synchronous update rule (21) https://minyoungg.github.io/vqtorch/assets/draft_050523.pdf\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.heads = heads\n",
    "        self.separate_codebook_per_head = separate_codebook_per_head\n",
    "\n",
    "        codebook_dim = codebook_dim if (codebook_dim is not None) else dim\n",
    "        codebook_input_dim = codebook_dim * heads\n",
    "\n",
    "        requires_projection = codebook_input_dim != dim\n",
    "        self.project_in = nn.Linear(dim, codebook_input_dim) if requires_projection else None\n",
    "        self.project_out = nn.Linear(codebook_input_dim, dim) if requires_projection else None\n",
    "\n",
    "        self.eps = eps\n",
    "        self.commitment_weight = commitment_weight\n",
    "        self.commitment_use_cross_entropy_loss = commitment_use_cross_entropy_loss  # whether to use cross entropy loss to codebook as commitment loss\n",
    "\n",
    "        self.learnable_codebook = learnable_codebook\n",
    "\n",
    "        has_codebook_orthogonal_loss = orthogonal_reg_weight > 0\n",
    "        self.has_codebook_orthogonal_loss = has_codebook_orthogonal_loss\n",
    "        self.orthogonal_reg_weight = orthogonal_reg_weight\n",
    "        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only\n",
    "        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes\n",
    "\n",
    "        assert not (ema_update and learnable_codebook), \"learnable codebook not compatible with EMA update\"\n",
    "\n",
    "        assert 0 <= sync_update_v <= 1.0\n",
    "        assert not (sync_update_v > 0.0 and not learnable_codebook), \"learnable codebook must be turned on\"\n",
    "\n",
    "        self.sync_update_v = sync_update_v\n",
    "\n",
    "        gumbel_sample_fn = partial(\n",
    "            gumbel_sample,\n",
    "            stochastic=stochastic_sample_codes,\n",
    "            reinmax=reinmax,\n",
    "            straight_through=straight_through,\n",
    "        )\n",
    "\n",
    "        if sync_codebook is None:\n",
    "            sync_codebook = False\n",
    "\n",
    "        codebook_kwargs = {\n",
    "            \"dim\": codebook_dim,\n",
    "            \"num_codebooks\": heads if separate_codebook_per_head else 1,\n",
    "            \"codebook_size\": codebook_size,\n",
    "            \"kmeans_init\": kmeans_init,\n",
    "            \"kmeans_iters\": kmeans_iters,\n",
    "            \"sync_kmeans\": sync_kmeans,\n",
    "            \"decay\": decay,\n",
    "            \"eps\": eps,\n",
    "            \"threshold_ema_dead_code\": threshold_ema_dead_code,\n",
    "            \"use_ddp\": sync_codebook,\n",
    "            \"learnable_codebook\": has_codebook_orthogonal_loss or learnable_codebook,\n",
    "            \"sample_codebook_temp\": sample_codebook_temp,\n",
    "            \"gumbel_sample\": gumbel_sample_fn,\n",
    "            \"ema_update\": ema_update,\n",
    "        }\n",
    "\n",
    "        if affine_param:\n",
    "            codebook_kwargs = dict(\n",
    "                **codebook_kwargs,\n",
    "                affine_param=True,\n",
    "                sync_affine_param=sync_affine_param,\n",
    "                affine_param_batch_decay=affine_param_batch_decay,\n",
    "                affine_param_codebook_decay=affine_param_codebook_decay,\n",
    "            )\n",
    "\n",
    "        self._codebook = EuclideanCodebook(**codebook_kwargs)\n",
    "\n",
    "        self.in_place_codebook_optimizer = (\n",
    "            in_place_codebook_optimizer(nn.state.get_parameters(self._codebook))\n",
    "            if (in_place_codebook_optimizer is not None)\n",
    "            else None\n",
    "        )\n",
    "\n",
    "        self.codebook_size = codebook_size\n",
    "\n",
    "        self.accept_image_fmap = accept_image_fmap\n",
    "        self.channel_last = channel_last\n",
    "\n",
    "    @property\n",
    "    def codebook(self):\n",
    "        codebook = self._codebook.embed\n",
    "\n",
    "        if self.separate_codebook_per_head:\n",
    "            return codebook\n",
    "\n",
    "        return rearrange(codebook, \"1 ... -> ...\")\n",
    "\n",
    "    @codebook.setter\n",
    "    def codebook(self, codes):\n",
    "        if not self.separate_codebook_per_head:\n",
    "            codes = rearrange(codes, \"... -> 1 ...\")\n",
    "\n",
    "        self._codebook.embed.copy_(codes)\n",
    "\n",
    "    def get_codebook_vector_from_indices(self, indices):\n",
    "        codebook = self.codebook\n",
    "        is_multiheaded = codebook.ndim > 2\n",
    "\n",
    "        if not is_multiheaded:\n",
    "            codes = codebook[indices]\n",
    "            return rearrange(codes, \"... h d -> ... (h d)\")\n",
    "\n",
    "        indices, ps = pack_one(indices, \"b * h\")\n",
    "        indices = rearrange(indices, \"b n h -> b h n\")\n",
    "\n",
    "        indices = repeat(indices, \"b h n -> b h n d\", d=codebook.shape[-1])\n",
    "        codebook = repeat(codebook, \"h n d -> b h n d\", b=indices.shape[0])\n",
    "\n",
    "        codes = codebook.gather(2, indices)\n",
    "        codes = rearrange(codes, \"b h n d -> b n (h d)\")\n",
    "        codes = unpack_one(codes, ps, \"b * d\")\n",
    "        return codes\n",
    "\n",
    "    def __call__(\n",
    "        self,\n",
    "        x: Tensor,\n",
    "        indices=None,\n",
    "        mask=None,\n",
    "        sample_codebook_temp=None,\n",
    "        freeze_codebook=False,\n",
    "    ):\n",
    "        orig_input = x\n",
    "\n",
    "        only_one = x.ndim == 2\n",
    "\n",
    "        if only_one:\n",
    "            assert mask is None\n",
    "            x = rearrange(x, \"b d -> b 1 d\")\n",
    "\n",
    "        shape, device, heads, is_multiheaded, _codebook_size, return_loss = (\n",
    "            x.shape,\n",
    "            x.device,\n",
    "            self.heads,\n",
    "            self.heads > 1,\n",
    "            self.codebook_size,\n",
    "            (indices is not None),\n",
    "        )\n",
    "\n",
    "        need_transpose = not self.channel_last and not self.accept_image_fmap\n",
    "        should_inplace_optimize = self.in_place_codebook_optimizer is not None\n",
    "\n",
    "        # rearrange inputs\n",
    "\n",
    "        if self.accept_image_fmap:\n",
    "            height, width = x.shape[-2:]\n",
    "            x = rearrange(x, \"b c h w -> b (h w) c\")\n",
    "\n",
    "        if need_transpose:\n",
    "            x = rearrange(x, \"b d n -> b n d\")\n",
    "\n",
    "        # project input\n",
    "\n",
    "        x = self.project_in(x) if self.project_in is not None else x\n",
    "\n",
    "        # handle multi-headed separate codebooks\n",
    "\n",
    "        if is_multiheaded:\n",
    "            ein_rhs_eq = \"h b n d\" if self.separate_codebook_per_head else \"1 (b h) n d\"\n",
    "            x = rearrange(x, f\"b n (h d) -> {ein_rhs_eq}\", h=heads)\n",
    "\n",
    "        # l2norm for cosine sim, otherwise identity\n",
    "\n",
    "        x = self._codebook.transform_input(x)\n",
    "\n",
    "        # codebook forward kwargs\n",
    "\n",
    "        codebook_forward_kwargs = {\n",
    "            \"sample_codebook_temp\": sample_codebook_temp,\n",
    "            \"mask\": mask,\n",
    "            \"freeze_codebook\": freeze_codebook,\n",
    "        }\n",
    "\n",
    "        # quantize\n",
    "\n",
    "        quantize, embed_ind, distances = self._codebook(x, **codebook_forward_kwargs)\n",
    "\n",
    "        # one step in-place update\n",
    "\n",
    "        if should_inplace_optimize and Tensor.training and not freeze_codebook:\n",
    "            if mask is not None:\n",
    "                loss = (quantize - x.detach()).square().sum()\n",
    "\n",
    "                loss_mask = mask\n",
    "                if is_multiheaded:\n",
    "                    loss_mask = repeat(\n",
    "                        mask,\n",
    "                        \"b n -> c (b h) n\",\n",
    "                        c=loss.shape[0],\n",
    "                        h=loss.shape[1] // mask.shape[0],\n",
    "                    )\n",
    "\n",
    "                loss = loss[loss_mask].mean()\n",
    "\n",
    "            else:\n",
    "                loss = (quantize - x.detach()).square().sum().mean()\n",
    "\n",
    "            loss.backward()\n",
    "            self.in_place_codebook_optimizer.step()\n",
    "            self.in_place_codebook_optimizer.zero_grad()\n",
    "\n",
    "            # quantize again\n",
    "\n",
    "            quantize, embed_ind, distances = self._codebook(x, **codebook_forward_kwargs)\n",
    "\n",
    "        if Tensor.training:\n",
    "            # determine code to use for commitment loss\n",
    "            commit_quantize = quantize.detach() if not self.learnable_codebook or freeze_codebook else quantize\n",
    "\n",
    "            # straight through\n",
    "\n",
    "            quantize = x + (quantize - x).detach()\n",
    "\n",
    "            if self.sync_update_v > 0.0:\n",
    "                # (21) in https://minyoungg.github.io/vqtorch/assets/draft_050523.pdf\n",
    "                quantize = quantize + self.sync_update_v * (quantize - quantize.detach())\n",
    "\n",
    "        # function for calculating cross entropy loss to distance matrix\n",
    "        # used for (1) naturalspeech2 training residual vq latents to be close to the correct codes and (2) cross-entropy based commitment loss\n",
    "\n",
    "        def calculate_ce_loss(codes):\n",
    "            if not is_multiheaded:\n",
    "                dist_einops_eq = \"1 b n l -> b l n\"\n",
    "            elif self.separate_codebook_per_head:\n",
    "                dist_einops_eq = \"c b n l -> b l n c\"\n",
    "            else:\n",
    "                dist_einops_eq = \"1 (b h) n l -> b l n h\"\n",
    "\n",
    "            ce_loss = rearrange(distances, dist_einops_eq, b=shape[0]).cross_entropy(\n",
    "                codes #, ignore_index=-1\n",
    "            )\n",
    "\n",
    "            return ce_loss\n",
    "\n",
    "        # if returning cross entropy loss on codes that were passed in\n",
    "\n",
    "        if return_loss:\n",
    "            return quantize, calculate_ce_loss(indices)\n",
    "\n",
    "        # transform embedding indices\n",
    "\n",
    "        if is_multiheaded:\n",
    "            if self.separate_codebook_per_head:\n",
    "                embed_ind = rearrange(embed_ind, \"h b n -> b n h\", h=heads)\n",
    "            else:\n",
    "                embed_ind = rearrange(embed_ind, \"1 (b h) n -> b n h\", h=heads)\n",
    "\n",
    "        if self.accept_image_fmap:\n",
    "            embed_ind = rearrange(embed_ind, \"b (h w) ... -> b h w ...\", h=height, w=width)\n",
    "\n",
    "        if only_one:\n",
    "            embed_ind = rearrange(embed_ind, \"b 1 -> b\")\n",
    "\n",
    "        # aggregate loss\n",
    "\n",
    "        loss = Tensor([0.0], requires_grad=Tensor.training)\n",
    "\n",
    "        if Tensor.training:\n",
    "            if self.commitment_weight > 0:\n",
    "                if self.commitment_use_cross_entropy_loss:\n",
    "                    if mask is not None:\n",
    "                        ce_loss_mask = mask\n",
    "                        if is_multiheaded:\n",
    "                            ce_loss_mask = repeat(ce_loss_mask, \"b n -> b n h\", h=heads)\n",
    "                        \n",
    "                        embed_ind = embed_ind.masked_fill(ce_loss_mask.logical_not(), -1)\n",
    "\n",
    "                    commit_loss = calculate_ce_loss(embed_ind)\n",
    "                else:\n",
    "                    if mask is not None:\n",
    "                        # with variable lengthed sequences\n",
    "                        commit_loss = (commit_quantize - x).square().sum()\n",
    "\n",
    "                        loss_mask = mask\n",
    "                        if is_multiheaded:\n",
    "                            loss_mask = repeat(\n",
    "                                loss_mask,\n",
    "                                \"b n -> c (b h) n\",\n",
    "                                c=commit_loss.shape[0],\n",
    "                                h=commit_loss.shape[1] // mask.shape[0],\n",
    "                            )\n",
    "\n",
    "                        commit_loss = commit_loss[loss_mask].mean()\n",
    "                    else:\n",
    "                        commit_loss = (commit_quantize - x).square().sum().mean()\n",
    "\n",
    "                loss = loss + commit_loss * self.commitment_weight\n",
    "\n",
    "            if self.has_codebook_orthogonal_loss:\n",
    "                codebook = self._codebook.embed\n",
    "\n",
    "                # only calculate orthogonal loss for the activated codes for this batch\n",
    "\n",
    "                if self.orthogonal_reg_active_codes_only:\n",
    "                    assert not (\n",
    "                        is_multiheaded and self.separate_codebook_per_head\n",
    "                    ), \"orthogonal regularization for only active codes not compatible with multi-headed with separate codebooks yet\"\n",
    "                    unique_code_ids = Tensor(embed_ind.numpy().unique(), requires_grad=False)\n",
    "                    codebook = codebook[:, unique_code_ids]\n",
    "\n",
    "                num_codes = codebook.shape[-2]\n",
    "\n",
    "                if (self.orthogonal_reg_max_codes is not None) and num_codes > self.orthogonal_reg_max_codes:\n",
    "                    rand_ids = list(np.randperm(num_codes)[: self.orthogonal_reg_max_codes])\n",
    "                    codebook = codebook[:, rand_ids]\n",
    "\n",
    "                orthogonal_reg_loss = orthogonal_loss_fn(codebook)\n",
    "                loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight\n",
    "\n",
    "        # handle multi-headed quantized embeddings\n",
    "\n",
    "        if is_multiheaded:\n",
    "            if self.separate_codebook_per_head:\n",
    "                quantize = rearrange(quantize, \"h b n d -> b n (h d)\", h=heads)\n",
    "            else:\n",
    "                quantize = rearrange(quantize, \"1 (b h) n d -> b n (h d)\", h=heads)\n",
    "\n",
    "        # project out\n",
    "\n",
    "        quantize = self.project_out(quantize) if self.project_out is not None else quantize\n",
    "\n",
    "        # rearrange quantized embeddings\n",
    "\n",
    "        if need_transpose:\n",
    "            quantize = rearrange(quantize, \"b n d -> b d n\")\n",
    "\n",
    "        if self.accept_image_fmap:\n",
    "            quantize = rearrange(quantize, \"b (h w) c -> b c h w\", h=height, w=width)\n",
    "\n",
    "        if only_one:\n",
    "            quantize = rearrange(quantize, \"b 1 d -> b d\")\n",
    "\n",
    "        # if masking, only return quantized for where mask has True\n",
    "\n",
    "        if mask is not None:\n",
    "            quantize = rearrange(mask, \"... -> ... 1\").where(quantize, orig_input)\n",
    "\n",
    "        return quantize, embed_ind, loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac38817-9c7c-4226-80a2-248ba5f1ef59",
   "metadata": {},
   "source": [
    "# ResidualVQ: Multi-Stage Vector Quantization for VQ-BeTs\n",
    "\n",
    "The `ResidualVQ` class implements a crucial component of Vector-Quantized Behavior Transformers (VQ-BeTs): multi-stage vector quantization. This approach, also known as residual vector quantization, enhances the model's ability to capture fine-grained details in the action space.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **Multiple Quantization Layers**:\n",
    "   - Uses a cascade of VectorQuantize layers to iteratively quantize the input and its residuals.\n",
    "   - Allows for more precise representation of complex behaviors by capturing details at different scales.\n",
    "\n",
    "2. **Shared or Independent Codebooks**:\n",
    "   - Option to use shared codebooks across layers for parameter efficiency or independent codebooks for more expressive power.\n",
    "   - Crucial for balancing model complexity and representation capacity.\n",
    "\n",
    "3. **Quantize Dropout**:\n",
    "   - Enables stochastic dropping of fine-grained quantizations during training.\n",
    "   - Enhances model robustness and generalization by forcing it to work with partial information.\n",
    "\n",
    "4. **Flexible Dimensionality**:\n",
    "   - Supports different input and codebook dimensions with optional projection layers.\n",
    "   - Allows adaptation to various action space configurations in different environments.\n",
    "\n",
    "5. **Image Feature Map Support**:\n",
    "   - Can handle inputs structured as image feature maps.\n",
    "   - Useful for tasks involving visual inputs, such as robotic manipulation or autonomous driving.\n",
    "\n",
    "6. **Codebook Freezing**:\n",
    "   - Allows freezing of codebooks during certain phases of training.\n",
    "   - Important for fine-tuning or transfer learning scenarios in behavior modeling.\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Hierarchical Representation**: Enables capturing behavior patterns at multiple levels of abstraction.\n",
    "- **Improved Reconstruction**: Residual quantization often leads to better reconstruction of the original input, crucial for accurate behavior modeling.\n",
    "- **Scalability**: Allows handling of complex, high-dimensional action spaces common in advanced robotics and control tasks.\n",
    "- **Adaptability**: The flexible architecture allows the model to adapt to various types of behavior data and task requirements.\n",
    "\n",
    "Let's examine the implementation of the ResidualVQ class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f61d3d0f-c28d-4d4f-a737-f3f661c575c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "class ResidualVQ:\n",
    "    \"\"\"\n",
    "    Residual VQ is composed of multiple VectorQuantize layers.\n",
    "\n",
    "    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf\n",
    "        \"Residual Vector Quantizer (a.k.a. multi-stage vector quantizer [36]) cascades Nq layers of VQ as follows. The unquantized input vector is\n",
    "        passed through a first VQ and quantization residuals are computed. The residuals are then iteratively quantized by a sequence of additional\n",
    "        Nq -1 vector quantizers, as described in Algorithm 1.\"\n",
    "\n",
    "\n",
    "    self.project_in: function for projecting input to codebook dimension\n",
    "    self.project_out: function for projecting codebook dimension to output dimension\n",
    "    self.layers: nn.ModuleList of VectorQuantize layers that contains Nq layers of VQ as described in the paper.\n",
    "    self.freeze_codebook: buffer to save an indicator whether the codebook is frozen or not. VQ-BeT will check this to determine whether to update the codebook or not.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        dim,\n",
    "        num_quantizers,\n",
    "        codebook_dim=None,\n",
    "        shared_codebook=False,\n",
    "        heads=1,\n",
    "        quantize_dropout=False,\n",
    "        quantize_dropout_cutoff_index=0,\n",
    "        quantize_dropout_multiple_of=1,\n",
    "        accept_image_fmap=False,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        assert heads == 1, \"residual vq is not compatible with multi-headed codes\"\n",
    "        codebook_dim = codebook_dim if (codebook_dim is not None) else dim\n",
    "        codebook_input_dim = codebook_dim * heads\n",
    "        \n",
    "        requires_projection = codebook_input_dim != dim\n",
    "        self.project_in = nn.Linear(dim, codebook_input_dim) if requires_projection else lambda x: x\n",
    "        self.project_out = nn.Linear(codebook_input_dim, dim) if requires_projection else lambda x: x\n",
    "\n",
    "        self.num_quantizers = num_quantizers\n",
    "\n",
    "        self.accept_image_fmap = accept_image_fmap\n",
    "        self.layers = [\n",
    "                VectorQuantize(\n",
    "                    dim=codebook_dim, codebook_dim=codebook_dim, accept_image_fmap=accept_image_fmap, **kwargs\n",
    "                )\n",
    "                for _ in range(num_quantizers)\n",
    "            ]\n",
    "\n",
    "        self.quantize_dropout = quantize_dropout and num_quantizers > 1\n",
    "\n",
    "        assert quantize_dropout_cutoff_index >= 0\n",
    "\n",
    "        # freeze_cookbook is like a buffer. Don't strict_load or strict_save\n",
    "        self.freeze_codebook = False\n",
    "        self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index\n",
    "        self.quantize_dropout_multiple_of = quantize_dropout_multiple_of  # encodec paper proposes structured dropout, believe this was set to 4\n",
    "\n",
    "        if not shared_codebook:\n",
    "            return\n",
    "\n",
    "        first_vq, *rest_vq = self.layers\n",
    "        codebook = first_vq._codebook\n",
    "\n",
    "        for vq in rest_vq:\n",
    "            vq._codebook = codebook\n",
    "\n",
    "    @property\n",
    "    def codebooks(self):\n",
    "        codebooks = [layer._codebook.embed for layer in self.layers]\n",
    "        codebooks = Tensor.stack(*codebooks, dim=0)\n",
    "        codebooks = rearrange(codebooks, \"q 1 c d -> q c d\")\n",
    "        return codebooks\n",
    "\n",
    "    def get_codebook_vector_from_indices(self, indices):\n",
    "        # this function will return the codes from all codebooks across layers corresponding to the indices\n",
    "        batch, quantize_dim = indices.shape[0], indices.shape[-1]\n",
    "\n",
    "        # may also receive indices in the shape of 'b h w q' (accept_image_fmap)\n",
    "\n",
    "        indices, ps = pack([indices], \"b * q\")\n",
    "\n",
    "        # because of quantize dropout, one can pass in indices that are coarse\n",
    "        # and the network should be able to reconstruct\n",
    "\n",
    "        if quantize_dim < self.num_quantizers:\n",
    "            assert (\n",
    "                self.quantize_dropout > 0.0\n",
    "            ), \"quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations\"\n",
    "            indices = indices.pad((0, self.num_quantizers - quantize_dim), value=-1)\n",
    "\n",
    "        # get ready for gathering\n",
    "\n",
    "        codebooks = repeat(self.codebooks, \"q c d -> q b c d\", b=batch)\n",
    "        gather_indices = repeat(indices, \"b n q -> q b n d\", d=codebooks.shape[-1])\n",
    "\n",
    "        # take care of quantizer dropout\n",
    "\n",
    "        # have it fetch a dummy code to be masked out later\n",
    "        mask = gather_indices == -1.0\n",
    "        gather_indices = gather_indices.masked_fill(\n",
    "            mask, 0\n",
    "        )  # have it fetch a dummy code to be masked out later\n",
    "\n",
    "        all_codes = codebooks.gather(2, gather_indices)  # gather all codes\n",
    "\n",
    "        # mask out any codes that were dropout-ed\n",
    "        all_codes = all_codes.masked_fill(mask, 0.0)\n",
    "\n",
    "        # if (accept_image_fmap = True) then return shape (quantize, batch, height, width, dimension)\n",
    "\n",
    "        (all_codes,) = unpack(all_codes, ps, \"q b * d\")\n",
    "\n",
    "        return all_codes\n",
    "\n",
    "    def __call__(self, x:Tensor, indices=None, return_all_codes=False, sample_codebook_temp=None):\n",
    "        \"\"\"\n",
    "        For given input tensor x, this function will return the quantized output, the indices of the quantized output, and the loss.\n",
    "        First, the input tensor x is projected to the codebook dimension. Then, the input tensor x is passed through Nq layers of VectorQuantize.\n",
    "        The residual value of each layer is fed to the next layer.\n",
    "        \"\"\"\n",
    "        num_quant, quant_dropout_multiple_of, return_loss = (\n",
    "            self.num_quantizers,\n",
    "            self.quantize_dropout_multiple_of,\n",
    "            (indices is not None)\n",
    "        )\n",
    "\n",
    "        x = self.project_in(x)\n",
    "\n",
    "        assert not (self.accept_image_fmap and (indices is not None))\n",
    "\n",
    "        quantized_out = 0.0\n",
    "        residual = x\n",
    "\n",
    "        all_losses = []\n",
    "        all_indices = []\n",
    "\n",
    "        if return_loss:\n",
    "            assert not (indices == -1).any().item(), \"some of the residual vq indices were dropped out. please use indices derived when the module is in eval mode to derive cross entropy loss\"\n",
    "            ce_losses = []\n",
    "\n",
    "        should_quantize_dropout = Tensor.training and self.quantize_dropout and not return_loss\n",
    "\n",
    "        # sample a layer index at which to dropout further residual quantization\n",
    "        # also prepare null indices and loss\n",
    "\n",
    "        if should_quantize_dropout:\n",
    "            rand_quantize_dropout_index = randrange(self.quantize_dropout_cutoff_index, num_quant)\n",
    "\n",
    "            if quant_dropout_multiple_of != 1:\n",
    "                rand_quantize_dropout_index = (\n",
    "                    ceil((rand_quantize_dropout_index + 1) / quant_dropout_multiple_of)\n",
    "                    * quant_dropout_multiple_of\n",
    "                    - 1\n",
    "                )\n",
    "\n",
    "            null_indices_shape = (x.shape[0], *x.shape[-2:]) if self.accept_image_fmap else tuple(x.shape[:2])\n",
    "            null_indices = Tensor.full(null_indices_shape, -1.0, dtype=dtypes.long, requires_grad=False)\n",
    "            null_loss = Tensor.full((1,), 0.0, dtype=x.dtype, requires_grad=False)\n",
    "\n",
    "        # go through the layers\n",
    "\n",
    "        for quantizer_index, layer in enumerate(self.layers):\n",
    "            if should_quantize_dropout and quantizer_index > rand_quantize_dropout_index:\n",
    "                all_indices.append(null_indices)\n",
    "                all_losses.append(null_loss)\n",
    "                continue\n",
    "\n",
    "            layer_indices = None\n",
    "            if return_loss:\n",
    "                layer_indices = indices[..., quantizer_index]\n",
    "\n",
    "            quantized, *rest = layer(\n",
    "                residual,\n",
    "                indices=layer_indices,\n",
    "                sample_codebook_temp=sample_codebook_temp,\n",
    "                freeze_codebook=self.freeze_codebook,\n",
    "            )\n",
    "\n",
    "            residual = residual - quantized.detach()\n",
    "            quantized_out = quantized_out + quantized\n",
    "\n",
    "            if return_loss:\n",
    "                ce_loss = rest[0]\n",
    "                ce_losses.append(ce_loss)\n",
    "                continue\n",
    "\n",
    "            embed_indices, loss = rest\n",
    "\n",
    "            all_indices.append(embed_indices)\n",
    "            all_losses.append(loss)\n",
    "\n",
    "        # project out, if needed\n",
    "\n",
    "        quantized_out = self.project_out(quantized_out)\n",
    "\n",
    "        # whether to early return the cross entropy loss\n",
    "\n",
    "        if return_loss:\n",
    "            return quantized_out, sum(ce_losses)\n",
    "\n",
    "        # stack all losses and indices\n",
    "\n",
    "        # be cautious here: might need to decouple the all_losses through a point into individual objects\n",
    "        all_losses = Tensor.stack(*all_losses, dim=-1)\n",
    "        all_indices = Tensor.stack(*all_indices, dim=-1)\n",
    "\n",
    "        ret = (quantized_out, all_indices, all_losses)\n",
    "\n",
    "        if return_all_codes:\n",
    "            # whether to return all codes from all codebooks across layers\n",
    "            all_codes = self.get_codebook_vector_from_indices(all_indices)\n",
    "\n",
    "            # will return all codes in shape (quantizer, batch, sequence length, codebook dimension)\n",
    "            ret = (*ret, all_codes)\n",
    "\n",
    "        return ret\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffaba650-ca02-4710-82d6-4bbec8c1d057",
   "metadata": {},
   "source": [
    "# Creating a Config\n",
    "\n",
    "By now, we've gotten all the basic building blocks assembled for making the critical parts of VQ-BeT. Let's make a config that helps us understand and tweak parameters in our system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "833567e0-20ad-49c8-bc9b-ff4377e9a64c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass, field\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class VQBeTConfig:\n",
    "    \"\"\"Configuration class for VQ-BeT.\n",
    "\n",
    "    Defaults are configured for training with PushT providing proprioceptive and single camera observations.\n",
    "\n",
    "    The parameters you will most likely need to change are the ones which depend on the environment / sensors.\n",
    "    Those are: `input_shapes` and `output_shapes`.\n",
    "\n",
    "    Notes on the inputs and outputs:\n",
    "        - \"observation.state\" is required as an input key.\n",
    "        - At least one key starting with \"observation.image is required as an input.\n",
    "        - If there are multiple keys beginning with \"observation.image\" they are treated as multiple camera\n",
    "          views. Right now we only support all images having the same shape.\n",
    "        - \"action\" is required as an output key.\n",
    "\n",
    "    Args:\n",
    "        n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the\n",
    "            current step and additional steps going back).\n",
    "        n_action_pred_token: Total number of current token and future tokens that VQ-BeT predicts.\n",
    "        action_chunk_size: Action chunk size of each action prediction token.\n",
    "        input_shapes: A dictionary defining the shapes of the input data for the policy.\n",
    "            The key represents the input data name, and the value is a list indicating the dimensions\n",
    "            of the corresponding data. For example, \"observation.image\" refers to an input from\n",
    "            a camera with dimensions [3, 96, 96], indicating it has three color channels and 96x96 resolution.\n",
    "            Importantly, shapes doesnt include batch dimension or temporal dimension.\n",
    "        output_shapes: A dictionary defining the shapes of the output data for the policy.\n",
    "            The key represents the output data name, and the value is a list indicating the dimensions\n",
    "            of the corresponding data. For example, \"action\" refers to an output shape of [14], indicating\n",
    "            14-dimensional actions. Importantly, shapes doesnt include batch dimension or temporal dimension.\n",
    "        input_normalization_modes: A dictionary with key representing the modality (e.g. \"observation.state\"),\n",
    "            and the value specifies the normalization mode to apply. The two available modes are \"mean_std\"\n",
    "            which subtracts the mean and divides by the standard deviation and \"min_max\" which rescale in a\n",
    "            [-1, 1] range.\n",
    "        output_normalization_modes: Similar dictionary as `normalize_input_modes`, but to unnormalize to the\n",
    "            original scale. Note that this is also used for normalizing the training targets.\n",
    "        vision_backbone: Name of the torchvision resnet backbone to use for encoding images.\n",
    "        crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit\n",
    "            within the image size. If None, no cropping is done.\n",
    "        crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval\n",
    "            mode).\n",
    "        pretrained_backbone_weights: Pretrained weights from torchvision to initalize the backbone.\n",
    "            `None` means no pretrained weights.\n",
    "        use_group_norm: Whether to replace batch normalization with group normalization in the backbone.\n",
    "            The group sizes are set to be about 16 (to be precise, feature_dim // 16).\n",
    "        spatial_softmax_num_keypoints: Number of keypoints for SpatialSoftmax.\n",
    "        n_vqvae_training_steps: Number of optimization steps for training Residual VQ.\n",
    "        vqvae_n_embed: Number of embedding vectors in the RVQ dictionary (each layer).\n",
    "        vqvae_embedding_dim: Dimension of each embedding vector in the RVQ dictionary.\n",
    "        vqvae_enc_hidden_dim: Size of hidden dimensions of Encoder / Decoder part of Residaul VQ-VAE\n",
    "        gpt_block_size: Max block size of minGPT (should be larger than the number of input tokens)\n",
    "        gpt_input_dim: Size of output input of GPT. This is also used as the dimension of observation features.\n",
    "        gpt_output_dim: Size of output dimension of GPT. This is also used as a input dimension of offset / bin prediction headers.\n",
    "        gpt_n_layer: Number of layers of GPT\n",
    "        gpt_n_head: Number of headers of GPT\n",
    "        gpt_hidden_dim: Size of hidden dimensions of GPT\n",
    "        dropout: Dropout rate for GPT\n",
    "        mlp_hidden_dim: Size of hidden dimensions of offset header / bin prediction headers parts of VQ-BeT\n",
    "        offset_loss_weight:  A constant that is multiplied to the offset loss\n",
    "        primary_code_loss_weight: A constant that is multiplied to the primary code prediction loss\n",
    "        secondary_code_loss_weight: A constant that is multiplied to the secondary code prediction loss\n",
    "        bet_softmax_temperature: Sampling temperature of code for rollout with VQ-BeT\n",
    "        sequentially_select: Whether select code of primary / secondary as sequentially (pick primary code,\n",
    "            and then select secodnary code), or at the same time.\n",
    "    \"\"\"\n",
    "\n",
    "    # Inputs / output structure.\n",
    "    n_obs_steps: int = 5\n",
    "    n_action_pred_token: int = 7\n",
    "    action_chunk_size: int = 5\n",
    "\n",
    "    input_shapes: dict[str, list[int]] = field(\n",
    "        default_factory=lambda: {\n",
    "            \"observation.image\": [3, 96, 96],\n",
    "            \"observation.state\": [2],\n",
    "        }\n",
    "    )\n",
    "    output_shapes: dict[str, list[int]] = field(\n",
    "        default_factory=lambda: {\n",
    "            \"action\": [2],\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # Normalization / Unnormalization\n",
    "    input_normalization_modes: dict[str, str] = field(\n",
    "        default_factory=lambda: {\n",
    "            \"observation.image\": \"mean_std\",\n",
    "            \"observation.state\": \"min_max\",\n",
    "        }\n",
    "    )\n",
    "    output_normalization_modes: dict[str, str] = field(default_factory=lambda: {\"action\": \"min_max\"})\n",
    "\n",
    "    # Architecture / modeling.\n",
    "    # Vision backbone.\n",
    "    vision_backbone: str = \"resnet18\"\n",
    "    crop_shape: tuple[int, int] | None = (84, 84)\n",
    "    crop_is_random: bool = True\n",
    "    pretrained_backbone_weights: str | None = None\n",
    "    use_group_norm: bool = True\n",
    "    spatial_softmax_num_keypoints: int = 32\n",
    "    # VQ-VAE\n",
    "    # change this to make it super quick to iterate on the main GPT training portion\n",
    "    n_vqvae_training_steps: int = 0\n",
    "    vqvae_n_embed: int = 16\n",
    "    vqvae_embedding_dim: int = 256\n",
    "    vqvae_enc_hidden_dim: int = 128\n",
    "    # VQ-BeT\n",
    "    gpt_block_size: int = 500\n",
    "    gpt_input_dim: int = 512\n",
    "    gpt_output_dim: int = 512\n",
    "    gpt_n_layer: int = 8\n",
    "    gpt_n_head: int = 8\n",
    "    gpt_hidden_dim: int = 512\n",
    "    dropout: float = 0.1\n",
    "    mlp_hidden_dim: int = 1024\n",
    "    offset_loss_weight: float = 10000.0\n",
    "    primary_code_loss_weight: float = 5.0\n",
    "    secondary_code_loss_weight: float = 0.5\n",
    "    bet_softmax_temperature: float = 0.1\n",
    "    sequentially_select: bool = False\n",
    "\n",
    "    def __post_init__(self):\n",
    "        \"\"\"Input validation (not exhaustive).\"\"\"\n",
    "        if not self.vision_backbone.startswith(\"resnet\"):\n",
    "            raise ValueError(\n",
    "                f\"`vision_backbone` must be one of the ResNet variants. Got {self.vision_backbone}.\"\n",
    "            )\n",
    "        image_keys = {k for k in self.input_shapes if k.startswith(\"observation.image\")}\n",
    "        if self.crop_shape is not None:\n",
    "            for image_key in image_keys:\n",
    "                if (\n",
    "                    self.crop_shape[0] > self.input_shapes[image_key][1]\n",
    "                    or self.crop_shape[1] > self.input_shapes[image_key][2]\n",
    "                ):\n",
    "                    raise ValueError(\n",
    "                        f\"`crop_shape` should fit within `input_shapes[{image_key}]`. Got {self.crop_shape} \"\n",
    "                        f\"for `crop_shape` and {self.input_shapes[image_key]} for \"\n",
    "                        \"`input_shapes[{image_key}]`.\"\n",
    "                    )\n",
    "        # Check that all input images have the same shape.\n",
    "        first_image_key = next(iter(image_keys))\n",
    "        for image_key in image_keys:\n",
    "            if self.input_shapes[image_key] != self.input_shapes[first_image_key]:\n",
    "                raise ValueError(\n",
    "                    f\"`input_shapes[{image_key}]` does not match `input_shapes[{first_image_key}]`, but we \"\n",
    "                    \"expect all image shapes to match.\"\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "921f1bec-4c61-447b-8939-a351fbfc11b2",
   "metadata": {},
   "source": [
    "# VqVae: The Core of Vector-Quantized Behavior Transformers (VQ-BeTs)\n",
    "\n",
    "The `VqVae` class implements the Vector Quantized Variational Autoencoder (VQ-VAE) component of VQ-BeTs. This is a crucial part of the architecture that enables the discretization of continuous action spaces, which is fundamental to the VQ-BeT approach.\n",
    "\n",
    "Key Components and Their Importance:\n",
    "\n",
    "1. **Encoder and Decoder**:\n",
    "   - Implemented as MLPs (Multi-Layer Perceptrons)\n",
    "   - Encoder: Transforms continuous actions into latent representations\n",
    "   - Decoder: Reconstructs actions from quantized latent representations\n",
    "   - Critical for learning meaningful representations of behaviors\n",
    "\n",
    "2. **Residual VQ Layer**:\n",
    "   - Uses multiple layers of vector quantization\n",
    "   - Allows for more expressive and fine-grained discretization of the action space\n",
    "   - Key to capturing complex behavior patterns\n",
    "\n",
    "3. **Two-Phase Training**:\n",
    "   - Phase 1: Trains the encoder, decoder, and VQ layer\n",
    "   - Phase 2: Uses the trained VQ-VAE to provide discrete codes for training the transformer\n",
    "\n",
    "4. **Code Generation**:\n",
    "   - Ability to generate discrete codes from continuous actions\n",
    "   - Essential for creating the \"vocabulary\" of behaviors that the transformer will learn to predict\n",
    "\n",
    "5. **Reconstruction and Commitment Losses**:\n",
    "   - Ensures that the discretized representations can accurately reconstruct the original actions\n",
    "   - Balances the trade-off between reconstruction accuracy and discretization\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Discretization**: Enables the transformation of continuous actions into discrete tokens, allowing the use of transformer architectures for behavior modeling\n",
    "- **Hierarchical Representation**: The residual VQ approach allows for capturing behavior at multiple levels of abstraction\n",
    "- **Efficient Learning**: By clustering similar actions into discrete codes, it helps in identifying and learning common behavior patterns\n",
    "- **Bridging Continuous and Discrete**: Allows the model to work with discrete tokens while still being able to generate continuous actions\n",
    "\n",
    "Let's examine the implementation of the VqVae class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b371647a-f736-48e8-80fe-cf8bcfa08a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VqVae:\n",
    "    def __init__(\n",
    "        self,\n",
    "        config: VQBeTConfig,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        VQ-VAE is composed of three parts: encoder, vq_layer, and decoder.\n",
    "        Encoder and decoder are MLPs consisting of an input, output layer, and hidden layer, respectively.\n",
    "        The vq_layer uses residual VQs.\n",
    "\n",
    "        This class contains functions for training the encoder and decoder along with the residual VQ layer (for trainign phase 1),\n",
    "        as well as functions to help BeT training part in training phase 2.\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        # 'discretized' indicates whether the Residual VQ part is trained or not. (After finishing the training, we set discretized=True)\n",
    "        self.discretized = False\n",
    "        self.optimized_steps = 0\n",
    "        # we use the fixed number of layers for Residual VQ across all environments.\n",
    "        self.vqvae_num_layers = 2\n",
    "\n",
    "        self.vq_layer = ResidualVQ(\n",
    "            dim=config.vqvae_embedding_dim,\n",
    "            num_quantizers=self.vqvae_num_layers,\n",
    "            codebook_size=config.vqvae_n_embed,\n",
    "        )\n",
    "\n",
    "        self.encoder = MLP(\n",
    "            in_channels=self.config.output_shapes[\"action\"][0] * self.config.action_chunk_size,\n",
    "            hidden_channels=[\n",
    "                config.vqvae_enc_hidden_dim,\n",
    "                config.vqvae_enc_hidden_dim,\n",
    "                config.vqvae_embedding_dim,\n",
    "            ],\n",
    "        )\n",
    "        self.decoder = MLP(\n",
    "            in_channels=config.vqvae_embedding_dim,\n",
    "            hidden_channels=[\n",
    "                config.vqvae_enc_hidden_dim,\n",
    "                config.vqvae_enc_hidden_dim,\n",
    "                self.config.output_shapes[\"action\"][0] * self.config.action_chunk_size,\n",
    "            ],\n",
    "        )\n",
    "\n",
    "    def get_embeddings_from_code(self, encoding_indices):\n",
    "            # This function gets code indices as inputs, and outputs embedding vectors corresponding to the code indices.\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        z_embed = self.vq_layer.get_codebook_vector_from_indices(encoding_indices)\n",
    "        # since the RVQ has multiple layers, it adds the vectors in the axis of layers to provide a vector for that code combination.\n",
    "        z_embed = z_embed.sum(axis=0)\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        return z_embed\n",
    "\n",
    "    def get_action_from_latent(self, latent):\n",
    "        # given latent vector, this function outputs the decoded action.\n",
    "        output = self.decoder(latent)\n",
    "        if self.config.action_chunk_size == 1:\n",
    "            return einops.rearrange(output, \"N (T A) -> N T A\", A=self.config.output_shapes[\"action\"][0])\n",
    "        else:\n",
    "            return einops.rearrange(output, \"N (T A) -> N T A\", A=self.config.output_shapes[\"action\"][0])\n",
    "\n",
    "    def get_code(self, state):\n",
    "        # in phase 2 of VQ-BeT training, we need a `ground truth labels of action data` to calculate the Focal loss for code prediction head. (please refer to section 3.3 in the paper https://arxiv.org/pdf/2403.03181)\n",
    "        # this function outputs the `GT code` of given action using frozen encoder and quantization layers. (please refer to Figure 2. in the paper https://arxiv.org/pdf/2403.03181)\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        state = einops.rearrange(state, \"N T A -> N (T A)\")\n",
    "        state_rep = self.encoder(state)\n",
    "        state_rep_shape = state_rep.shape[:-1]\n",
    "        state_rep_flat = state_rep.view(state_rep.size(0), -1, state_rep.size(1))\n",
    "        state_rep_flat, vq_code, vq_loss_state = self.vq_layer(state_rep_flat)\n",
    "        state_vq = state_rep_flat.view(*state_rep_shape, -1)\n",
    "        vq_code = vq_code.view(*state_rep_shape, -1)\n",
    "        vq_loss_state = vq_loss_state.sum()\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        return state_vq, vq_code\n",
    "\n",
    "    \n",
    "    def vqvae_forward(self, state):\n",
    "        # This function passes the given data through Residual VQ with Encoder and Decoder. Please refer to section 3.2 in the paper https://arxiv.org/pdf/2403.03181).\n",
    "        state = einops.rearrange(state, \"N T A -> N (T A)\")\n",
    "        # We start with passing action (or action chunk) at:t+n through the encoder ϕ.\n",
    "        state_rep = self.encoder(state)\n",
    "        state_rep_shape = state_rep.shape[:-1]\n",
    "        state_rep_flat = state_rep.view(state_rep.size(0), -1, state_rep.size(1))\n",
    "        # The resulting latent embedding vector x = ϕ(at:t+n) is then mapped to an embedding vector in the codebook of the RVQ layers by the nearest neighbor look-up.\n",
    "        state_rep_flat, vq_code, vq_loss_state = self.vq_layer(state_rep_flat)\n",
    "        #vq_loss_state.requires_grad = False\n",
    "        state_vq = state_rep_flat.view(*state_rep_shape, -1)\n",
    "        vq_code = vq_code.view(*state_rep_shape, -1)\n",
    "        # since the RVQ has multiple layers, it adds the vectors in the axis of layers to provide a vector for that code combination.\n",
    "        vq_loss_state = vq_loss_state.sum()\n",
    "        # Then, the discretized vector zq(x) is reconstructed as ψ(zq(x)) by passing through the decoder ψ.\n",
    "        dec_out = self.decoder(state_vq)\n",
    "        # Calculate L1 reconstruction loss\n",
    "        encoder_loss = (state - dec_out).abs().mean()\n",
    "        # add encoder reconstruction loss and commitment loss\n",
    "        rep_loss = encoder_loss + vq_loss_state * 5\n",
    "        \n",
    "        metric = (\n",
    "            encoder_loss.detach(),\n",
    "            vq_loss_state.detach(),\n",
    "            vq_code,\n",
    "            rep_loss.item(),\n",
    "        )\n",
    "        return rep_loss, metric"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8c91e5b-525b-4373-a5a7-2a0bdc7b4cb8",
   "metadata": {},
   "source": [
    "# GPT: The Transformer Component of VQ-BeTs\n",
    "\n",
    "The `GPT` class implements the Transformer component of Vector-Quantized Behavior Transformers (VQ-BeTs). This is a crucial part of the architecture that learns to predict future behavior tokens based on past observations and actions.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **Transformer Architecture**:\n",
    "   - Based on the GPT (Generative Pre-trained Transformer) model\n",
    "   - Crucial for learning long-range dependencies in behavior sequences\n",
    "\n",
    "2. **Flexible Configuration**:\n",
    "   - Uses a `VQBeTConfig` object for hyperparameters\n",
    "   - Allows easy experimentation with different model sizes and architectures\n",
    "\n",
    "3. **Token and Position Embeddings**:\n",
    "   - Separate embeddings for tokens (actions/observations) and positions\n",
    "   - Enables the model to understand both the content and sequence of behaviors\n",
    "\n",
    "4. **Multi-Layer Structure**:\n",
    "   - Uses multiple transformer blocks for deep representation learning\n",
    "   - Essential for capturing complex patterns in behavior data\n",
    "\n",
    "5. **Output Projection**:\n",
    "   - Projects the final hidden states to the output dimension\n",
    "   - Crucial for predicting the next action token\n",
    "\n",
    "6. **Parameter Configuration**:\n",
    "   - Separates parameters into those with and without weight decay\n",
    "   - Important for effective regularization during training\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Sequence Modeling**: Enables the model to learn and predict sequences of behavior tokens\n",
    "- **Context Understanding**: Allows the model to consider long-term context when predicting future actions\n",
    "- **Scalability**: The transformer architecture can handle variable-length input sequences\n",
    "- **Transfer Learning**: The GPT-style architecture allows for potential pre-training on large behavior datasets\n",
    "\n",
    "Let's examine the implementation of the GPT class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bef9d835-2faf-4591-985b-a46ee23ee4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPT section, here we gooooo\n",
    "\n",
    "from tinygrad.nn.state import get_state_dict\n",
    "\n",
    "class CausalSelfAttention:\n",
    "  def __init__(self, config):\n",
    "    assert config.gpt_hidden_dim % config.gpt_n_head == 0\n",
    "    # key, query, value projections for all heads, but in a batch\n",
    "    self.c_attn = nn.Linear(config.gpt_hidden_dim, 3 * config.gpt_hidden_dim)\n",
    "    # output projection\n",
    "    self.c_proj = nn.Linear(config.gpt_hidden_dim, config.gpt_hidden_dim)\n",
    "    # regularization\n",
    "    self.gpt_n_head = config.gpt_n_head\n",
    "    self.gpt_hidden_dim = config.gpt_hidden_dim\n",
    "    self.dropout_rate = config.dropout\n",
    "    # not really a 'bias', more of a mask, but following the OpenAI/HF naming though\n",
    "    self.bias = Tensor.ones(config.gpt_block_size, config.gpt_block_size, requires_grad=False).view(\n",
    "                1, 1, config.gpt_block_size, config.gpt_block_size\n",
    "            ).tril().contiguous().detach().realize()\n",
    "  def __call__(self, x: Tensor):\n",
    "    (\n",
    "        B,\n",
    "        T,\n",
    "        C,\n",
    "    ) = x.size()  # batch size, sequence length, embedding dimensionality (gpt_hidden_dim)\n",
    "    # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
    "    q, k, v = self.c_attn(x).split(self.gpt_hidden_dim, dim=2)\n",
    "    k = k.view(B, T, self.gpt_n_head, C // self.gpt_n_head).transpose(1, 2)  # (B, nh, T, hs)\n",
    "    q = q.view(B, T, self.gpt_n_head, C // self.gpt_n_head).transpose(1, 2)  # (B, nh, T, hs)\n",
    "    v = v.view(B, T, self.gpt_n_head, C // self.gpt_n_head).transpose(1, 2)  # (B, nh, T, hs)\n",
    "    # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
    "    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
    "    att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float(\"-inf\"))\n",
    "    att = att.softmax(axis=-1)\n",
    "    att = att.dropout(self.dropout_rate)\n",
    "    y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
    "    y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side\n",
    "    # output projection\n",
    "    y = self.c_proj(y).dropout(self.dropout_rate)\n",
    "    return y\n",
    "\n",
    "class Block:\n",
    "  # causual self-attention block for GPT\n",
    "  def __init__(self, config):\n",
    "    self.ln_1 = nn.LayerNorm(config.gpt_hidden_dim)\n",
    "    self.attn = CausalSelfAttention(config)\n",
    "    self.ln_2 = nn.LayerNorm(config.gpt_hidden_dim)\n",
    "    self.mlp = [\n",
    "        nn.Linear(config.gpt_hidden_dim, 4 * config.gpt_hidden_dim),\n",
    "        Tensor.gelu,\n",
    "        nn.Linear(4 * config.gpt_hidden_dim, config.gpt_hidden_dim),\n",
    "    ]\n",
    "    self.dropout_rate = config.dropout\n",
    "\n",
    "  def __call__(self, x: Tensor) -> Tensor:\n",
    "    x = x + self.attn(self.ln_1(x.realize()))\n",
    "    x = x.realize() + self.ln_2(x.realize()).sequential(self.mlp).realize().dropout(self.dropout_rate)\n",
    "    return x\n",
    "\n",
    "class GPT:\n",
    "    \"\"\"\n",
    "    Original comments:\n",
    "    Full definition of a GPT Language Model, all of it in this single file.\n",
    "    References: the official GPT-2 TensorFlow implementation released by OpenAI:\n",
    "    https://github.com/openai/gpt-2/blob/master/src/model.py\n",
    "    \"\"\"\n",
    "    def __init__(self, config: VQBeTConfig):\n",
    "        \"\"\"\n",
    "        GPT model gets hyperparameters from a config object. Please refer configuration_vqbet.py for more details.\n",
    "        \"\"\"\n",
    "        assert config.gpt_output_dim is not None\n",
    "        assert config.gpt_block_size is not None\n",
    "        self.config = config\n",
    "        self.wte = nn.Linear(config.gpt_input_dim, config.gpt_hidden_dim)\n",
    "        self.wpe = nn.Embedding(config.gpt_block_size, config.gpt_hidden_dim)\n",
    "        self.drop = config.dropout\n",
    "        self.h = [Block(config) for _ in range(config.gpt_n_layer)]\n",
    "        self.ln_f = nn.LayerNorm(config.gpt_hidden_dim)\n",
    "        self.lm_head = nn.Linear(config.gpt_hidden_dim, config.gpt_output_dim, bias=False)\n",
    "        for pn, p in get_state_dict(self).items():\n",
    "            if pn.endswith(\"c_proj.weight\"):\n",
    "                print(f'replacing GPT parameter: {pn} with a normal')\n",
    "                p.assign(Tensor.normal(p.shape, mean=0.0, std=0.02 / math.sqrt(2 * config.gpt_n_layer)))\n",
    "\n",
    "    def __call__(self, input_to_wte:Tensor, targets=None) -> Tensor:\n",
    "        b, t, d = input_to_wte.size()\n",
    "        assert (\n",
    "            t <= self.config.gpt_block_size\n",
    "        ), f\"Cannot forward sequence of length {t}, block size is only {self.config.gpt_block_size}\"\n",
    "        # positional encodings that are added to the input embeddings\n",
    "        pos = Tensor.arange(0, t, dtype=dtypes.long).unsqueeze(0)  # shape (1, t)\n",
    "        # forward the GPT model itself\n",
    "        tok_emb = self.wte(input_to_wte)  # token embeddings of shape (b, t, gpt_hidden_dim)\n",
    "        pos_emb = self.wpe(pos)  # position embeddings of shape (1, t, gpt_hidden_dim)\n",
    "        x = (tok_emb + pos_emb).dropout(self.drop).realize()\n",
    "        for block in self.h:\n",
    "            x = block(x).realize()\n",
    "        x = self.ln_f(x).realize()\n",
    "        logits = self.lm_head(x)\n",
    "        return logits\n",
    "    def configure_parameters(self):\n",
    "        \"\"\"\n",
    "        This long function is unfortunately doing something very simple and is being very defensive:\n",
    "        We are separating out all parameters of the model into two buckets: those that will experience\n",
    "        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).\n",
    "        \"\"\"\n",
    "\n",
    "        # separate out all parameters to those that will and won't experience regularizing weight decay\n",
    "        decay = set()\n",
    "        no_decay = set()\n",
    "        whitelist_weight_modules = (nn.Linear,)\n",
    "        blacklist_weight_modules = (nn.LayerNorm, nn.Embedding)\n",
    "        for pn, _p in nn.state.get_state_dict(self).items():\n",
    "            fpn = pn  # full param name\n",
    "            # Kind of strange, but - attn's bias is a *not* \n",
    "            if pn.endswith(\"bias\"):\n",
    "                # all biases will not be decayed\n",
    "                no_decay.add(fpn)\n",
    "                continue\n",
    "            if pn.endswith(\"mlp.0.weight\") or pn.endswith(\"mlp.2.weight\"):\n",
    "                # all biases will not be decayed\n",
    "                decay.add(fpn)\n",
    "                continue\n",
    "            if pn.endswith(\"c_attn.weight\") or pn.endswith(\"c_proj.weight\"):\n",
    "                # all biases will not be decayed\n",
    "                decay.add(fpn)\n",
    "                continue\n",
    "            if pn.endswith(\"ln_1.weight\") or pn.endswith(\"ln_2.weight\"):\n",
    "                # all biases will not be decayed\n",
    "                no_decay.add(fpn)\n",
    "                continue\n",
    "        decay.add('wte.weight')\n",
    "        decay.add('lm_head.weight')\n",
    "        no_decay.add('wpe.weight')\n",
    "        no_decay.add('ln_f.weight')\n",
    "\n",
    "        # validate that we considered every parameter\n",
    "        param_dict = nn.state.get_state_dict(self)\n",
    "        inter_params = decay & no_decay\n",
    "        union_params = decay | no_decay\n",
    "        assert len(inter_params) == 0, \"parameters {} made it into both decay/no_decay sets!\".format(\n",
    "            str(inter_params)\n",
    "        )\n",
    "        assert (\n",
    "            len(param_dict.keys() - union_params) == 0\n",
    "        ), \"parameters {} were not separated into either decay/no_decay set!\".format(\n",
    "            str(param_dict.keys() - union_params),\n",
    "        )\n",
    "        print(f'decay list: {sorted(decay)}')\n",
    "        print(f'no_decay list: {sorted(no_decay)}')\n",
    "\n",
    "        decay = [param_dict[pn] for pn in sorted(decay)]\n",
    "        no_decay = [param_dict[pn] for pn in sorted(no_decay)]\n",
    "        # return the parameters that require weight decay, and the parameters that don't separately.\n",
    "        return decay, no_decay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5954642-3b2e-4cd8-a749-bb07f4483344",
   "metadata": {},
   "source": [
    "# VQBeTRgbEncoder: Image Encoding for VQ-BeTs\n",
    "\n",
    "The `VQBeTRgbEncoder` class is a crucial component of Vector-Quantized Behavior Transformers (VQ-BeTs) designed to process RGB images. This encoder transforms raw image inputs into compact feature vectors, which can then be used by the VQ-BeT model for behavior prediction and generation.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **Flexible Preprocessing**:\n",
    "   - Optional cropping (random or center) based on configuration\n",
    "   - Allows for consistent input sizes and potential data augmentation\n",
    "\n",
    "2. **ResNet Backbone**:\n",
    "   - Uses a ResNet18 architecture for feature extraction\n",
    "   - Can be configured with GroupNorm for improved performance in some scenarios\n",
    "\n",
    "3. **Spatial Softmax Pooling**:\n",
    "   - Converts spatial feature maps into a fixed-size vector\n",
    "   - Helps in focusing on the most important spatial features\n",
    "\n",
    "4. **Configurable Architecture**:\n",
    "   - Uses a `VQBeTConfig` object for flexible setup\n",
    "   - Allows easy experimentation with different encoder configurations\n",
    "\n",
    "5. **Feature Dimensionality Reduction**:\n",
    "   - Final linear layer to adjust the feature dimension\n",
    "   - Crucial for matching the input requirements of the VQ-BeT model\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Visual Input Processing**: Enables VQ-BeTs to work with raw image inputs, essential for visual-based tasks\n",
    "- **Feature Extraction**: Transforms high-dimensional image data into compact, informative feature vectors\n",
    "- **Spatial Understanding**: The combination of CNN and spatial softmax allows the model to capture spatial relationships in the input\n",
    "- **Adaptability**: The configurable nature allows the encoder to be tailored for different visual tasks and environments\n",
    "\n",
    "Let's examine the implementation of the VQBeTRgbEncoder class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9cd96b85-ab1b-4c4f-9f3a-0b9cb4680b1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "###\n",
    "### And now our real hill-climb. Here's where the main part of VQ-BeT lies\n",
    "### This is a knockoff of the same image -> audio encoder that DiffusionPolicy's RGB encoder uses.\n",
    "### I'm gonna assume this is the best way to get ResNet18 to get an image to an encoded latent that the audio encoder understands\n",
    "###\n",
    "\n",
    "import tinygrad.nn as nn\n",
    "from tinygrad import Tensor, dtypes\n",
    "from tinygrad.nn.state import torch_load\n",
    "from tinygrad.helpers import fetch, get_child\n",
    "\n",
    "# allow monkeypatching in layer implementations\n",
    "BatchNorm = nn.BatchNorm2d\n",
    "Conv2d = nn.Conv2d\n",
    "Linear = nn.Linear\n",
    "\n",
    "class ResNet18Backbone:\n",
    "    def __init__(self, use_group_norm=False):\n",
    "        self.conv1 = Conv2d(3, 64, kernel_size=7, stride=2, bias=False, padding=3)\n",
    "        self.bn1 = nn.GroupNorm(num_groups=4, num_channels=64) if use_group_norm else BatchNorm(64)\n",
    "\n",
    "        self.layer1_0 = [\n",
    "            # 0\n",
    "            Conv2d(64, 64, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=4, num_channels=64) if use_group_norm else BatchNorm(64),\n",
    "            Tensor.relu,\n",
    "            Conv2d(64, 64, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=4, num_channels=64) if use_group_norm else BatchNorm(64),\n",
    "        ]\n",
    "        self.layer1_1 = [\n",
    "            # 1\n",
    "            Conv2d(64, 64, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=4, num_channels=64) if use_group_norm else BatchNorm(64),\n",
    "            Tensor.relu,\n",
    "            Conv2d(64, 64, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=4, num_channels=64) if use_group_norm else BatchNorm(64),\n",
    "        ]\n",
    "        self.layer2_0 = [\n",
    "            # 0\n",
    "            Conv2d(64, 128, kernel_size=3, stride=2, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=128) if use_group_norm else BatchNorm(128),\n",
    "            Tensor.relu,\n",
    "            Conv2d(128, 128, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=128) if use_group_norm else BatchNorm(128)\n",
    "        ]\n",
    "        self.layer2_d = [\n",
    "            # 0 downsample\n",
    "            Conv2d(64, 128, kernel_size=1, stride=2, bias=False),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=128) if use_group_norm else BatchNorm(128)\n",
    "        ]\n",
    "        self.layer2_1 = [\n",
    "            # 1\n",
    "            Conv2d(128, 128, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=128) if use_group_norm else BatchNorm(64),\n",
    "            Tensor.relu,\n",
    "            Conv2d(128, 128, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=128) if use_group_norm else BatchNorm(128),\n",
    "        ]\n",
    "        self.layer3_0 = [\n",
    "            # 0\n",
    "            Conv2d(128, 256, kernel_size=3, stride=2, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=16, num_channels=256) if use_group_norm else BatchNorm(256),\n",
    "            Tensor.relu,\n",
    "            Conv2d(256, 256, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=16, num_channels=256) if use_group_norm else BatchNorm(256)\n",
    "        ]\n",
    "        self.layer3_d = [\n",
    "            # 0 downsample\n",
    "            Conv2d(128, 256, kernel_size=1, stride=2, bias=False),\n",
    "            nn.GroupNorm(num_groups=16, num_channels=256) if use_group_norm else BatchNorm(256)\n",
    "        ]\n",
    "        self.layer3_1 = [\n",
    "            # 1\n",
    "            Conv2d(256, 256, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=16, num_channels=256) if use_group_norm else BatchNorm(256),\n",
    "            Tensor.relu,\n",
    "            Conv2d(256, 256, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=16, num_channels=256) if use_group_norm else BatchNorm(256),\n",
    "        ]\n",
    "        self.layer4_0 = [\n",
    "            # 0\n",
    "            Conv2d(256, 512, kernel_size=3, stride=2, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=32, num_channels=512) if use_group_norm else BatchNorm(512),\n",
    "            Tensor.relu,\n",
    "            Conv2d(512, 512, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=32, num_channels=512) if use_group_norm else BatchNorm(512)\n",
    "        ]\n",
    "        self.layer4_d = [\n",
    "            # 0 downsample\n",
    "            Conv2d(256, 512, kernel_size=1, stride=2, bias=False),\n",
    "            nn.GroupNorm(num_groups=32, num_channels=512) if use_group_norm else BatchNorm(512)\n",
    "        ]\n",
    "        self.layer4_1 = [\n",
    "            # 1\n",
    "            Conv2d(512, 512, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=32, num_channels=512) if use_group_norm else BatchNorm(512),\n",
    "            Tensor.relu,\n",
    "            Conv2d(512, 512, kernel_size=3, stride=1, bias=False, padding=1),\n",
    "            nn.GroupNorm(num_groups=32, num_channels=512) if use_group_norm else BatchNorm(512),\n",
    "        ]\n",
    "\n",
    "    \n",
    "    def __call__(self, x:Tensor) -> Tensor:\n",
    "        out = self.bn1(self.conv1(x)).relu()\n",
    "        out = out.max_pool2d(kernel_size=(3,3), stride=2, padding=1, dilation=1)\n",
    "        out = out.sequential(self.layer1_0)\n",
    "        out = out.sequential(self.layer1_1)\n",
    "        out = out.sequential(self.layer2_0) + out.sequential(self.layer2_d)\n",
    "        out = out.sequential(self.layer2_1)\n",
    "        out = out.sequential(self.layer3_0) + out.sequential(self.layer3_d)\n",
    "        out = out.sequential(self.layer3_1)\n",
    "        out = out.sequential(self.layer4_0) + out.sequential(self.layer4_d)\n",
    "        out = out.sequential(self.layer4_1)\n",
    "        return out\n",
    "   \n",
    "def center_crop_f(size: tuple, img: Tensor) -> Tensor:\n",
    "    BS, C, H, W = img.shape\n",
    "    crop_h, crop_w = size\n",
    "\n",
    "    # Calculate the starting points for the crop\n",
    "    start_h = (H - crop_h) // 2\n",
    "    start_w = (W - crop_w) // 2\n",
    "\n",
    "    # Create a boolean mask for the center crop\n",
    "    mask = Tensor.zeros((BS, C, H, W), dtype=dtypes.bool)\n",
    "    mask[:, :, start_h:start_h+crop_h, start_w:start_w+crop_w] = True\n",
    "\n",
    "    # Apply the mask and reshape\n",
    "    X_cropped = Tensor(img.numpy()[mask.numpy()])\n",
    "    return X_cropped.reshape((BS, C, crop_h, crop_w))\n",
    "\n",
    "def make_square_mask_2d(shape, mask_size) -> Tensor:\n",
    "    H, W = shape\n",
    "    low_x = Tensor.randint(1, low=0, high=W-mask_size).item()\n",
    "    low_y = Tensor.randint(1, low=0, high=H-mask_size).item()\n",
    "    return low_y, low_x\n",
    "\n",
    "def random_crop_f_2d(size: tuple, img: Tensor) -> Tensor:\n",
    "    H, W = img.shape\n",
    "    y1, x1 = make_square_mask_2d(img.shape, size[0])\n",
    "    \n",
    "    # Perform the crop\n",
    "    X_cropped = img[y1:y1+size[0], x1:x1+size[0]]\n",
    "    \n",
    "    return X_cropped\n",
    "\n",
    "def make_square_mask_3d(shape, mask_size) -> tuple:\n",
    "    C, H, W = shape\n",
    "    low_x = Tensor.randint(1, low=0, high=W - mask_size).item()\n",
    "    low_y = Tensor.randint(1, low=0, high=H - mask_size).item()\n",
    "    return low_y, low_x\n",
    "\n",
    "def random_crop_f_3d(size: tuple, img: Tensor) -> Tensor:\n",
    "    C, H, W = img.shape\n",
    "    y1, x1 = make_square_mask_3d(img.shape, size[0])\n",
    "    \n",
    "    # Perform the crop\n",
    "    X_cropped = img[:, y1:y1 + size[0], x1:x1 + size[0]]\n",
    "    \n",
    "    return X_cropped\n",
    "\n",
    "def make_square_mask(shape, mask_size) -> Tensor:\n",
    "    BS, _, H, W = shape\n",
    "    low_x = Tensor.randint(BS, low=0, high=W-mask_size, requires_grad=False).reshape(BS,1,1,1)\n",
    "    low_y = Tensor.randint(BS, low=0, high=H-mask_size, requires_grad=False).reshape(BS,1,1,1)\n",
    "    idx_x = Tensor.arange(W, dtype=dtypes.int32, requires_grad=False).reshape((1,1,1,W))\n",
    "    idx_y = Tensor.arange(H, dtype=dtypes.int32, requires_grad=False).reshape((1,1,H,1))\n",
    "    return (idx_x >= low_x) * (idx_x < (low_x + mask_size)) * (idx_y >= low_y) * (idx_y < (low_y + mask_size))\n",
    "\n",
    "def random_crop_f(size:tuple, img: Tensor) -> Tensor:\n",
    "    mask = make_square_mask(img.shape, size[0])\n",
    "    mask = mask.expand((-1,3,-1,-1))\n",
    "    X_cropped = Tensor(img.numpy()[mask.numpy()])\n",
    "    return X_cropped.reshape((-1, 3, size[0], size[0]))\n",
    "\n",
    "class VQBeTRgbEncoder:\n",
    "    \"\"\"Encoder an RGB image into a 1D feature vector.\n",
    "\n",
    "    Includes the ability to normalize and crop the image first.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config: VQBeTConfig):\n",
    "        # Set up optional preprocessing.\n",
    "        if config.crop_shape is not None:\n",
    "            self.do_crop = True\n",
    "            # Always use center crop for eval\n",
    "            self.center_crop = lambda x: center_crop_f(config.crop_shape, x)\n",
    "            if config.crop_is_random:\n",
    "                self.maybe_random_crop = lambda x: random_crop_f(config.crop_shape, x)\n",
    "            else:\n",
    "                self.maybe_random_crop = self.center_crop\n",
    "        else:\n",
    "            self.do_crop = False\n",
    "\n",
    "        # Set up backbone.\n",
    "        # Note: This assumes that the layer4 feature map is children()[-3]\n",
    "        # TODO(alexander-soare): Use a safer alternative.\n",
    "        self.backbone = ResNet18Backbone(use_group_norm=config.use_group_norm) \n",
    "        #self.backbone = generate_resnet_model(config.vision_backbone, use_group_norm=config.use_group_norm)\n",
    "        if config.use_group_norm and config.pretrained_backbone_weights:\n",
    "            raise ValueError(\n",
    "                \"You can't replace BatchNorm in a pretrained model without ruining the weights!\"\n",
    "            )\n",
    "\n",
    "        # Set up pooling and final layers.\n",
    "        # Use a dry run to get the feature map shape.\n",
    "        # The dummy input should take the number of image channels from `config.input_shapes` and it should\n",
    "        # use the height and width from `config.crop_shape` if it is provided, otherwise it should use the\n",
    "        # height and width from `config.input_shapes`.\n",
    "        image_keys = [k for k in config.input_shapes if k.startswith(\"observation.image\")]\n",
    "        # Note: we have a check in the config class to make sure all images have the same shape.\n",
    "        image_key = image_keys[0]\n",
    "        dummy_input_h_w = (\n",
    "            config.crop_shape if config.crop_shape is not None else config.input_shapes[image_key][1:]\n",
    "        )\n",
    "        dummy_input = Tensor.zeros(1, config.input_shapes[image_key][0], *dummy_input_h_w, requires_grad=False)\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        dummy_feature_map = self.backbone(dummy_input)\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        feature_map_shape = tuple(dummy_feature_map.shape[1:])\n",
    "        self.pool = SpatialSoftmax(feature_map_shape, num_kp=config.spatial_softmax_num_keypoints)\n",
    "        self.feature_dim = config.spatial_softmax_num_keypoints * 2\n",
    "        self.out = nn.Linear(config.spatial_softmax_num_keypoints * 2, self.feature_dim)\n",
    "\n",
    "    def __call__(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: (B, C, H, W) image tensor with pixel values in [0, 1].\n",
    "        Returns:\n",
    "            (B, D) image feature.\n",
    "        \"\"\"\n",
    "        # Preprocess: maybe crop (if it was set up in the __init__).\n",
    "        if self.do_crop:\n",
    "            if Tensor.training:  # noqa: SIM108\n",
    "                x = self.maybe_random_crop(x)\n",
    "            else:\n",
    "                # Always use center crop for eval.\n",
    "                x = self.center_crop(x)\n",
    "        # Extract backbone feature.\n",
    "        x = self.pool(self.backbone(x)).flatten(start_dim=1)\n",
    "        # Final linear layer with non-linearity.\n",
    "        x = self.out(x).relu()\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8486289-287d-495e-b7aa-b7ee6c5cb502",
   "metadata": {},
   "source": [
    "# VQBeTScheduler: Learning Rate Scheduler for VQ-BeTs\n",
    "\n",
    "The `VQBeTScheduler` class implements a custom learning rate scheduler specifically designed for training Vector-Quantized Behavior Transformers (VQ-BeTs). This scheduler manages the learning rate throughout the training process, adapting it based on the current training phase and progress.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **Two-Phase Training Support**:\n",
    "   - Maintains a constant learning rate during VQ-VAE training\n",
    "   - Implements a more complex schedule for the main VQ-BeT training phase\n",
    "   - Essential for handling the distinct requirements of each training phase\n",
    "\n",
    "2. **Warmup Period**:\n",
    "   - Gradually increases the learning rate at the start of the main training phase\n",
    "   - Helps stabilize training in the early stages, particularly important for transformer models\n",
    "\n",
    "3. **Cosine Decay with Restarts**:\n",
    "   - Implements a cyclical learning rate schedule after the warmup period\n",
    "   - Allows for better exploration of the loss landscape, potentially leading to improved convergence\n",
    "\n",
    "4. **Configurable Parameters**:\n",
    "   - Uses configuration object for flexible setup (number of steps, warmup period, etc.)\n",
    "   - Enables easy experimentation with different scheduling strategies\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Phase-Specific Optimization**: Addresses the different learning dynamics of the VQ-VAE and transformer components\n",
    "- **Training Stability**: The warmup period helps prevent early training instability, especially important for the transformer part\n",
    "- **Improved Convergence**: Cyclical learning rates can lead to better final performance and faster convergence\n",
    "- **Flexibility**: Allows for easy adjustment of the learning rate schedule to suit different datasets and task complexities\n",
    "\n",
    "Let's examine the implementation of the VQBeTScheduler class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d41acc66-599e-409a-8555-50ceb078e423",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tinygrad.nn.optim import Optimizer\n",
    "from tinygrad.tensor import Tensor\n",
    "\n",
    "class VQBeTScheduler:\n",
    "    def __init__(self, optimizer_group, cfg):\n",
    "        self.optimizer_group = optimizer_group\n",
    "        self.n_vqvae_training_steps = cfg.n_vqvae_training_steps\n",
    "        self.base_lr = 1.0e-4\n",
    "        self.num_warmup_steps = 500\n",
    "        self.num_training_steps = 250000\n",
    "        self.num_cycles = 0.5\n",
    "        self.epoch_counter = 0\n",
    "        \n",
    "    def step(self) -> None:\n",
    "        self.epoch_counter += 1\n",
    "        for optimizer in self.optimizer_group:\n",
    "            optimizer.lr.assign(Tensor([self.base_lr * self.get_lr()], requires_grad=False)).realize()\n",
    "\n",
    "\n",
    "    # Breakdown of this:\n",
    "    # If the current step is less than the VQ-VAE training step, just don't decay\n",
    "    # If the current step is greater, then you see:\n",
    "    # 1. how much greater\n",
    "    # 2. if it's warmed up yet. If not, return a percentage of the warmup\n",
    "    # 3. If it's warm, do a cos based on the progress and the number of cycles. Add, never let it get above 1. Should be cyclical\n",
    "    def get_lr(self):\n",
    "        if self.epoch_counter < self.n_vqvae_training_steps:\n",
    "            return 1.0\n",
    "        else:\n",
    "            current_step = self.epoch_counter - self.n_vqvae_training_steps\n",
    "            if current_step < self.num_warmup_steps:\n",
    "                return float(current_step) / float(max(1, self.num_warmup_steps))\n",
    "            progress = float(current_step - self.num_warmup_steps) / float(\n",
    "                max(1, self.num_training_steps - self.num_warmup_steps)\n",
    "            )\n",
    "            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(self.num_cycles) * 2.0 * progress)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de2ac07f-bbb6-4955-a6b1-3b7dde64c3bf",
   "metadata": {},
   "source": [
    "# VQBeTOptimizer: Customized Optimization for VQ-BeTs\n",
    "\n",
    "The `VQBeTOptimizer` class implements a specialized optimization strategy for Vector-Quantized Behavior Transformers (VQ-BeTs). This optimizer is crucial for effectively training the various components of the VQ-BeT model, each with its own learning requirements.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **Parameter Grouping**:\n",
    "   - Separates parameters into different groups (VQ-VAE, decay, no decay)\n",
    "   - Allows for different optimization strategies for different parts of the model\n",
    "\n",
    "2. **AdamW Optimizer**:\n",
    "   - Uses AdamW, which combines the benefits of Adam with weight decay regularization\n",
    "   - Crucial for training large transformer models effectively\n",
    "\n",
    "3. **Different Learning Rates**:\n",
    "   - Applies different learning rates to VQ-VAE and other components\n",
    "   - Enables fine-tuned optimization for each part of the model\n",
    "\n",
    "4. **Weight Decay Configuration**:\n",
    "   - Applies weight decay selectively to certain parameter groups\n",
    "   - Helps in preventing overfitting while maintaining important features\n",
    "\n",
    "5. **Flexible Architecture Support**:\n",
    "   - Adapts to different VQ-BeT configurations (e.g., sequential selection)\n",
    "   - Ensures the optimizer can handle various model architectures\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Component-Specific Optimization**: Different parts of VQ-BeTs (e.g., VQ-VAE, transformer) may require different optimization strategies\n",
    "- **Regularization Control**: Selective application of weight decay helps in managing the complexity of the model\n",
    "- **Training Stability**: Carefully tuned optimization parameters contribute to stable and effective training\n",
    "- **Adaptability**: The flexible design allows for easy adaptation to different VQ-BeT variants and configurations\n",
    "\n",
    "Let's examine the implementation of the VQBeTOptimizer class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5e3f0191-d980-4960-b630-fa644f6a139c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I...don't think this merits its own optimizer.\n",
    "# Instead, segregate it into the VQBeT Optimization step.\n",
    "\n",
    "from tinygrad.nn.state import get_parameters\n",
    "from tinygrad.nn.optim import *\n",
    "class VQBeTOptimizer:\n",
    "    @staticmethod\n",
    "    def create(policy, cfg):\n",
    "        vqvae_params = (\n",
    "            get_parameters(policy.vqbet.action_head.vqvae_model.encoder)\n",
    "            + get_parameters(policy.vqbet.action_head.vqvae_model.decoder)\n",
    "            + get_parameters(policy.vqbet.action_head.vqvae_model.vq_layer)\n",
    "        )\n",
    "        decay_params, no_decay_params = policy.vqbet.policy.configure_parameters()\n",
    "        decay_params = (\n",
    "            decay_params\n",
    "            + get_parameters(policy.vqbet.rgb_encoder)\n",
    "            + get_parameters(policy.vqbet.state_projector)\n",
    "            + get_parameters(policy.vqbet.rgb_feature_projector)\n",
    "            + [policy.vqbet.action_token]\n",
    "            + get_parameters(policy.vqbet.action_head.map_to_cbet_preds_offset)\n",
    "        )\n",
    "\n",
    "        if cfg.sequentially_select:\n",
    "            decay_params = (\n",
    "                decay_params\n",
    "                + get_parameters(policy.vqbet.action_head.map_to_cbet_preds_primary_bin)\n",
    "                + get_parameters(policy.vqbet.action_head.map_to_cbet_preds_secondary_bin)\n",
    "            )\n",
    "        else:\n",
    "            decay_params = decay_params + get_parameters(policy.vqbet.action_head.map_to_cbet_preds_bin)\n",
    "\n",
    "        optim_and_param_groups = [\n",
    "            AdamW(\n",
    "                decay_params, \n",
    "                lr=1.0e-4, \n",
    "                b1=0.95,\n",
    "                b2=0.999,\n",
    "                eps=1.0e-8,\n",
    "                weight_decay=1.0e-6\n",
    "            ),\n",
    "            AdamW(\n",
    "                vqvae_params, \n",
    "                lr=1.0e-3, \n",
    "                b1=0.9,\n",
    "                b2=0.999,\n",
    "                eps=1.0e-8,\n",
    "                weight_decay=0.0001\n",
    "            ),\n",
    "            AdamW(\n",
    "                no_decay_params, \n",
    "                lr=1.0e-4, \n",
    "                b1=0.95,\n",
    "                b2=0.999,\n",
    "                eps=1.0e-8,\n",
    "                weight_decay=0.0\n",
    "            ),\n",
    "            decay_params, vqvae_params, no_decay_params\n",
    "        ]\n",
    "        return optim_and_param_groups"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec9f4df-a550-41d4-862c-c17ab78ecf22",
   "metadata": {},
   "source": [
    "# VQBeTHead: Action Prediction Component for VQ-BeTs\n",
    "\n",
    "The `VQBeTHead` class is a crucial component of Vector-Quantized Behavior Transformers (VQ-BeTs), responsible for predicting actions based on the output of the transformer layers. This class integrates the VQ-VAE model with prediction heads for both discrete code selection and continuous offset prediction.\n",
    "\n",
    "Key Features and Their Importance for VQ-BeTs:\n",
    "\n",
    "1. **VQ-VAE Integration**:\n",
    "   - Incorporates a pre-trained VQ-VAE model for action discretization and reconstruction\n",
    "   - Essential for bridging the gap between discrete tokens and continuous actions\n",
    "\n",
    "2. **Bin Prediction**:\n",
    "   - Predicts probabilities for each discrete code in the VQ-VAE codebook\n",
    "   - Supports both simultaneous and sequential code selection strategies\n",
    "\n",
    "3. **Offset Prediction**:\n",
    "   - Predicts continuous offsets for fine-tuning the selected discrete actions\n",
    "   - Allows for more precise action generation within the quantized action space\n",
    "\n",
    "4. **Flexible Architecture**:\n",
    "   - Adapts to different VQ-VAE configurations (e.g., number of layers, codebook size)\n",
    "   - Supports both simultaneous and sequential code selection strategies\n",
    "\n",
    "5. **Focal Loss**:\n",
    "   - Uses Focal Loss for training the bin prediction head\n",
    "   - Helps in addressing class imbalance issues in code prediction\n",
    "\n",
    "Why is this crucial for VQ-BeTs?\n",
    "\n",
    "- **Action Generation**: Enables the model to generate both discrete and continuous components of actions\n",
    "- **Hierarchical Prediction**: The two-stage prediction (bin and offset) allows for more nuanced action generation\n",
    "- **Adaptability**: Can work with different VQ-VAE configurations and selection strategies\n",
    "- **Fine-grained Control**: The offset prediction allows for precise adjustments to the quantized actions\n",
    "\n",
    "Let's examine the implementation of the VQBeTHead class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "11a313df-bc03-41e5-8a45-56a4d5903f17",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Main method (finally)\n",
    "class VQBeTHead:\n",
    "    def __init__(self, config: VQBeTConfig):\n",
    "        \"\"\"\n",
    "        VQBeTHead takes output of GPT layers, and pass the feature through bin prediction head (`self.map_to_cbet_preds_bin`), and offset prediction head (`self.map_to_cbet_preds_offset`)\n",
    "\n",
    "        self.map_to_cbet_preds_bin: outputs probability of each code (for each layer).\n",
    "            The input dimension of `self.map_to_cbet_preds_bin` is same with the output of GPT,\n",
    "            and the output dimension of `self.map_to_cbet_preds_bin` is `self.vqvae_model.vqvae_num_layers (=fixed as 2) * self.config.vqvae_n_embed`.\n",
    "            if the agent select the code sequentially, we use self.map_to_cbet_preds_primary_bin and self.map_to_cbet_preds_secondary_bin instead of self._map_to_cbet_preds_bin.\n",
    "\n",
    "        self.map_to_cbet_preds_offset: output the predicted offsets for all the codes in all the layers.\n",
    "            The input dimension of ` self.map_to_cbet_preds_offset` is same with the output of GPT,\n",
    "            and the output dimension of ` self.map_to_cbet_preds_offset` is `self.vqvae_model.vqvae_num_layers (=fixed as 2) * self.config.vqvae_n_embed * config.action_chunk_size * config.output_shapes[\"action\"][0]`.\n",
    "        \"\"\"\n",
    "\n",
    "        self.config = config\n",
    "        # init vqvae\n",
    "        self.vqvae_model = VqVae(config)\n",
    "        if config.sequentially_select:\n",
    "            self.map_to_cbet_preds_primary_bin = MLP(\n",
    "                in_channels=config.gpt_output_dim,\n",
    "                hidden_channels=[self.config.vqvae_n_embed],\n",
    "            )\n",
    "            self.map_to_cbet_preds_secondary_bin = MLP(\n",
    "                in_channels=config.gpt_output_dim + self.config.vqvae_n_embed,\n",
    "                hidden_channels=[self.config.vqvae_n_embed],\n",
    "            )\n",
    "        else:\n",
    "            self.map_to_cbet_preds_bin = MLP(\n",
    "                in_channels=config.gpt_output_dim,\n",
    "                hidden_channels=[self.vqvae_model.vqvae_num_layers * self.config.vqvae_n_embed],\n",
    "            )\n",
    "        self.map_to_cbet_preds_offset = MLP(\n",
    "            in_channels=config.gpt_output_dim,\n",
    "            hidden_channels=[\n",
    "                self.vqvae_model.vqvae_num_layers\n",
    "                * self.config.vqvae_n_embed\n",
    "                * config.action_chunk_size\n",
    "                * config.output_shapes[\"action\"][0],\n",
    "            ],\n",
    "        )\n",
    "        # loss\n",
    "        self._focal_loss_fn = FocalLoss(gamma=2.0)\n",
    "\n",
    "    def discretize(self, n_vqvae_training_steps, actions):\n",
    "        # Resize the action sequence data to fit the action chunk size using a sliding window approach.        \n",
    "        actions = Tensor.cat(\n",
    "            *[\n",
    "                actions[:, j : j + self.config.action_chunk_size, :]\n",
    "                for j in range(actions.shape[1] + 1 - self.config.action_chunk_size)\n",
    "            ],\n",
    "            dim=0,\n",
    "        )\n",
    "        # `actions` is a tensor of shape (new_batch, action_chunk_size, action_dim) where new_batch is the number of possible chunks created from the original sequences using the sliding window.\n",
    "\n",
    "        loss, metric = self.vqvae_model.vqvae_forward(actions)\n",
    "        n_different_codes = Tensor(sum(\n",
    "            [len(np.unique(metric[2][:, i].numpy())) for i in range(self.vqvae_model.vqvae_num_layers)]\n",
    "        ), requires_grad=False)\n",
    "        n_different_combinations = len(np.unique(metric[2].numpy(), axis=0))\n",
    "        recon_l1_error = metric[0].detach().item()\n",
    "        self.vqvae_model.optimized_steps += 1\n",
    "        # if we updated RVQ more than `n_vqvae_training_steps` steps, we freeze the RVQ part.\n",
    "        if self.vqvae_model.optimized_steps >= n_vqvae_training_steps:\n",
    "            self.vqvae_model.discretized = True\n",
    "            self.vqvae_model.vq_layer.freeze_codebook = True\n",
    "            print(\"Finished discretizing action data!\")\n",
    "            for param in nn.state.get_parameters(self.vqvae_model.vq_layer):\n",
    "                param.requires_grad = False\n",
    "        return loss, n_different_codes, n_different_combinations, recon_l1_error\n",
    "\n",
    "    def __call__(self, x:Tensor, **kwargs) -> Tensor:\n",
    "        # N is the batch size, and T is number of action query tokens, which are process through same GPT\n",
    "        N, T, _ = x.shape\n",
    "        # we calculate N and T side parallely. Thus, the dimensions would be\n",
    "        # (batch size * number of action query tokens, action chunk size, action dimension)\n",
    "        x = einops.rearrange(x, \"N T WA -> (N T) WA\")\n",
    "\n",
    "        # sample offsets\n",
    "        cbet_offsets = self.map_to_cbet_preds_offset(x)\n",
    "        cbet_offsets = einops.rearrange(\n",
    "            cbet_offsets,\n",
    "            \"(NT) (G C WA) -> (NT) G C WA\",\n",
    "            G=self.vqvae_model.vqvae_num_layers,\n",
    "            C=self.config.vqvae_n_embed,\n",
    "        )\n",
    "        # if self.config.sequentially_select is True, bin prediction head first sample the primary code, and then sample secondary code\n",
    "        if self.config.sequentially_select:\n",
    "            cbet_primary_logits = self.map_to_cbet_preds_primary_bin(x)\n",
    "\n",
    "            # select primary bin first\n",
    "            cbet_primary_probs = (cbet_primary_logits / self.config.bet_softmax_temperature).softmax()\n",
    "            NT, choices = cbet_primary_probs.shape\n",
    "            sampled_primary_centers = einops.rearrange(\n",
    "                cbet_primary_probs.view(-1, choices).multinomial(num_samples=1),\n",
    "                \"(NT) 1 -> NT\",\n",
    "                NT=NT,\n",
    "            )\n",
    "\n",
    "            cbet_secondary_logits = self.map_to_cbet_preds_secondary_bin(\n",
    "                Tensor.cat(\n",
    "                    x, sampled_primary_centers.one_hot(num_classes=self.config.vqvae_n_embed),\n",
    "                    axis=1,\n",
    "                )\n",
    "            )\n",
    "            cbet_secondary_probs = (cbet_secondary_logits / self.config.bet_softmax_temperature).softmax(\n",
    "                dim=-1\n",
    "            )\n",
    "            sampled_secondary_centers = einops.rearrange(\n",
    "                cbet_secondary_probs.view(-1, choices).multinomial(num_samples=1),\n",
    "                \"(NT) 1 -> NT\",\n",
    "                NT=NT,\n",
    "            )\n",
    "            sampled_centers = Tensor.stack(sampled_primary_centers, sampled_secondary_centers, dim=1)\n",
    "            cbet_logits = Tensor.stack(cbet_primary_logits, cbet_secondary_logits, dim=1)\n",
    "        # if self.config.sequentially_select is False, bin prediction head samples primary and secondary code at once.\n",
    "        else:\n",
    "            cbet_logits = self.map_to_cbet_preds_bin(x)\n",
    "            cbet_logits = einops.rearrange(\n",
    "                cbet_logits, \"(NT) (G C) -> (NT) G C\", G=self.vqvae_model.vqvae_num_layers\n",
    "            )\n",
    "            cbet_probs = (cbet_logits / self.config.bet_softmax_temperature).softmax(axis=-1)\n",
    "            NT, G, choices = cbet_probs.shape\n",
    "            sampled_centers = einops.rearrange(\n",
    "                cbet_probs.view(-1, choices).multinomial(num_samples=1),\n",
    "                \"(NT G) 1 -> NT G\",\n",
    "                NT=NT,\n",
    "            ).contiguous()\n",
    "\n",
    "        indices = (\n",
    "            Tensor.arange(NT, requires_grad=False).unsqueeze(1),\n",
    "            Tensor.arange(self.vqvae_model.vqvae_num_layers, requires_grad=False).unsqueeze(0),\n",
    "            sampled_centers.detach(),\n",
    "        )\n",
    "        # Use advanced indexing to sample the values (Extract the only offsets corresponding to the sampled codes.)\n",
    "        sampled_offsets = cbet_offsets[indices]\n",
    "        # Then, sum the offsets over the RVQ layers to get a net offset for the bin prediction\n",
    "        sampled_offsets = sampled_offsets.sum(axis=1)\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        # Get the centroids (= vectors corresponding to the codes) of each layer to pass it through RVQ decoder\n",
    "        return_decoder_input = self.vqvae_model.get_embeddings_from_code(sampled_centers)\n",
    "        # pass the centroids through decoder to get actions.\n",
    "        decoded_action = self.vqvae_model.get_action_from_latent(return_decoder_input)\n",
    "        decoded_action = decoded_action.detach()\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        # reshaped extracted offset to match with decoded centroids\n",
    "        sampled_offsets = einops.rearrange(\n",
    "            sampled_offsets, \"NT (W A) -> NT W A\", W=self.config.action_chunk_size\n",
    "        )\n",
    "        # add offset and decoded centroids\n",
    "        predicted_action = decoded_action + sampled_offsets\n",
    "        predicted_action = einops.rearrange(\n",
    "            predicted_action,\n",
    "            \"(N T) W A -> N T (W A)\",\n",
    "            N=N,\n",
    "            T=T,\n",
    "            W=self.config.action_chunk_size,\n",
    "        )\n",
    "\n",
    "        return {\n",
    "            \"cbet_logits\": cbet_logits,\n",
    "            \"predicted_action\": predicted_action,\n",
    "            \"sampled_centers\": sampled_centers,\n",
    "            \"decoded_action\": decoded_action,\n",
    "        }\n",
    "\n",
    "    def loss_fn(self, pred, target, **kwargs):\n",
    "        \"\"\"\n",
    "        for given ground truth action values (target), and prediction (pred) this function calculates the overall loss.\n",
    "\n",
    "        predicted_action: predicted action chunk (offset + decoded centroids)\n",
    "        sampled_centers: sampled centroids (code of RVQ)\n",
    "        decoded_action: decoded action, which is produced by passing sampled_centers through RVQ decoder\n",
    "        NT: batch size * T\n",
    "        T: number of action query tokens, which are process through same GPT\n",
    "        cbet_logits: probability of all codes in each layer\n",
    "        \"\"\"\n",
    "        action_seq = target\n",
    "        predicted_action = pred[\"predicted_action\"]\n",
    "        sampled_centers = pred[\"sampled_centers\"]\n",
    "        decoded_action = pred[\"decoded_action\"]\n",
    "        NT = predicted_action.shape[0] * predicted_action.shape[1]\n",
    "\n",
    "        cbet_logits = pred[\"cbet_logits\"]\n",
    "\n",
    "        predicted_action = einops.rearrange(\n",
    "            predicted_action, \"N T (W A) -> (N T) W A\", W=self.config.action_chunk_size\n",
    "        )\n",
    "\n",
    "        action_seq = einops.rearrange(action_seq, \"N T W A -> (N T) W A\")\n",
    "\n",
    "        # offset loss is L1 distance between the predicted action and ground truth action\n",
    "        offset_loss = (action_seq - predicted_action).abs().mean()\n",
    "        \n",
    "        # Figure out the loss for the actions.\n",
    "        # First, we need to find the closest cluster center for each ground truth action.\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        state_vq, action_bins = self.vqvae_model.get_code(action_seq)  # action_bins: NT, G\n",
    "        action_bins = action_bins.detach()\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "\n",
    "        # Now we can compute the loss.\n",
    "\n",
    "        # calculate primary code prediction loss\n",
    "        cbet_loss1 = self._focal_loss_fn(\n",
    "            cbet_logits[:, 0, :],\n",
    "            action_bins[:, 0],\n",
    "        )\n",
    "        # calculate secondary code prediction loss\n",
    "        cbet_loss2 = self._focal_loss_fn(\n",
    "            cbet_logits[:, 1, :],\n",
    "            action_bins[:, 1],\n",
    "        )\n",
    "        # add all the prediction loss\n",
    "        cbet_loss = cbet_loss1 * self.config.primary_code_loss_weight + cbet_loss2 * self.config.secondary_code_loss_weight\n",
    "\n",
    "        equal_primary_code_rate = (action_bins[:, 0] == sampled_centers[:, 0]).int().sum() / (NT)\n",
    "        equal_secondary_code_rate = (action_bins[:, 1] == sampled_centers[:, 1]).int().sum() / (NT)\n",
    "\n",
    "        action_mse_error = (action_seq - predicted_action).square().mean()\n",
    "        vq_action_error = (action_seq - decoded_action).abs().mean()\n",
    "        offset_action_error = (action_seq - predicted_action).abs().mean()\n",
    "        action_error_max = (action_seq - predicted_action).abs().max()\n",
    "\n",
    "        loss = cbet_loss + self.config.offset_loss_weight * offset_loss\n",
    "\n",
    "        loss_dict = {\n",
    "            \"loss\": loss,\n",
    "            \"classification_loss\": cbet_loss.detach().item(),\n",
    "            \"offset_loss\": offset_loss.detach().item(),\n",
    "            \"equal_primary_code_rate\": equal_primary_code_rate.detach().item(),\n",
    "            \"equal_secondary_code_rate\": equal_secondary_code_rate.detach().item(),\n",
    "            \"vq_action_error\": vq_action_error.detach().item(),\n",
    "            \"offset_action_error\": offset_action_error.detach().item(),\n",
    "            \"action_error_max\": action_error_max.detach().item(),\n",
    "            \"action_mse_error\": action_mse_error.detach().item(),\n",
    "        }\n",
    "        return loss_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ebff820-710e-4fd8-a664-aaa3d73e4838",
   "metadata": {},
   "source": [
    "# VQBeTModel: The Core Architecture of Vector-Quantized Behavior Transformers\n",
    "\n",
    "The `VQBeTModel` class represents the complete architecture of Vector-Quantized Behavior Transformers (VQ-BeTs). This model integrates various components to process visual and state observations, generate action predictions, and handle the vector quantization of the action space.\n",
    "\n",
    "Key Components and Their Roles:\n",
    "\n",
    "1. **RGB Encoder**:\n",
    "   - Processes image observations into compact feature vectors\n",
    "   - Essential for handling visual input in tasks like robotic manipulation or autonomous driving\n",
    "\n",
    "2. **State and RGB Feature Projectors**:\n",
    "   - Project state observations and RGB features to the appropriate dimension for the GPT model\n",
    "   - Ensure all inputs are compatible with the transformer architecture\n",
    "\n",
    "3. **Action Query Token**:\n",
    "   - Serves as a prompt for querying action chunks\n",
    "   - Crucial for the model to distinguish between observation and action prediction tasks\n",
    "\n",
    "4. **GPT (Generative Pre-trained Transformer)**:\n",
    "   - Core component for sequence modeling\n",
    "   - Processes the sequence of observations and action queries\n",
    "\n",
    "5. **VQBeTHead**:\n",
    "   - Handles the prediction of discrete codes and continuous offsets\n",
    "   - Integrates the Vector Quantized Variational Autoencoder (VQ-VAE) for action discretization\n",
    "\n",
    "6. **Training Phases**:\n",
    "   - Phase 1: Discretize actions using Residual VQ\n",
    "   - Phase 2: Train the full model to predict future actions\n",
    "\n",
    "Why is this architecture crucial for VQ-BeTs?\n",
    "\n",
    "- **End-to-End Learning**: Integrates perception, sequence modeling, and action generation in a single architecture\n",
    "- **Hierarchical Representation**: Combines continuous and discrete representations of actions\n",
    "- **Long-term Dependencies**: The transformer architecture allows for capturing long-range dependencies in behavior sequences\n",
    "- **Flexibility**: Can handle various types of inputs (images, states) and output complex, multi-step action predictions\n",
    "\n",
    "Let's examine the implementation of the VQBeTModel class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9899cc57-eedf-4bc6-82b1-c10b46a3c1ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VQBeTModel:\n",
    "    \"\"\"VQ-BeT: The underlying neural network for VQ-BeT\n",
    "\n",
    "    Note: In this code we use the terms `rgb_encoder`, 'policy', `action_head`. The meanings are as follows.\n",
    "        - The `rgb_encoder` process rgb-style image observations to one-dimensional embedding vectors\n",
    "        - A `policy` is a minGPT architecture, that takes observation sequences and action query tokens to generate `features`.\n",
    "        - These `features` pass through the action head, which passes through the code prediction, offset prediction head,\n",
    "        and finally generates a prediction for the action chunks.\n",
    "\n",
    "        -------------------------------** legend **-------------------------------\n",
    "        │   n = n_obs_steps, p = n_action_pred_token, c = action_chunk_size)   │\n",
    "        │   o_{t} : visual observation at timestep {t}                           │\n",
    "        │   s_{t} : state observation at timestep {t}                            │\n",
    "        │   a_{t} : action at timestep {t}                                       │\n",
    "        │   A_Q : action_query_token                                             │\n",
    "        --------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "        Training Phase 1. Discretize action using Residual VQ (for config.n_vqvae_training_steps steps)\n",
    "\n",
    "\n",
    "        ┌─────────────────┐            ┌─────────────────┐            ┌─────────────────┐\n",
    "        │                 │            │                 │            │                 │\n",
    "        │   RVQ encoder   │    ─►      │     Residual    │    ─►      │   RVQ Decoder   │\n",
    "        │ (a_{t}~a_{t+p}) │            │  Code Quantizer │            │                 │\n",
    "        │                 │            │                 │            │                 │\n",
    "        └─────────────────┘            └─────────────────┘            └─────────────────┘\n",
    "\n",
    "        Training Phase 2.\n",
    "\n",
    "          timestep {t-n+1}   timestep {t-n+2}                timestep {t}\n",
    "            ┌─────┴─────┐     ┌─────┴─────┐                 ┌─────┴─────┐\n",
    "\n",
    "        o_{t-n+1}         o_{t-n+2}           ...         o_{t}\n",
    "            │                 │                             │\n",
    "            │ s_{t-n+1}       │ s_{t-n+2}         ...       │   s_{t}           p\n",
    "            │     │           │     │                       │     │     ┌───────┴───────┐\n",
    "            │     │    A_Q    │     │    A_Q          ...   │     │    A_Q     ...     A_Q\n",
    "            │     │     │     │     │     │                 │     │     │               │\n",
    "        ┌───▼─────▼─────▼─────▼─────▼─────▼─────────────────▼─────▼─────▼───────────────▼───┐\n",
    "        │                                                                                   │\n",
    "        │                                       GPT                                         │       =>    policy\n",
    "        │                                                                                   │\n",
    "        └───────────────▼─────────────────▼─────────────────────────────▼───────────────▼───┘\n",
    "                        │                 │                             │               │\n",
    "                    ┌───┴───┐         ┌───┴───┐                     ┌───┴───┐       ┌───┴───┐\n",
    "                  code    offset    code    offset                code    offset  code    offset\n",
    "                    ▼       │         ▼       │                     ▼       │       ▼       │       =>    action_head\n",
    "               RVQ Decoder  │    RVQ Decoder  │                RVQ Decoder  │  RVQ Decoder  │\n",
    "                    └── + ──┘         └── + ──┘                     └── + ──┘       └── + ──┘\n",
    "                        ▼                 ▼                             ▼               ▼\n",
    "                   action chunk      action chunk                  action chunk     action chunk\n",
    "                    a_{t-n+1} ~       a_{t-n+2} ~                   a_{t} ~     ...  a_{t+p-1} ~\n",
    "                     a_{t-n+c}         a_{t-n+c+1}                   a_{t+c-1}        a_{t+p+c-1}\n",
    "\n",
    "                                                                        ▼\n",
    "                                                      ONLY this chunk is used in rollout!\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config: VQBeTConfig):\n",
    "        self.config = config\n",
    "\n",
    "        self.rgb_encoder = VQBeTRgbEncoder(config)\n",
    "        self.num_images = len([k for k in config.input_shapes if k.startswith(\"observation.image\")])\n",
    "        # This action query token is used as a prompt for querying action chunks. Please refer to \"A_Q\" in the image above.\n",
    "        # Note: During the forward pass, this token is repeated as many times as needed. The authors also experimented with initializing the necessary number of tokens independently and observed inferior results.\n",
    "        self.action_token = Tensor.randn(1, 1, self.config.gpt_input_dim)\n",
    "\n",
    "        # To input state and observation features into GPT layers, we first project the features to fit the shape of input size of GPT.\n",
    "        self.state_projector = MLP(\n",
    "            config.input_shapes[\"observation.state\"][0], hidden_channels=[self.config.gpt_input_dim]\n",
    "        )\n",
    "        self.rgb_feature_projector = MLP(\n",
    "            self.rgb_encoder.feature_dim, hidden_channels=[self.config.gpt_input_dim]\n",
    "        )\n",
    "\n",
    "        # GPT part of VQ-BeT\n",
    "        self.policy = GPT(config)\n",
    "        # bin prediction head / offset prediction head part of VQ-BeT\n",
    "        self.action_head = VQBeTHead(config)\n",
    "\n",
    "        # Action tokens for: each observation step, the current action token, and all future action tokens.\n",
    "        num_tokens = self.config.n_action_pred_token + self.config.n_obs_steps - 1\n",
    "        ranges = [np.arange(i, i + self.config.action_chunk_size) for i in range(num_tokens)]\n",
    "        self.select_target_actions_indices = Tensor(np.vstack(ranges), requires_grad=False)\n",
    "\n",
    "    def call_first_step(self, observation_images, observation_state):\n",
    "        batch_size, n_obs_steps = observation_state.shape[:2]\n",
    "        assert n_obs_steps == self.config.n_obs_steps\n",
    "\n",
    "        # Extract image feature (first combine batch and sequence dims).\n",
    "        img_features = self.rgb_encoder(\n",
    "            einops.rearrange(observation_images, \"b s n ... -> (b s n) ...\")\n",
    "        )\n",
    "        # Separate batch and sequence dims.\n",
    "        img_features = einops.rearrange(\n",
    "            img_features, \"(b s n) ... -> b s n ...\", b=batch_size, s=n_obs_steps, n=self.num_images\n",
    "        )\n",
    "\n",
    "        # Arrange prior and current observation step tokens as shown in the class docstring.\n",
    "        # First project features to token dimension.\n",
    "        rgb_tokens = self.rgb_feature_projector(\n",
    "            img_features\n",
    "        )  # (batch, obs_step, number of different cameras, projection dims)\n",
    "        input_tokens = [rgb_tokens[:, :, i] for i in range(rgb_tokens.size(2))]\n",
    "        input_tokens.append(\n",
    "            self.state_projector(observation_state)\n",
    "        )  # (batch, obs_step, projection dims)\n",
    "        input_tokens.append(einops.repeat(self.action_token, \"1 1 d -> b n d\", b=batch_size, n=n_obs_steps))\n",
    "        # Interleave tokens by stacking and rearranging.\n",
    "        input_tokens = Tensor.stack(*input_tokens, dim=2)\n",
    "        input_tokens = einops.rearrange(input_tokens, \"b n t d -> b (n t) d\")\n",
    "\n",
    "        len_additional_action_token = self.config.n_action_pred_token - 1\n",
    "        future_action_tokens = self.action_token.repeat(batch_size, len_additional_action_token, 1)\n",
    "\n",
    "        # add additional action query tokens for predicting future action chunks\n",
    "        input_tokens = input_tokens.cat(future_action_tokens, dim=1)\n",
    "\n",
    "        # get action features (pass through GPT)\n",
    "        features = self.policy(input_tokens)\n",
    "        # len(self.config.input_shapes) is the number of different observation modes.\n",
    "        # this line gets the index of action prompt tokens.\n",
    "        historical_act_pred_index = Tensor.arange(0, n_obs_steps, requires_grad=False) * (len(self.config.input_shapes) + 1) + len(\n",
    "            self.config.input_shapes\n",
    "        )\n",
    "\n",
    "        # only extract the output tokens at the position of action query:\n",
    "        # Behavior Transformer (BeT), and VQ-BeT are both sequence-to-sequence prediction models,\n",
    "        # mapping sequential observation to sequential action (please refer to section 2.2 in BeT paper https://arxiv.org/pdf/2206.11251).\n",
    "        # Thus, it predicts a historical action sequence, in addition to current and future actions (predicting future actions : optional).\n",
    "        if len_additional_action_token > 0:\n",
    "            features = Tensor.cat(\n",
    "                *[features[:, historical_act_pred_index], features[:, -len_additional_action_token:]], dim=1\n",
    "            )\n",
    "        else:\n",
    "            features = features[:, historical_act_pred_index]\n",
    "        # pass through action head\n",
    "        return self.action_head(features)\n",
    "\n",
    "\n",
    "    def __call__(self, batch: dict[str, Tensor], rollout: bool) -> Tensor:\n",
    "        # Input validation.\n",
    "        assert set(batch).issuperset({\"observation.state\", \"observation.images\"})\n",
    "        action_head_output = self.call_first_step(batch[\"observation.images\"], batch[\"observation.state\"])\n",
    "        # if rollout, VQ-BeT don't calculate loss\n",
    "        if rollout:\n",
    "            return action_head_output[\"predicted_action\"][:, n_obs_steps - 1, :].reshape(\n",
    "                batch_size, self.config.action_chunk_size, -1\n",
    "            )\n",
    "        # else, it calculate overall loss (bin prediction loss, and offset loss)\n",
    "        else:\n",
    "            output = batch[\"action\"][:, self.select_target_actions_indices]\n",
    "            loss = self.action_head.loss_fn(action_head_output, output, reduction=\"mean\")\n",
    "            return action_head_output, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "610d655e-29fc-43aa-a084-33c498e3a090",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpatialSoftmax:\n",
    "    \"\"\"\n",
    "    Spatial Soft Argmax operation described in \"Deep Spatial Autoencoders for Visuomotor Learning\" by Finn et al.\n",
    "    (https://arxiv.org/pdf/1509.06113). A minimal port of the robomimic implementation.\n",
    "\n",
    "    At a high level, this takes 2D feature maps (from a convnet/ViT) and returns the \"center of mass\"\n",
    "    of activations of each channel, i.e., keypoints in the image space for the policy to focus on.\n",
    "\n",
    "    Example: take feature maps of size (512x10x12). We generate a grid of normalized coordinates (10x12x2):\n",
    "    -----------------------------------------------------\n",
    "    | (-1., -1.)   | (-0.82, -1.)   | ... | (1., -1.)   |\n",
    "    | (-1., -0.78) | (-0.82, -0.78) | ... | (1., -0.78) |\n",
    "    | ...          | ...            | ... | ...         |\n",
    "    | (-1., 1.)    | (-0.82, 1.)    | ... | (1., 1.)    |\n",
    "    -----------------------------------------------------\n",
    "    This is achieved by applying channel-wise softmax over the activations (512x120) and computing the dot\n",
    "    product with the coordinates (120x2) to get expected points of maximal activation (512x2).\n",
    "\n",
    "    The example above results in 512 keypoints (corresponding to the 512 input channels). We can optionally\n",
    "    provide num_kp != None to control the number of keypoints. This is achieved by a first applying a learnable\n",
    "    linear mapping (in_channels, H, W) -> (num_kp, H, W).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_shape, num_kp=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            input_shape (list): (C, H, W) input feature map shape.\n",
    "            num_kp (int): number of keypoints in output. If None, output will have the same number of channels as input.\n",
    "        \"\"\"\n",
    "\n",
    "        assert len(input_shape) == 3\n",
    "        self._in_c, self._in_h, self._in_w = input_shape\n",
    "\n",
    "        if num_kp is not None:\n",
    "            self.nets = nn.Conv2d(self._in_c, num_kp, kernel_size=1)\n",
    "            self._out_c = num_kp\n",
    "        else:\n",
    "            self.nets = None\n",
    "            self._out_c = self._in_c\n",
    "\n",
    "        pos_x, pos_y = np.meshgrid(\n",
    "            np.linspace(-1.0, 1.0, self._in_w), \n",
    "            np.linspace(-1.0, 1.0, self._in_h)\n",
    "        )\n",
    "        pos_x = Tensor(pos_x.reshape(self._in_h * self._in_w, 1), dtype=dtypes.float, requires_grad=False)\n",
    "        pos_y = Tensor(pos_y.reshape(self._in_h * self._in_w, 1), dtype=dtypes.float, requires_grad=False)\n",
    "        # register as buffer so it's moved to the correct device.\n",
    "        self.pos_grid = Tensor.cat(pos_x, pos_y, dim=1).detach()\n",
    "        self.pos_grid.requires_grad = False\n",
    "\n",
    "    def __call__(self, features: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            features: (B, C, H, W) input feature maps.\n",
    "        Returns:\n",
    "            (B, K, 2) image-space coordinates of keypoints.\n",
    "        \"\"\"\n",
    "        if self.nets is not None:\n",
    "            features = self.nets(features)\n",
    "\n",
    "        # [B, K, H, W] -> [B * K, H * W] where K is number of keypoints\n",
    "        features = features.reshape(-1, self._in_h * self._in_w)\n",
    "        # 2d softmax normalization\n",
    "        attention = features.softmax(axis=-1)\n",
    "        # [B * K, H * W] x [H * W, 2] -> [B * K, 2] for spatial coordinate mean in x and y dimensions\n",
    "        expected_xy = attention @ self.pos_grid\n",
    "        # reshape to [B, K, 2]\n",
    "        feature_keypoints = expected_xy.view(-1, self._out_c, 2)\n",
    "\n",
    "        return feature_keypoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "eb0ceeba-2529-47f6-ac88-057bf0b1dc8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tinygrad import Tensor, nn, dtypes\n",
    "\n",
    "def create_stats_buffers(\n",
    "    shapes: dict[str, list[int]],\n",
    "    modes: dict[str, str],\n",
    "    stats: dict[str, dict[str, Tensor]] | None = None,\n",
    ") -> dict[str, dict[str, dict]]:\n",
    "    \"\"\"\n",
    "    Create buffers per modality (e.g. \"observation.image\", \"action\") containing their mean, std, min, max\n",
    "    statistics.\n",
    "\n",
    "    Args: (see Normalize and Unnormalize)\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary where keys are modalities and values are `nn.ParameterDict` containing\n",
    "            `nn.Parameters` set to `requires_grad=False`, suitable to not be updated during backpropagation.\n",
    "    \"\"\"\n",
    "    stats_buffers = {}\n",
    "\n",
    "    for key, mode in modes.items():\n",
    "        assert mode in [\"mean_std\", \"min_max\"]\n",
    "\n",
    "        shape = tuple(shapes[key])\n",
    "\n",
    "        if \"image\" in key:\n",
    "            # sanity checks\n",
    "            assert len(shape) == 3, f\"number of dimensions of {key} != 3 ({shape=}\"\n",
    "            c, h, w = shape\n",
    "            assert c < h and c < w, f\"{key} is not channel first ({shape=})\"\n",
    "            # override image shape to be invariant to height and width\n",
    "            shape = (c, 1, 1)\n",
    "\n",
    "        # Note: we initialize mean, std, min, max to infinity. They should be overwritten\n",
    "        # downstream by `stats` or `policy.load_state_dict`, as expected. During forward,\n",
    "        # we assert they are not infinity anymore.\n",
    "\n",
    "        buffer = {}\n",
    "        if mode == \"mean_std\":\n",
    "            buffer = {\n",
    "                \"mean\": Tensor.ones(shape, dtype=dtypes.float, requires_grad=False) * float('inf'),\n",
    "                \"std\": Tensor.ones(shape, dtype=dtypes.float, requires_grad=False) * float('inf'),\n",
    "            }\n",
    "        elif mode == \"min_max\":\n",
    "            buffer = {\n",
    "                \"min\": Tensor.ones(shape, dtype=dtypes.float, requires_grad=False) * float('inf'),\n",
    "                \"max\": Tensor.ones(shape, dtype=dtypes.float, requires_grad=False) * float('inf'),\n",
    "            }\n",
    "\n",
    "        if stats is not None:\n",
    "            # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated\n",
    "            # tensors anywhere (for example, when we use the same stats for normalization and\n",
    "            # unnormalization). See the logic here\n",
    "            # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.\n",
    "            if mode == \"mean_std\":\n",
    "                buffer[\"mean\"].assign(stats[key][\"mean\"].numpy())\n",
    "                buffer[\"mean\"].requires_grad = False\n",
    "                buffer[\"std\"].assign(stats[key][\"std\"].numpy())\n",
    "                buffer[\"std\"].requires_grad = False\n",
    "            elif mode == \"min_max\":\n",
    "                buffer[\"min\"].assign(stats[key][\"min\"].numpy())\n",
    "                buffer[\"min\"].requires_grad = False\n",
    "                buffer[\"max\"].assign(stats[key][\"max\"].numpy())\n",
    "                buffer[\"max\"].requires_grad = False\n",
    "\n",
    "        stats_buffers[key] = buffer\n",
    "    return stats_buffers\n",
    "\n",
    "def _no_stats_error_str(name: str) -> str:\n",
    "    return (\n",
    "        f\"`{name}` is infinity. You should either initialize with `stats` as an argument, or use a \"\n",
    "        \"pretrained model.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class Normalize():\n",
    "    \"\"\"Normalizes data (e.g. \"observation.image\") for more stable and faster convergence during training.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        shapes: dict[str, list[int]],\n",
    "        modes: dict[str, str],\n",
    "        stats: dict[str, dict[str, Tensor]] | None = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            shapes (dict): A dictionary where keys are input modalities (e.g. \"observation.image\") and values\n",
    "            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing\n",
    "            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape\n",
    "            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.\n",
    "            modes (dict): A dictionary where keys are output modalities (e.g. \"observation.image\") and values\n",
    "                are their normalization modes among:\n",
    "                    - \"mean_std\": subtract the mean and divide by standard deviation.\n",
    "                    - \"min_max\": map to [-1, 1] range.\n",
    "            stats (dict, optional): A dictionary where keys are output modalities (e.g. \"observation.image\")\n",
    "                and values are dictionaries of statistic types and their values (e.g.\n",
    "                `{\"mean\": torch.randn(3,1,1)}, \"std\": torch.randn(3,1,1)}`). If provided, as expected for\n",
    "                training the model for the first time, these statistics will overwrite the default buffers. If\n",
    "                not provided, as expected for finetuning or evaluation, the default buffers should to be\n",
    "                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the\n",
    "                dataset is not needed to get the stats, since they are already in the policy state_dict.\n",
    "        \"\"\"\n",
    "        self.shapes = shapes\n",
    "        self.modes = modes\n",
    "        self.stats = stats\n",
    "        stats_buffers = create_stats_buffers(shapes, modes, stats)\n",
    "        for key, buffer in stats_buffers.items():\n",
    "            setattr(self, \"buffer_\" + key.replace(\".\", \"_\"), buffer)\n",
    "\n",
    "    # TODO(rcadene): should we remove Tensor.no_grad?\n",
    "    # @Tensor.no_grad\n",
    "    def __call__(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        batch = dict(batch)  # shallow copy avoids mutating the input batch\n",
    "        for key, mode in self.modes.items():\n",
    "            buffer = getattr(self, \"buffer_\" + key.replace(\".\", \"_\"))\n",
    "\n",
    "            if mode == \"mean_std\":\n",
    "                mean = buffer[\"mean\"]\n",
    "                std = buffer[\"std\"]\n",
    "                assert not (mean == float('inf')).any().numpy(), _no_stats_error_str(\"mean\")\n",
    "                assert not (std == float('inf')).any().numpy(), _no_stats_error_str(\"std\")\n",
    "                #print(f'mean: {mean.numpy()}, std: {std.numpy()}')\n",
    "                #print(f'batch[{key}] before normalization: {batch[key].numpy()}')\n",
    "                batch[key] = (batch[key] - mean) / (std + 1e-8)\n",
    "                #print(f'batch[{key}] after normalization: {batch[key].numpy()}')\n",
    "            elif mode == \"min_max\":\n",
    "                min = buffer[\"min\"]\n",
    "                max = buffer[\"max\"]\n",
    "                assert not (min == float('inf')).any().numpy(), _no_stats_error_str(\"min\")\n",
    "                assert not (max == float('inf')).any().numpy(), _no_stats_error_str(\"max\")\n",
    "                # normalize to [0,1]\n",
    "                #print(f'max: {max.numpy()}, min: {min.numpy()}')\n",
    "                #print(f'batch[{key}] before normalization: {batch[key].numpy()}')\n",
    "                batch[key] = (batch[key] - min) / (max - min + 1e-8)\n",
    "                # normalize to [-1, 1]\n",
    "                batch[key] = batch[key] * 2 - 1\n",
    "                #print(f'batch[{key}] after normalization: {batch[key].numpy()}')\n",
    "            else:\n",
    "                raise ValueError(mode)\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        return batch\n",
    "\n",
    "\n",
    "class Unnormalize():\n",
    "    \"\"\"\n",
    "    Similar to `Normalize` but unnormalizes output data (e.g. `{\"action\": torch.randn(b,c)}`) in their\n",
    "    original range used by the environment.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        shapes: dict[str, list[int]],\n",
    "        modes: dict[str, str],\n",
    "        stats: dict[str, dict[str, Tensor]] | None = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            shapes (dict): A dictionary where keys are input modalities (e.g. \"observation.image\") and values\n",
    "            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing\n",
    "            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape\n",
    "            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.\n",
    "            modes (dict): A dictionary where keys are output modalities (e.g. \"observation.image\") and values\n",
    "                are their normalization modes among:\n",
    "                    - \"mean_std\": subtract the mean and divide by standard deviation.\n",
    "                    - \"min_max\": map to [-1, 1] range.\n",
    "            stats (dict, optional): A dictionary where keys are output modalities (e.g. \"observation.image\")\n",
    "                and values are dictionaries of statistic types and their values (e.g.\n",
    "                `{\"mean\": torch.randn(3,1,1)}, \"std\": torch.randn(3,1,1)}`). If provided, as expected for\n",
    "                training the model for the first time, these statistics will overwrite the default buffers. If\n",
    "                not provided, as expected for finetuning or evaluation, the default buffers should to be\n",
    "                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the\n",
    "                dataset is not needed to get the stats, since they are already in the policy state_dict.\n",
    "        \"\"\"\n",
    "        self.shapes = shapes\n",
    "        self.modes = modes\n",
    "        self.stats = stats\n",
    "        # `self.buffer_observation_state[\"mean\"]` contains `torch.tensor(state_dim)`\n",
    "        stats_buffers = create_stats_buffers(shapes, modes, stats)\n",
    "        for key, buffer in stats_buffers.items():\n",
    "            setattr(self, \"buffer_\" + key.replace(\".\", \"_\"), buffer)\n",
    "\n",
    "    # TODO(rcadene): should we remove torch.no_grad?\n",
    "    def __call__(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        batch = dict(batch)  # shallow copy avoids mutating the input batch\n",
    "        for key, mode in self.modes.items():\n",
    "            buffer = getattr(self, \"buffer_\" + key.replace(\".\", \"_\"))\n",
    "\n",
    "            if mode == \"mean_std\":\n",
    "                mean = buffer[\"mean\"]\n",
    "                std = buffer[\"std\"]\n",
    "                assert not (mean == float('inf')).any().numpy(), _no_stats_error_str(\"mean\")\n",
    "                assert not (std == float('inf')).any().numpy(), _no_stats_error_str(\"std\")\n",
    "                batch[key] = batch[key] * std + mean\n",
    "            elif mode == \"min_max\":\n",
    "                min = buffer[\"min\"]\n",
    "                max = buffer[\"max\"]\n",
    "                assert not (min == float('inf')).any().numpy(), _no_stats_error_str(\"min\")\n",
    "                assert not (max == float('inf')).any().numpy(), _no_stats_error_str(\"max\")\n",
    "                batch[key] = (batch[key] + 1) / 2\n",
    "                batch[key] = batch[key] * (max - min) + min\n",
    "            else:\n",
    "                raise ValueError(mode)\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        return batch\n",
    "\n",
    "def populate_queues(queues, batch):\n",
    "    for key in batch:\n",
    "        # Ignore keys not in the queues already (leaving the responsibility to the caller to make sure the\n",
    "        # queues have the keys they want).\n",
    "        if key not in queues:\n",
    "            continue\n",
    "        if len(queues[key]) != queues[key].maxlen:\n",
    "            # initialize by copying the first observation several times until the queue is full\n",
    "            while len(queues[key]) != queues[key].maxlen:\n",
    "                queues[key].append(batch[key])\n",
    "        else:\n",
    "            # add latest observation to the queue\n",
    "            queues[key].append(batch[key])\n",
    "    return queues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "20ed9e34-fe34-424d-95f5-057d6b94cfc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VQBeTPolicy:\n",
    "    \"\"\"\n",
    "    VQ-BeT Policy as per \"Behavior Generation with Latent Actions\"\n",
    "    \"\"\"\n",
    "\n",
    "    name = \"vqbet\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        config: VQBeTConfig | None = None,\n",
    "        dataset_stats: dict[str, dict[str, Tensor]] | None = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            config: Policy configuration class instance or None, in which case the default instantiation of\n",
    "                the configuration class is used.\n",
    "            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected\n",
    "                that they will be passed with a call to `load_state_dict` before the policy is used.\n",
    "        \"\"\"\n",
    "        if config is None:\n",
    "            config = VQBeTConfig()\n",
    "        self.config = config\n",
    "        self.normalize_inputs = Normalize(\n",
    "            config.input_shapes, config.input_normalization_modes, dataset_stats\n",
    "        )\n",
    "        self.normalize_targets = Normalize(\n",
    "            config.output_shapes, config.output_normalization_modes, dataset_stats\n",
    "        )\n",
    "        self.unnormalize_outputs = Unnormalize(\n",
    "            config.output_shapes, config.output_normalization_modes, dataset_stats\n",
    "        )\n",
    "\n",
    "        self.vqbet = VQBeTModel(config)\n",
    "\n",
    "        self.expected_image_keys = [k for k in config.input_shapes if k.startswith(\"observation.image\")]\n",
    "\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Clear observation and action queues. Should be called on `env.reset()`\n",
    "        queues are populated during rollout of the policy, they contain the n latest observations and actions\n",
    "        \"\"\"\n",
    "        self._queues = {\n",
    "            \"observation.images\": deque(maxlen=self.config.n_obs_steps),\n",
    "            \"observation.state\": deque(maxlen=self.config.n_obs_steps),\n",
    "            \"action\": deque(maxlen=self.config.action_chunk_size),\n",
    "        }\n",
    "\n",
    "    def select_action(self, batch: dict[str, Tensor]) -> Tensor:\n",
    "        Tensor.no_grad = True\n",
    "        training_prev, Tensor.training = Tensor.training, False\n",
    "        \"\"\"Select a single action given environment observations.\n",
    "\n",
    "        This method wraps `select_actions` in order to return one action at a time for execution in the\n",
    "        environment. It works by managing the actions in a queue and only calling `select_actions` when the\n",
    "        queue is empty.\n",
    "        \"\"\"\n",
    "\n",
    "        batch = self.normalize_inputs(batch)\n",
    "        batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original\n",
    "        batch[\"observation.images\"] = Tensor.stack(*[batch[k] for k in self.expected_image_keys], dim=-4)\n",
    "        # Note: It's important that this happens after stacking the images into a single key.\n",
    "        self._queues = populate_queues(self._queues, batch)\n",
    "\n",
    "        if not self.vqbet.action_head.vqvae_model.discretized:\n",
    "            warnings.warn(\n",
    "                \"To evaluate in the environment, your VQ-BeT model should contain a pretrained Residual VQ.\",\n",
    "                stacklevel=1,\n",
    "            )\n",
    "\n",
    "        if len(self._queues[\"action\"]) == 0:\n",
    "            batch = {k: Tensor.stack(self._queues[k], dim=1) for k in batch if k in self._queues}\n",
    "            actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]\n",
    "\n",
    "            # the dimension of returned action is (batch_size, action_chunk_size, action_dim)\n",
    "            actions = self.unnormalize_outputs({\"action\": actions})[\"action\"]\n",
    "            # since the data in the action queue's dimension is (action_chunk_size, batch_size, action_dim), we transpose the action and fill the queue\n",
    "            self._queues[\"action\"].extend(actions.transpose(0, 1))\n",
    "\n",
    "        action = self._queues[\"action\"].popleft()\n",
    "        Tensor.no_grad = False\n",
    "        Tensor.training = training_prev\n",
    "        return action\n",
    "\n",
    "    def __call__(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:\n",
    "        \"\"\"Run the batch through the model and compute the loss for training or validation.\"\"\"\n",
    "        batch = self.normalize_inputs(batch)\n",
    "        batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original\n",
    "        batch[\"observation.images\"] = Tensor.stack(*[batch[k] for k in self.expected_image_keys], dim=-4)\n",
    "        batch = self.normalize_targets(batch)\n",
    "        # VQ-BeT discretizes action using VQ-VAE before training BeT (please refer to section 3.2 in the VQ-BeT paper https://arxiv.org/pdf/2403.03181)\n",
    "        if not self.vqbet.action_head.vqvae_model.discretized:\n",
    "            # loss: total loss of training RVQ\n",
    "            # n_different_codes: how many of the total possible VQ codes are being used in single batch (how many of them have at least one encoder embedding as a nearest neighbor). This can be at most `vqvae_n_embed * number of layers of RVQ (=2)`.\n",
    "            # n_different_combinations: how many different code combinations are being used out of all possible combinations in single batch. This can be at most `vqvae_n_embed ^ number of layers of RVQ (=2)` (hint consider the RVQ as a decision tree).\n",
    "            loss, n_different_codes, n_different_combinations, recon_l1_error = (\n",
    "                self.vqbet.action_head.discretize(self.config.n_vqvae_training_steps, batch[\"action\"])\n",
    "            )\n",
    "            return {\n",
    "                \"loss\": loss,\n",
    "                \"n_different_codes\": n_different_codes,\n",
    "                \"n_different_combinations\": n_different_combinations,\n",
    "                \"recon_l1_error\": recon_l1_error,\n",
    "            }\n",
    "        # if Residual VQ is already trained, VQ-BeT trains its GPT and bin prediction head / offset prediction head parts.\n",
    "        _, loss_dict = self.vqbet(batch, rollout=False)\n",
    "\n",
    "        return loss_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b641fefd-33a5-4443-b17a-d481a5d68a57",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d1b7c611c6d54689a967787090a3a6f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 212 files:   0%|          | 0/212 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "beam2  : \u001b[34m25\u001b[0m \u001b[34m100\u001b[0m \u001b[33m5\u001b[0m \u001b[33m5\u001b[0m \u001b[33m4\u001b[0m                   :    25.12 us  <  hc     : \u001b[34m125\u001b[0m \u001b[34m125\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                    :    70.71 us\n",
      "\u001b[32m*** METAL      7\u001b[0m E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.00 GB tm     56.37us/     0.16ms (     5.41 GFLOPS   17.7|17.7    GB/s) ['tril']\n",
      "*** METAL      8 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.00 GB tm     55.54us/     0.21ms (     5.49 GFLOPS   18.0|18.0    GB/s) ['tril']\n",
      "*** METAL      9 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.00 GB tm     61.12us/     0.28ms (     4.99 GFLOPS   16.4|16.4    GB/s) ['tril']\n",
      "*** METAL     10 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.00 GB tm     58.00us/     0.33ms (     5.26 GFLOPS   17.2|17.2    GB/s) ['tril']\n",
      "*** METAL     11 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.01 GB tm     56.21us/     0.39ms (     5.43 GFLOPS   17.8|17.8    GB/s) ['tril']\n",
      "*** METAL     12 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.01 GB tm     54.92us/     0.44ms (     5.55 GFLOPS   18.2|18.2    GB/s) ['tril']\n",
      "*** METAL     13 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.01 GB tm    540.25us/     0.98ms (     0.56 GFLOPS    1.9|1.9     GB/s) ['tril']\n",
      "*** METAL     14 E_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[34m100\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  0.01 GB tm     55.00us/     1.04ms (     5.55 GFLOPS   18.2|18.2    GB/s) ['tril']\n",
      "replacing GPT parameter: h.0.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.1.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.2.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.3.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.4.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.5.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.6.attn.c_proj.weight with a normal\n",
      "replacing GPT parameter: h.7.attn.c_proj.weight with a normal\n",
      "decay list: ['h.0.attn.c_attn.weight', 'h.0.attn.c_proj.weight', 'h.0.mlp.0.weight', 'h.0.mlp.2.weight', 'h.1.attn.c_attn.weight', 'h.1.attn.c_proj.weight', 'h.1.mlp.0.weight', 'h.1.mlp.2.weight', 'h.2.attn.c_attn.weight', 'h.2.attn.c_proj.weight', 'h.2.mlp.0.weight', 'h.2.mlp.2.weight', 'h.3.attn.c_attn.weight', 'h.3.attn.c_proj.weight', 'h.3.mlp.0.weight', 'h.3.mlp.2.weight', 'h.4.attn.c_attn.weight', 'h.4.attn.c_proj.weight', 'h.4.mlp.0.weight', 'h.4.mlp.2.weight', 'h.5.attn.c_attn.weight', 'h.5.attn.c_proj.weight', 'h.5.mlp.0.weight', 'h.5.mlp.2.weight', 'h.6.attn.c_attn.weight', 'h.6.attn.c_proj.weight', 'h.6.mlp.0.weight', 'h.6.mlp.2.weight', 'h.7.attn.c_attn.weight', 'h.7.attn.c_proj.weight', 'h.7.mlp.0.weight', 'h.7.mlp.2.weight', 'lm_head.weight', 'wte.weight']\n",
      "no_decay list: ['h.0.attn.bias', 'h.0.attn.c_attn.bias', 'h.0.attn.c_proj.bias', 'h.0.ln_1.bias', 'h.0.ln_1.weight', 'h.0.ln_2.bias', 'h.0.ln_2.weight', 'h.0.mlp.0.bias', 'h.0.mlp.2.bias', 'h.1.attn.bias', 'h.1.attn.c_attn.bias', 'h.1.attn.c_proj.bias', 'h.1.ln_1.bias', 'h.1.ln_1.weight', 'h.1.ln_2.bias', 'h.1.ln_2.weight', 'h.1.mlp.0.bias', 'h.1.mlp.2.bias', 'h.2.attn.bias', 'h.2.attn.c_attn.bias', 'h.2.attn.c_proj.bias', 'h.2.ln_1.bias', 'h.2.ln_1.weight', 'h.2.ln_2.bias', 'h.2.ln_2.weight', 'h.2.mlp.0.bias', 'h.2.mlp.2.bias', 'h.3.attn.bias', 'h.3.attn.c_attn.bias', 'h.3.attn.c_proj.bias', 'h.3.ln_1.bias', 'h.3.ln_1.weight', 'h.3.ln_2.bias', 'h.3.ln_2.weight', 'h.3.mlp.0.bias', 'h.3.mlp.2.bias', 'h.4.attn.bias', 'h.4.attn.c_attn.bias', 'h.4.attn.c_proj.bias', 'h.4.ln_1.bias', 'h.4.ln_1.weight', 'h.4.ln_2.bias', 'h.4.ln_2.weight', 'h.4.mlp.0.bias', 'h.4.mlp.2.bias', 'h.5.attn.bias', 'h.5.attn.c_attn.bias', 'h.5.attn.c_proj.bias', 'h.5.ln_1.bias', 'h.5.ln_1.weight', 'h.5.ln_2.bias', 'h.5.ln_2.weight', 'h.5.mlp.0.bias', 'h.5.mlp.2.bias', 'h.6.attn.bias', 'h.6.attn.c_attn.bias', 'h.6.attn.c_proj.bias', 'h.6.ln_1.bias', 'h.6.ln_1.weight', 'h.6.ln_2.bias', 'h.6.ln_2.weight', 'h.6.mlp.0.bias', 'h.6.mlp.2.bias', 'h.7.attn.bias', 'h.7.attn.c_attn.bias', 'h.7.attn.c_proj.bias', 'h.7.ln_1.bias', 'h.7.ln_1.weight', 'h.7.ln_2.bias', 'h.7.ln_2.weight', 'h.7.mlp.0.bias', 'h.7.mlp.2.bias', 'ln_f.bias', 'ln_f.weight', 'wpe.weight', 'wte.bias']\n",
      "\u001b[32m*** METAL     15\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.01 GB tm     31.17us/     1.07ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     16\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.01 GB tm     10.00us/     1.08ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     17\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.01 GB tm     25.50us/     1.11ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     18\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.01 GB tm     10.21us/     1.12ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     19\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.01 GB tm     10.75us/     1.13ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     20\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.01 GB tm      8.58us/     1.14ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "Starting training loop\n",
      "\u001b[32m*** METAL     21\u001b[0m \u001b[33mcopy   35.39M,   METAL <- NPY    \u001b[0m        mem  0.08 GB tm   5128.96us/     6.26ms (     0.00 GFLOPS    6.9|6.9     GB/s) \n",
      "\u001b[32m*** METAL     22\u001b[0m \u001b[33mcopy     2560,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     25.71us/     6.29ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL     23\u001b[0m \u001b[33mcopy     7680,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     13.17us/     6.30ms (     0.00 GFLOPS    0.6|0.6     GB/s) \n",
      "\u001b[32m*** METAL     24\u001b[0m \u001b[33mcopy      512,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     11.33us/     6.32ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     25\u001b[0m \u001b[33mcopy      512,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      9.33us/     6.32ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL     26\u001b[0m \u001b[33mcopy      256,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      9.29us/     6.33ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     27\u001b[0m \u001b[33mcopy      256,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      8.12us/     6.34ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     28\u001b[0m \u001b[33mcopy       64,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      7.79us/     6.35ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     29\u001b[0m \u001b[33mcopy       64,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      7.58us/     6.36ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     30\u001b[0m \u001b[33mcopy      512,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      8.58us/     6.37ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL     31\u001b[0m \u001b[33mcopy      320,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      9.29us/     6.38ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     32\u001b[0m \u001b[33mcopy      320,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      7.67us/     6.38ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     33\u001b[0m \u001b[33mcopy      960,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      7.79us/     6.39ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL     34\u001b[0m \u001b[33mcopy       12,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     10.62us/     6.40ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[35m3\u001b[0m                              :     6.71 us  <  hc     : \u001b[35m3\u001b[0m                              :     6.71 us\n",
      "\u001b[32m*** METAL     35\u001b[0m r_\u001b[35m3\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      7.50us/     6.41ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG     36\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.04 GB tm     41.38us/     6.45ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     37\u001b[0m \u001b[33mcopy       12,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     14.67us/     6.46ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL     38 r_\u001b[35m3\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      6.58us/     6.47ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG     39\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.04 GB tm     12.13us/     6.48ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     40\u001b[0m \u001b[33mcopy        8,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     12.08us/     6.50ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[35m2\u001b[0m                              :     6.62 us  <  hc     : \u001b[35m2\u001b[0m                              :     6.62 us\n",
      "\u001b[32m*** METAL     41\u001b[0m r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      7.58us/     6.50ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG     42\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.04 GB tm     14.42us/     6.52ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     43\u001b[0m \u001b[33mcopy        8,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     12.29us/     6.53ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL     44 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      6.58us/     6.54ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG     45\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.04 GB tm     13.29us/     6.55ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     46\u001b[0m \u001b[33mcopy        8,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     10.29us/     6.56ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL     47 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      6.25us/     6.57ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG     48\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.04 GB tm     10.46us/     6.58ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     49\u001b[0m \u001b[33mcopy        8,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      9.96us/     6.59ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL     50 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      6.54us/     6.59ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG     51\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.04 GB tm     16.17us/     6.61ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "scheduled 64 kernels\n",
      "beam2  : \u001b[36m2\u001b[0m                              :     6.58 us  <  hc     : \u001b[36m2\u001b[0m                              :     6.58 us\n",
      "\u001b[32m*** METAL     52\u001b[0m E_\u001b[36m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.04 GB tm      7.50us/     6.62ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__sub__', '__add__']\n",
      "\u001b[32m*** CUSTOM    53\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    54\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    55\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    56\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    57\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    58\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    59\u001b[0m custom_random                            mem  0.04 GB \n",
      "beam2  : \u001b[31m8\u001b[0m \u001b[33m16\u001b[0m \u001b[35m4\u001b[0m                         :     6.71 us  <  hc     : \u001b[34m16\u001b[0m \u001b[32m16\u001b[0m                          :     7.50 us\n",
      "\u001b[32m*** METAL     60\u001b[0m r_\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.04 GB tm      6.88us/     6.62ms (     0.20 GFLOPS    0.0|0.0     GB/s) ['argmax']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[31m16\u001b[0m                         :     6.13 us  <  hc     : \u001b[34m16\u001b[0m \u001b[32m16\u001b[0m                          :     6.88 us\n",
      "\u001b[32m*** METAL     61\u001b[0m r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                                 mem  0.04 GB tm      6.67us/     6.63ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['one_hot']\n",
      "\u001b[32m*** CUSTOM    62\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    63\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    64\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    65\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    66\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    67\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** CUSTOM    68\u001b[0m custom_random                            mem  0.04 GB \n",
      "\u001b[32m*** METAL     69\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm     10.96us/     6.64ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL     70\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.04 GB tm      8.17us/     6.65ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[34m30\u001b[0m \u001b[34m2\u001b[0m \u001b[36m32\u001b[0m                        :     6.92 us  <  hc     : \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m3\u001b[0m                      :     7.17 us\n",
      "\u001b[32m*** METAL     71\u001b[0m E_\u001b[34m30\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m\u001b[0m                                mem  0.04 GB tm      8.00us/     6.66ms (     1.44 GFLOPS    1.9|3.8     GB/s) ['__sub__', '__truediv__', '__mul__']\n",
      "hc     : \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.17 us  <  beam2  : \u001b[34m16\u001b[0m \u001b[36m16\u001b[0m \u001b[33m5\u001b[0m                        :     7.21 us\n",
      "\u001b[32m*** METAL     72\u001b[0m E_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.04 GB tm      8.13us/     6.67ms (     0.32 GFLOPS    1.3|1.3     GB/s) ['uniform']\n",
      "hc     : \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                           :     7.04 us  <  beam2  : \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m                       :     7.12 us\n",
      "\u001b[32m*** METAL     73\u001b[0m E_\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                   mem  0.04 GB tm      7.75us/     6.67ms (     0.03 GFLOPS    0.1|0.1     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     8.29 us  <  hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     8.29 us\n",
      "\u001b[32m*** METAL     74\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                               mem  0.04 GB tm      8.79us/     6.68ms (     3.73 GFLOPS   14.9|14.9    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m8\u001b[0m \u001b[36m16\u001b[0m                           :     6.88 us  <  hc     : \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                           :     7.04 us\n",
      "\u001b[32m*** METAL     75\u001b[0m E_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m\u001b[0m                                   mem  0.04 GB tm      7.83us/     6.69ms (     0.03 GFLOPS    0.1|0.1     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m256\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                      :     8.75 us  <  hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     8.92 us\n",
      "\u001b[32m*** METAL     76\u001b[0m E_\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.04 GB tm     13.50us/     6.70ms (     4.85 GFLOPS   19.4|19.4    GB/s) ['uniform']\n",
      "hc     : \u001b[34m2\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     6.50 us  <  beam2  : \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m                          :     7.12 us\n",
      "\u001b[32m*** METAL     77\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.04 GB tm      7.67us/     6.71ms (     0.07 GFLOPS    0.3|0.3     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.38 us  <  hc     : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.38 us\n",
      "\u001b[32m*** METAL     78\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.04 GB tm      8.25us/     6.72ms (     0.99 GFLOPS    4.0|4.0     GB/s) ['kaiming_uniform']\n",
      "*** METAL     79 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.04 GB tm      8.46us/     6.73ms (     0.97 GFLOPS    3.9|3.9     GB/s) ['kaiming_uniform']\n",
      "hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     8.25 us  <  beam2  : \u001b[34m64\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                    :     8.96 us\n",
      "\u001b[32m*** METAL     80\u001b[0m E_\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                               mem  0.04 GB tm      9.50us/     6.74ms (     6.90 GFLOPS   27.6|27.6    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[33m2\u001b[0m                          :     6.75 us  <  hc     : \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                           :     7.00 us\n",
      "\u001b[32m*** METAL     81\u001b[0m E_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                  mem  0.04 GB tm      7.75us/     6.75ms (     0.03 GFLOPS    0.1|0.1     GB/s) ['uniform']\n",
      "*** METAL     82 E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                               mem  0.04 GB tm      8.67us/     6.75ms (     3.78 GFLOPS   15.1|15.1    GB/s) ['uniform']\n",
      "*** METAL     83 E_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m\u001b[0m                                   mem  0.04 GB tm      7.62us/     6.76ms (     0.03 GFLOPS    0.1|0.1     GB/s) ['uniform']\n",
      "hc     : \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.13 us  <  beam2  : \u001b[34m64\u001b[0m \u001b[36m4\u001b[0m \u001b[33m5\u001b[0m                         :     7.38 us\n",
      "\u001b[32m*** METAL     84\u001b[0m E_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                              mem  0.04 GB tm      7.92us/     6.77ms (     0.32 GFLOPS    1.3|1.3     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[33m5\u001b[0m                            :     6.75 us  <  hc     : \u001b[34m5\u001b[0m \u001b[36m2\u001b[0m                            :     6.92 us\n",
      "\u001b[32m*** METAL     85\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                                    mem  0.04 GB tm      7.62us/     6.78ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['uniform']\n",
      "hc     : \u001b[34m22\u001b[0m \u001b[34m5\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m                      :     9.00 us  <  beam2  : \u001b[34m22\u001b[0m \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m                     :    10.46 us\n",
      "\u001b[32m*** METAL     86\u001b[0m E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m\u001b[0m                              mem  0.04 GB tm      9.75us/     6.79ms (    28.88 GFLOPS   11.5|34.7    GB/s) ['cat']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m                       :     8.88 us  <  hc     : \u001b[36m4\u001b[0m \u001b[32m8\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m                       :    14.54 us\n",
      "\u001b[32m*** METAL     87\u001b[0m r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                               mem  0.04 GB tm      9.54us/     6.80ms (     3.00 GFLOPS    1.7|5.4     GB/s) ['square', 'sum']\n",
      "*** METAL     88 r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                               mem  0.04 GB tm      9.38us/     6.81ms (     3.06 GFLOPS    1.8|5.5     GB/s) ['square', 'sum']\n",
      "hc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m10\u001b[0m               :    17.62 us  <  beam2  : \u001b[34m88\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m2\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m16\u001b[0m            :    19.38 us\n",
      "\u001b[32m*** METAL     89\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m10\u001b[0m\u001b[90m\u001b[0m                       mem  0.04 GB tm     18.04us/     6.82ms (   104.89 GFLOPS   21.9|124.9   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m8\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m          :    42.50 us  <  tc     : \u001b[34m22\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m        :    60.13 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   115.13 us\n",
      "\u001b[32m*** METAL     90\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.04 GB tm     42.13us/     6.87ms (   551.90 GFLOPS   18.7|290.9   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m8\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[35m16\u001b[0m \u001b[33m4\u001b[0m          :    68.00 us  <  tc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :    80.92 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   152.21 us\n",
      "\u001b[32m*** METAL     91\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     67.50us/     6.93ms (   683.52 GFLOPS   18.0|235.0   GB/s) ['linear']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m8\u001b[0m \u001b[32m16\u001b[0m \u001b[35m16\u001b[0m                     :    14.25 us  <  hc     : \u001b[34m44\u001b[0m \u001b[36m4\u001b[0m \u001b[32m8\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m                    :    19.62 us\n",
      "\u001b[32m*** METAL     92\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                             mem  0.05 GB tm     15.04us/     6.95ms (    35.20 GFLOPS   48.1|101.8   GB/s) ['square', 'sum']\n",
      "tc     : \u001b[34m22\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m          :    54.75 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m               :    78.25 us  <  beam2  : \u001b[34m22\u001b[0m \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m \u001b[31m256\u001b[0m                   :    83.71 us\n",
      "\u001b[32m*** METAL     93\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     51.25us/     7.00ms (   113.63 GFLOPS   15.3|22.4    GB/s) ['__add__', 'einsum', '__mul__', 'sqrt', '__neg__']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[35m16\u001b[0m                       :     7.50 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.42 us\n",
      "\u001b[32m*** METAL     94\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                               mem  0.05 GB tm      8.54us/     7.01ms (     1.24 GFLOPS    5.6|5.6     GB/s) ['argmax']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m                       :     9.04 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.54 us\n",
      "\u001b[32m*** METAL     95\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90mn1\u001b[0m                             mem  0.05 GB tm     10.54us/     7.02ms (     4.41 GFLOPS    4.8|9.1     GB/s) ['argmax']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[33m2\u001b[0m                    :     8.29 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                    :     8.63 us\n",
      "\u001b[32m*** METAL     96\u001b[0m E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  0.05 GB tm      8.25us/     7.03ms (     4.10 GFLOPS    5.8|13.7    GB/s) ['one_hot', 'cast']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m2\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m         :    19.75 us  <  tc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m2\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m       :    22.96 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m16\u001b[0m               :    66.29 us\n",
      "\u001b[32m*** METAL     97\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  0.05 GB tm     22.50us/     7.05ms (   272.34 GFLOPS   66.8|136.2   GB/s) ['einsum', '__sub__']\n",
      "beam2  : \u001b[34m2816\u001b[0m \u001b[36m32\u001b[0m \u001b[33m2\u001b[0m                      :    23.25 us  <  hc     : \u001b[34m1408\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    24.00 us\n",
      "\u001b[32m*** METAL     98\u001b[0m E_\u001b[34m2816\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     25.50us/     7.07ms (    28.27 GFLOPS   84.8|84.8    GB/s) ['__add__', '__sub__']\n",
      "beam2  : \u001b[34m176\u001b[0m \u001b[32m16\u001b[0m \u001b[31m64\u001b[0m                      :    14.46 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m4\u001b[0m \u001b[32m8\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m                   :    58.42 us\n",
      "\u001b[32m*** METAL     99\u001b[0m r_\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     14.67us/     7.09ms (    27.65 GFLOPS   49.2|63.0    GB/s) ['square', 'sum']\n",
      "*** METAL    100 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                             mem  0.05 GB tm     15.75us/     7.11ms (    33.61 GFLOPS   45.9|97.3    GB/s) ['square', 'sum']\n",
      "hc     : \u001b[32m16\u001b[0m \u001b[31m11\u001b[0m                          :     7.29 us  <  beam2  : \u001b[32m176\u001b[0m                            :    10.00 us\n",
      "\u001b[32m*** METAL    101\u001b[0m r_\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m11\u001b[0m\u001b[90m\u001b[0m                                  mem  0.05 GB tm      8.38us/     7.11ms (     0.05 GFLOPS    0.1|0.2     GB/s) ['sum', '__add__']\n",
      "*** METAL    102 r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     47.75us/     7.16ms (   121.96 GFLOPS   16.4|24.1    GB/s) ['__add__', 'einsum', '__mul__', 'sqrt', '__neg__']\n",
      "*** METAL    103 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                               mem  0.05 GB tm      7.92us/     7.17ms (     1.33 GFLOPS    6.0|6.0     GB/s) ['argmax']\n",
      "*** METAL    104 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90mn1\u001b[0m                             mem  0.05 GB tm      9.71us/     7.18ms (     4.79 GFLOPS    5.2|9.9     GB/s) ['argmax']\n",
      "*** METAL    105 E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  0.05 GB tm      8.62us/     7.19ms (     3.92 GFLOPS    5.6|13.1    GB/s) ['one_hot', 'cast']\n",
      "*** METAL    106 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  0.05 GB tm     23.38us/     7.21ms (   262.14 GFLOPS   64.3|131.1   GB/s) ['einsum', '__sub__']\n",
      "hc     : \u001b[34m1408\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    36.21 us  <  beam2  : \u001b[34m5632\u001b[0m \u001b[36m32\u001b[0m                        :    48.04 us\n",
      "\u001b[32m*** METAL    107\u001b[0m E_\u001b[34m1408\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     36.33us/     7.25ms (    14.88 GFLOPS   99.2|99.2    GB/s) ['__add__']\n",
      "*** METAL    108 r_\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     14.29us/     7.26ms (    28.37 GFLOPS   50.5|64.6    GB/s) ['square', 'sum']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m    :    85.46 us  <  tc     : \u001b[34m22\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m        :   118.63 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m64\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   164.58 us\n",
      "\u001b[32m*** METAL    109\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m            mem  0.05 GB tm     70.46us/     7.33ms (   657.37 GFLOPS   17.2|214.9   GB/s) ['linear', 'relu']\n",
      "*** METAL    110 r_\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m11\u001b[0m\u001b[90m\u001b[0m                                  mem  0.05 GB tm      7.62us/     7.34ms (     0.06 GFLOPS    0.1|0.3     GB/s) ['sum', '__add__']\n",
      "*** METAL    111 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     40.12us/     7.38ms (   579.41 GFLOPS   19.6|305.4   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[35m2\u001b[0m                              :     7.13 us  <  hc     : \u001b[35m2\u001b[0m                              :     7.13 us\n",
      "\u001b[32m*** METAL    112\u001b[0m r_\u001b[35m2\u001b[0m\u001b[90mn1\u001b[0m                                    mem  0.05 GB tm      7.92us/     7.39ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['stack', 'sum']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m          :    24.42 us  <  hc     : \u001b[34m11\u001b[0m \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m               :    49.79 us\n",
      "\u001b[32m*** METAL    113\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     22.37us/     7.41ms (   130.89 GFLOPS   18.9|70.5    GB/s) ['linear', '__sub__']\n",
      "beam2  : \u001b[32m64\u001b[0m \u001b[31m110\u001b[0m                         :    16.00 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m440\u001b[0m                         :    30.50 us\n",
      "\u001b[32m*** METAL    114\u001b[0m r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m110\u001b[0m\u001b[90m\u001b[0m                                 mem  0.05 GB tm     21.79us/     7.43ms (     1.81 GFLOPS    1.3|2.1     GB/s) ['abs', 'mean', '__mul__', '__add__']\n",
      "\u001b[32m*** CLANG    115\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     85.88us/     7.52ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m                         :     6.42 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                        :     6.87 us\n",
      "\u001b[32m*** METAL    116\u001b[0m E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.05 GB tm      7.63us/     7.53ms (     0.00 GFLOPS    0.7|0.7     GB/s) ['stack']\n",
      "\u001b[32m*** CLANG    117\u001b[0m \u001b[33mcopy     2816,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     35.46us/     7.56ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m8\u001b[0m                           :     7.25 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                        :     7.25 us\n",
      "\u001b[32m*** METAL    118\u001b[0m E_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m\u001b[0m                                   mem  0.05 GB tm      7.67us/     7.57ms (     0.00 GFLOPS    0.7|0.7     GB/s) ['stack']\n",
      "\u001b[32m*** CLANG    119\u001b[0m \u001b[33mcopy     2816,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     14.71us/     7.58ms (     0.00 GFLOPS    0.2|0.2     GB/s) \n",
      "hc     : \u001b[34m22\u001b[0m \u001b[36m32\u001b[0m \u001b[33m2\u001b[0m                        :     7.13 us  <  beam2  : \u001b[34m22\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                       :     7.33 us\n",
      "\u001b[32m*** METAL    120\u001b[0m E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  0.05 GB tm      7.62us/     7.59ms (     0.00 GFLOPS    1.5|1.5     GB/s) ['stack']\n",
      "\u001b[32m*** CLANG    121\u001b[0m \u001b[33mcopy     5632,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     26.29us/     7.62ms (     0.00 GFLOPS    0.2|0.2     GB/s) \n",
      "beam2  : \u001b[32m32\u001b[0m \u001b[32m4\u001b[0m \u001b[31m55\u001b[0m                        :    20.17 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m440\u001b[0m                         :    29.63 us\n",
      "\u001b[32m*** METAL    122\u001b[0m r_\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[32m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m55\u001b[0m\u001b[90m\u001b[0m                                mem  0.05 GB tm     20.62us/     7.64ms (     2.51 GFLOPS    1.4|4.6     GB/s) ['abs', 'mean']\n",
      "\u001b[32m*** CLANG    123\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     17.00us/     7.65ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "Finished discretizing action data!\n",
      "scheduled 45 kernels\n",
      "memory reduced from 0.37 MB -> 0.37 MB, 16 -> 12 bufs\n",
      "beam2  :                                :     7.12 us  <  hc     :                                :     7.12 us\n",
      "\u001b[32m*** METAL    124\u001b[0m E_\u001b[90m\u001b[0m                                       mem  0.05 GB tm      7.58us/     7.66ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__imul__']\n",
      "beam2  :                                :     6.92 us  <  hc     :                                :     6.92 us\n",
      "\u001b[32m*** METAL    125\u001b[0m E_\u001b[90mn1\u001b[0m                                     mem  0.05 GB tm    194.75us/     7.86ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__imul__']\n",
      "beam2  : \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[32m32\u001b[0m \u001b[31m22\u001b[0m                      :    11.17 us  <  hc     : \u001b[34m16\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                       :    11.50 us\n",
      "\u001b[32m*** METAL    126\u001b[0m r_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     14.00us/     7.87ms (     2.01 GFLOPS    3.2|8.2     GB/s) ['sum', '__mul__']\n",
      "*** METAL    127 r_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     11.50us/     7.88ms (     2.45 GFLOPS    3.9|10.0    GB/s) ['sum', '__mul__']\n",
      "beam2  : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m                   :     7.33 us  <  hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     7.63 us\n",
      "\u001b[32m*** METAL    128\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                           mem  0.05 GB tm      7.75us/     7.89ms (     0.00 GFLOPS   16.9|16.9    GB/s) \n",
      "hc     : \u001b[34m55\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     3.25 us  <  beam2  : \u001b[34m55\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                      :     3.54 us\n",
      "\u001b[32m*** METAL    129\u001b[0m E_\u001b[34m55\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.05 GB tm      8.58us/     7.90ms (     3.28 GFLOPS    6.6|6.6     GB/s) ['abs', 'abs bw', '__sub__ bw']\n",
      "hc     : \u001b[34m2\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     6.50 us  <  beam2  : \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m                           :     6.58 us\n",
      "\u001b[32m*** METAL    130\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                               mem  0.05 GB tm      6.62us/     7.91ms (     0.00 GFLOPS    0.2|0.2     GB/s) \n",
      "hc     : \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     6.12 us  <  beam2  : \u001b[34m256\u001b[0m \u001b[33m5\u001b[0m                          :     7.08 us\n",
      "\u001b[32m*** METAL    131\u001b[0m E_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn2\u001b[0m                              mem  0.05 GB tm      6.75us/     7.91ms (     0.00 GFLOPS    0.8|0.8     GB/s) \n",
      "hc     : \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                           :     6.25 us  <  beam2  : \u001b[34m32\u001b[0m \u001b[36m4\u001b[0m                           :     6.50 us\n",
      "\u001b[32m*** METAL    132\u001b[0m E_\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                                 mem  0.05 GB tm      6.67us/     7.92ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                      :     6.75 us  <  hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     6.88 us\n",
      "\u001b[32m*** METAL    133\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm      7.62us/     7.93ms (     0.00 GFLOPS    8.6|8.6     GB/s) \n",
      "hc     : \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     5.96 us  <  beam2  : \u001b[34m40\u001b[0m \u001b[33m2\u001b[0m \u001b[33m16\u001b[0m                        :     6.50 us\n",
      "\u001b[32m*** METAL    134\u001b[0m E_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn3\u001b[0m                              mem  0.05 GB tm      6.62us/     7.93ms (     0.00 GFLOPS    0.8|0.8     GB/s) \n",
      "beam2  : \u001b[34m5\u001b[0m \u001b[33m2\u001b[0m                            :     6.08 us  <  hc     : \u001b[34m5\u001b[0m \u001b[36m2\u001b[0m                            :     6.08 us\n",
      "\u001b[32m*** METAL    135\u001b[0m E_\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                    mem  0.05 GB tm      6.50us/     7.94ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m4\u001b[0m                     :     7.25 us  <  hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     7.88 us\n",
      "\u001b[32m*** METAL    136\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                             mem  0.05 GB tm      7.88us/     7.95ms (     0.00 GFLOPS   16.6|16.6    GB/s) \n",
      "\u001b[32m*** METAL    137\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm     34.75us/     7.98ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[34m32\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m            :    50.62 us  <  tc     : \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m           :   104.00 us  <  hc     : \u001b[34m4\u001b[0m \u001b[36m4\u001b[0m \u001b[36m16\u001b[0m \u001b[31m176\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m               :   134.50 us\n",
      "\u001b[32m*** METAL    138\u001b[0m r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                    mem  0.05 GB tm     38.54us/     8.02ms (   149.63 GFLOPS   20.3|75.2    GB/s) ['einsum']\n",
      "*** METAL    139 r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                    mem  0.05 GB tm     39.62us/     8.06ms (   145.54 GFLOPS   19.7|73.2    GB/s) ['einsum']\n",
      "beam2  :                                :     6.87 us  <  hc     :                                :     6.87 us\n",
      "\u001b[32m*** METAL    140\u001b[0m E_\u001b[90mn2\u001b[0m                                     mem  0.05 GB tm      7.33us/     8.07ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__rsub__']\n",
      "*** METAL    141 E_\u001b[90mn2\u001b[0m                                     mem  0.05 GB tm      7.25us/     8.08ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__rsub__']\n",
      "beam2  : \u001b[31m4\u001b[0m \u001b[35m4\u001b[0m                            :     7.37 us  <  hc     : \u001b[32m16\u001b[0m                             :     7.46 us\n",
      "\u001b[32m*** METAL    142\u001b[0m r_\u001b[31m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                                    mem  0.05 GB tm      7.79us/     8.08ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['sum']\n",
      "*** METAL    143 r_\u001b[31m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                                    mem  0.05 GB tm      7.62us/     8.09ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['sum']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m8\u001b[0m \u001b[31m2\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m           :    15.17 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m10\u001b[0m               :    19.88 us\n",
      "\u001b[32m*** METAL    144\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                   mem  0.05 GB tm     16.71us/     8.11ms (   199.55 GFLOPS   45.1|107.9   GB/s) ['relu bw', 'linear bw']\n",
      "beam2  : \u001b[34m128\u001b[0m \u001b[36m16\u001b[0m \u001b[33m2\u001b[0m                       :     7.62 us  <  hc     : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.71 us\n",
      "\u001b[32m*** METAL    145\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                               mem  0.05 GB tm     51.58us/     8.16ms (     0.24 GFLOPS    1.0|1.0     GB/s) ['mul', '__mul__', 'add']\n",
      "*** METAL    146 E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                               mem  0.05 GB tm      7.50us/     8.17ms (     1.64 GFLOPS    6.6|6.6     GB/s) ['mul', '__mul__', 'add']\n",
      "beam2  : \u001b[34m10\u001b[0m \u001b[32m32\u001b[0m \u001b[31m22\u001b[0m                       :    11.29 us  <  hc     : \u001b[34m10\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                       :    14.00 us\n",
      "\u001b[32m*** METAL    147\u001b[0m r3_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m\u001b[0m                              mem  0.05 GB tm     12.00us/     8.18ms (     1.95 GFLOPS    2.4|6.7     GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  :                                :     7.29 us  <  hc     :                                :     7.29 us\n",
      "\u001b[32m*** METAL    148\u001b[0m E_\u001b[90mn3\u001b[0m                                     mem  0.05 GB tm      7.71us/     8.19ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__add__']\n",
      "*** METAL    149 E_\u001b[90mn3\u001b[0m                                     mem  0.05 GB tm      6.58us/     8.19ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__add__']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[36m8\u001b[0m \u001b[36m2\u001b[0m \u001b[32m32\u001b[0m \u001b[35m22\u001b[0m                   :    15.87 us  <  hc     : \u001b[34m128\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                      :    34.92 us\n",
      "\u001b[32m*** METAL    150\u001b[0m r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m22\u001b[0m\u001b[90m\u001b[0m                          mem  0.05 GB tm     16.00us/     8.21ms (    39.94 GFLOPS   22.7|129.0   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m          :    45.42 us  <  tc     : \u001b[34m11\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :    63.54 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   117.46 us\n",
      "\u001b[32m*** METAL    151\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     40.08us/     8.25ms (   580.01 GFLOPS   28.6|305.7   GB/s) ['relu bw', 'linear bw']\n",
      "beam2  : \u001b[34m5\u001b[0m \u001b[34m8\u001b[0m \u001b[36m16\u001b[0m \u001b[32m8\u001b[0m \u001b[31m88\u001b[0m \u001b[33m2\u001b[0m                  :    24.54 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m128\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                   :   140.46 us\n",
      "\u001b[32m*** METAL    152\u001b[0m r3_\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                         mem  0.05 GB tm     25.12us/     8.27ms (    82.53 GFLOPS   16.5|240.5   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m                          :     8.12 us  <  hc     : \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                            :     8.33 us\n",
      "\u001b[32m*** METAL    153\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                  mem  0.05 GB tm      8.62us/     8.28ms (     0.01 GFLOPS    0.0|0.0     GB/s) ['__add__', '__truediv__', '__mul__']\n",
      "*** METAL    154 E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                  mem  0.05 GB tm      7.71us/     8.29ms (     0.01 GFLOPS    0.0|0.0     GB/s) ['__add__', '__truediv__', '__mul__']\n",
      "*** METAL    155 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m22\u001b[0m\u001b[90m\u001b[0m                          mem  0.05 GB tm     15.13us/     8.31ms (    42.25 GFLOPS   24.0|136.5   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m       :    67.42 us  <  tc     : \u001b[34m11\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :   128.25 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   194.96 us\n",
      "\u001b[32m*** METAL    156\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  0.05 GB tm     68.46us/     8.37ms (   687.11 GFLOPS   38.8|200.1   GB/s) ['square bw', '__add__', '__sub__ bw', 'linear bw']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m22\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[35m4\u001b[0m         :    71.29 us  <  tc     : \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m         :   255.88 us  <  hc     : \u001b[34m4\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m176\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   266.50 us\n",
      "\u001b[32m*** METAL    157\u001b[0m r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                mem  0.05 GB tm     54.38us/     8.43ms (   456.34 GFLOPS   19.3|220.0   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "hc     : \u001b[34m2\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                     :     7.96 us  <  beam2  : \u001b[34m16\u001b[0m \u001b[34m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                      :     9.08 us\n",
      "\u001b[32m*** METAL    158\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                             mem  0.05 GB tm      8.42us/     8.44ms (     0.61 GFLOPS    3.9|4.4     GB/s) ['__truediv__']\n",
      "*** METAL    159 E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                             mem  0.05 GB tm      8.88us/     8.45ms (     0.58 GFLOPS    3.7|4.2     GB/s) ['__truediv__']\n",
      "beam2  : \u001b[34m8\u001b[0m \u001b[36m32\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m \u001b[33m2\u001b[0m                   :    29.37 us  <  hc     : \u001b[34m256\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                      :    57.12 us\n",
      "\u001b[32m*** METAL    160\u001b[0m r3_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m44\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                          mem  0.05 GB tm     26.92us/     8.47ms (    24.80 GFLOPS   27.0|82.2    GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m4\u001b[0m \u001b[34m32\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m          :   117.46 us  <  tc     : \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m         :   148.58 us  <  hc     : \u001b[34m4\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m176\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   263.63 us\n",
      "\u001b[32m*** METAL    161\u001b[0m r3_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  0.05 GB tm     88.08us/     8.56ms (   530.67 GFLOPS   19.7|271.6   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m           :    57.00 us  <  tc     : \u001b[34m11\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :   127.04 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m64\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   182.25 us\n",
      "\u001b[32m*** METAL    162\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                   mem  0.05 GB tm    184.83us/     8.75ms (   250.59 GFLOPS    8.5|128.7   GB/s) ['relu bw', 'linear bw']\n",
      "*** METAL    163 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m22\u001b[0m\u001b[90m\u001b[0m                          mem  0.05 GB tm    209.00us/     8.95ms (     3.06 GFLOPS    1.7|9.9     GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    164 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  0.05 GB tm     40.12us/     8.99ms (   579.41 GFLOPS   28.6|305.4   GB/s) ['relu bw', 'linear bw']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m         :   116.58 us  <  hc     : \u001b[34m8\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m176\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   264.13 us  <  tc     : \u001b[34m2\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m       :   272.42 us\n",
      "\u001b[32m*** METAL    165\u001b[0m r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                mem  0.05 GB tm    109.96us/     9.10ms (   425.03 GFLOPS   15.8|164.2   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m4\u001b[0m \u001b[34m5\u001b[0m \u001b[36m32\u001b[0m \u001b[32m8\u001b[0m \u001b[31m22\u001b[0m \u001b[33m2\u001b[0m \u001b[35m4\u001b[0m                :    25.83 us  <  hc     : \u001b[34m128\u001b[0m \u001b[34m10\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                   :   141.33 us\n",
      "\u001b[32m*** METAL    166\u001b[0m r3_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[32m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                       mem  0.05 GB tm    103.00us/     9.21ms (    20.13 GFLOPS    4.0|58.7    GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    167 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m22\u001b[0m\u001b[90m\u001b[0m                          mem  0.05 GB tm     14.63us/     9.22ms (    43.69 GFLOPS   24.8|141.2   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    168 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                mem  0.05 GB tm     53.75us/     9.28ms (   461.65 GFLOPS   19.5|222.5   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "\u001b[32m*** CLANG    169\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     27.75us/     9.30ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "step: 500 loss: 10226.548\n",
      "The start time is : 570569.330503375\n",
      "\u001b[32m*** METAL    170\u001b[0m \u001b[33mcopy   35.39M,   METAL <- NPY    \u001b[0m        mem  0.08 GB tm   4065.25us/    13.37ms (     0.00 GFLOPS    8.7|8.7     GB/s) \n",
      "\u001b[32m*** METAL    171\u001b[0m \u001b[33mcopy     2560,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm     21.46us/    13.39ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL    172\u001b[0m \u001b[33mcopy     7680,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm     10.71us/    13.40ms (     0.00 GFLOPS    0.7|0.7     GB/s) \n",
      "\u001b[32m*** METAL    173\u001b[0m \u001b[33mcopy      512,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      9.79us/    13.41ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL    174\u001b[0m \u001b[33mcopy      512,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      8.75us/    13.42ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL    175\u001b[0m \u001b[33mcopy      256,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      8.62us/    13.43ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    176\u001b[0m \u001b[33mcopy      256,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      8.29us/    13.44ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    177\u001b[0m \u001b[33mcopy       64,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      7.75us/    13.44ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    178\u001b[0m \u001b[33mcopy       64,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      7.17us/    13.45ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    179\u001b[0m \u001b[33mcopy      512,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      9.00us/    13.46ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "\u001b[32m*** METAL    180\u001b[0m \u001b[33mcopy      320,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      8.25us/    13.47ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    181\u001b[0m \u001b[33mcopy      320,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      7.58us/    13.48ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    182\u001b[0m \u001b[33mcopy      960,   METAL <- NPY    \u001b[0m        mem  0.05 GB tm      7.79us/    13.48ms (     0.00 GFLOPS    0.1|0.1     GB/s) \n",
      "*** METAL    183 r_\u001b[35m3\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      7.50us/    13.49ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG    184\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     25.79us/    13.52ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL    185 r_\u001b[35m3\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      6.42us/    13.52ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG    186\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     14.67us/    13.54ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL    187 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      6.50us/    13.54ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG    188\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     12.25us/    13.56ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL    189 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      6.54us/    13.56ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG    190\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     12.67us/    13.58ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL    191 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      6.38us/    13.58ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG    192\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     12.29us/    13.59ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL    193 r_\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      6.50us/    13.60ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__', 'any']\n",
      "\u001b[32m*** CLANG    194\u001b[0m \u001b[33mcopy        1,   CLANG <- METAL  \u001b[0m        mem  0.05 GB tm     14.75us/    13.62ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[36m3\u001b[0m                              :     6.96 us  <  hc     : \u001b[36m3\u001b[0m                              :     6.96 us\n",
      "\u001b[32m*** METAL    195\u001b[0m E_\u001b[36m3\u001b[0m\u001b[90m\u001b[0m                                      mem  0.05 GB tm      7.33us/    13.62ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__add__']\n",
      "hc     : \u001b[34m40\u001b[0m \u001b[34m3\u001b[0m \u001b[34m192\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m                :   598.04 us  <  beam2  : \u001b[34m72\u001b[0m \u001b[34m3\u001b[0m \u001b[34m40\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m               :   612.96 us\n",
      "\u001b[32m*** METAL    196\u001b[0m E_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                        mem  0.08 GB tm    575.46us/    14.20ms (    41.00 GFLOPS  123.0|164.0   GB/s) ['__sub__', '__truediv__']\n",
      "\u001b[32m*** CLANG    197\u001b[0m \u001b[33mcopy   35.39M,   CLANG <- METAL  \u001b[0m        mem  0.12 GB tm   2991.62us/    17.19ms (     0.00 GFLOPS   11.8|11.8    GB/s) \n",
      "scheduled 13 kernels\n",
      "beam2  : \u001b[34m24\u001b[0m \u001b[31m96\u001b[0m \u001b[33m4\u001b[0m                        :     5.79 us  <  hc     : \u001b[34m96\u001b[0m \u001b[32m16\u001b[0m \u001b[31m6\u001b[0m                        :     7.25 us\n",
      "\u001b[32m*** METAL    198\u001b[0m r_\u001b[34m24\u001b[0m\u001b[90m_\u001b[0m\u001b[31m96\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.08 GB tm      6.75us/    17.20ms (     0.00 GFLOPS    0.1|0.1     GB/s) ['arange']\n",
      "\u001b[32m*** CUSTOM   199\u001b[0m custom_random                            mem  0.08 GB \n",
      "\u001b[32m*** CUSTOM   200\u001b[0m custom_random                            mem  0.08 GB \n",
      "beam2  : \u001b[34m5\u001b[0m \u001b[36m4\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                        :     7.25 us  <  hc     : \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                         :     7.25 us\n",
      "\u001b[32m*** METAL    201\u001b[0m E_\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.08 GB tm      7.83us/    17.21ms (     0.04 GFLOPS    0.3|0.3     GB/s) ['randint']\n",
      "*** METAL    202 E_\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.08 GB tm      7.13us/    17.21ms (     0.04 GFLOPS    0.4|0.4     GB/s) ['randint']\n",
      "hc     : \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                         :     6.62 us  <  beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m                          :     7.29 us\n",
      "\u001b[32m*** METAL    203\u001b[0m E_\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.08 GB tm      8.00us/    17.22ms (     0.04 GFLOPS    0.3|0.3     GB/s) ['__add__']\n",
      "beam2  : \u001b[34m80\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m6\u001b[0m                      :     7.50 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m                  :     8.17 us\n",
      "\u001b[32m*** METAL    204\u001b[0m E_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m6\u001b[0m\u001b[90m\u001b[0m                              mem  0.08 GB tm      8.29us/    17.23ms (     7.41 GFLOPS    3.9|21.0    GB/s) ['__ge__']\n",
      "*** METAL    205 E_\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.08 GB tm      7.62us/    17.24ms (     0.04 GFLOPS    0.3|0.3     GB/s) ['__add__']\n",
      "hc     : \u001b[34m10\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m                  :     4.13 us  <  beam2  : \u001b[34m5\u001b[0m \u001b[36m32\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m                  :     9.87 us\n",
      "\u001b[32m*** METAL    206\u001b[0m E_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                          mem  0.08 GB tm      9.29us/    17.25ms (    13.22 GFLOPS    3.6|15.4    GB/s) ['__ge__', '__lt__', '__mul__']\n",
      "hc     : \u001b[34m10\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m                  :     8.50 us  <  beam2  : \u001b[34m4\u001b[0m \u001b[34m6\u001b[0m \u001b[36m16\u001b[0m \u001b[36m16\u001b[0m \u001b[33m5\u001b[0m                    :     8.88 us\n",
      "\u001b[32m*** METAL    207\u001b[0m E_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                        mem  0.08 GB tm      9.29us/    17.25ms (     3.31 GFLOPS    3.5|11.0    GB/s) ['__lt__']\n",
      "beam2  : \u001b[34m320\u001b[0m \u001b[34m4\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m6\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m               :    51.04 us  <  hc     : \u001b[34m320\u001b[0m \u001b[34m4\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m3\u001b[0m               :    56.62 us\n",
      "\u001b[32m*** METAL    208\u001b[0m E_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m6\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                       mem  0.09 GB tm     50.33us/    17.31ms (   117.18 GFLOPS   60.4|85.4    GB/s) ['__mul__']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m72\u001b[0m \u001b[34m3\u001b[0m \u001b[36m4\u001b[0m \u001b[36m8\u001b[0m \u001b[33m5\u001b[0m \u001b[33m4\u001b[0m                :   126.75 us  <  hc     : \u001b[34m40\u001b[0m \u001b[34m576\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m                  :   286.25 us\n",
      "\u001b[32m*** METAL    209\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m72\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                        mem  0.10 GB tm     96.08us/    17.40ms (     0.00 GFLOPS  122.8|184.2   GB/s) \n",
      "\u001b[32m*** CLANG    210\u001b[0m \u001b[33mcopy    8.85M,   CLANG <- METAL  \u001b[0m        mem  0.11 GB tm    780.17us/    18.18ms (     0.00 GFLOPS   11.3|11.3    GB/s) \n",
      "scheduled 162 kernels\n",
      "\u001b[32m*** METAL    211\u001b[0m \u001b[33mcopy   27.10M,   METAL <- NPY    \u001b[0m        mem  0.14 GB tm   3009.83us/    21.19ms (     0.00 GFLOPS    9.0|9.0     GB/s) \n",
      "\u001b[32m*** CUSTOM   212\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   213\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   214\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   215\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   216\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   217\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   218\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   219\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   220\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   221\u001b[0m custom_random                            mem  0.11 GB \n",
      "\u001b[32m*** CUSTOM   222\u001b[0m custom_random                            mem  0.12 GB \n",
      "\u001b[32m*** CUSTOM   223\u001b[0m custom_random                            mem  0.12 GB \n",
      "\u001b[32m*** CUSTOM   224\u001b[0m custom_random                            mem  0.12 GB \n",
      "\u001b[32m*** CUSTOM   225\u001b[0m custom_random                            mem  0.12 GB \n",
      "\u001b[32m*** CUSTOM   226\u001b[0m custom_random                            mem  0.12 GB \n",
      "\u001b[32m*** CUSTOM   227\u001b[0m custom_random                            mem  0.13 GB \n",
      "\u001b[32m*** CUSTOM   228\u001b[0m custom_random                            mem  0.14 GB \n",
      "\u001b[32m*** CUSTOM   229\u001b[0m custom_random                            mem  0.14 GB \n",
      "\u001b[32m*** CUSTOM   230\u001b[0m custom_random                            mem  0.15 GB \n",
      "\u001b[32m*** CUSTOM   231\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   232\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   233\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** METAL    234\u001b[0m \u001b[33mcopy       36,   METAL <- NPY    \u001b[0m        mem  0.16 GB tm     14.17us/    21.21ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** METAL    235\u001b[0m \u001b[33mcopy       36,   METAL <- NPY    \u001b[0m        mem  0.16 GB tm      8.25us/    21.21ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** CUSTOM   236\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   237\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   238\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   239\u001b[0m custom_random                            mem  0.16 GB \n",
      "*** METAL    240 E_\u001b[36m2\u001b[0m\u001b[90m\u001b[0m                                      mem  0.16 GB tm      7.38us/    21.22ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__sub__', '__add__']\n",
      "\u001b[32m*** CUSTOM   241\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   242\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   243\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   244\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   245\u001b[0m custom_random                            mem  0.16 GB \n",
      "beam2  : \u001b[34m25\u001b[0m \u001b[31m500\u001b[0m \u001b[33m4\u001b[0m \u001b[33m5\u001b[0m                     :     6.50 us  <  hc     : \u001b[34m125\u001b[0m \u001b[36m4\u001b[0m \u001b[31m125\u001b[0m \u001b[35m4\u001b[0m                    :     6.79 us\n",
      "\u001b[32m*** METAL    246\u001b[0m r_\u001b[34m25\u001b[0m\u001b[90m_\u001b[0m\u001b[31m500\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                             mem  0.16 GB tm      7.13us/    21.23ms (     0.00 GFLOPS    0.3|0.3     GB/s) ['arange']\n",
      "beam2  : \u001b[34m7\u001b[0m \u001b[36m3\u001b[0m \u001b[31m21\u001b[0m                         :     6.25 us  <  hc     : \u001b[34m7\u001b[0m \u001b[36m3\u001b[0m \u001b[35m21\u001b[0m                         :     7.75 us\n",
      "\u001b[32m*** METAL    247\u001b[0m r_\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                                 mem  0.16 GB tm    117.75us/    21.35ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['arange']\n",
      "\u001b[32m*** CUSTOM   248\u001b[0m custom_random                            mem  0.16 GB \n",
      "\u001b[32m*** CUSTOM   249\u001b[0m custom_random                            mem  0.16 GB \n",
      "beam2  : \u001b[34m84\u001b[0m \u001b[36m16\u001b[0m \u001b[33m7\u001b[0m                        :     7.75 us  <  hc     : \u001b[34m98\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                        :     7.87 us\n",
      "\u001b[32m*** METAL    250\u001b[0m E_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  0.16 GB tm      9.00us/    21.35ms (     2.09 GFLOPS    8.4|8.4     GB/s) ['uniform']\n",
      "hc     : \u001b[34m384\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                       :     9.13 us  <  beam2  : \u001b[34m160\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[33m2\u001b[0m                     :     9.63 us\n",
      "\u001b[32m*** METAL    251\u001b[0m E_\u001b[34m384\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                               mem  0.16 GB tm     10.38us/    21.37ms (     7.11 GFLOPS   28.4|28.4    GB/s) ['uniform']\n",
      "*** METAL    252 E_\u001b[34m384\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                               mem  0.16 GB tm      9.25us/    21.37ms (     7.97 GFLOPS   31.9|31.9    GB/s) ['uniform']\n",
      "*** METAL    253 E_\u001b[34m384\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                               mem  0.16 GB tm      9.17us/    21.38ms (     8.04 GFLOPS   32.2|32.2    GB/s) ['uniform']\n",
      "*** METAL    254 E_\u001b[34m384\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                               mem  0.16 GB tm    520.46us/    21.90ms (     0.14 GFLOPS    0.6|0.6     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m288\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                     :    10.88 us  <  hc     : \u001b[34m768\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                       :    11.37 us\n",
      "\u001b[32m*** METAL    255\u001b[0m E_\u001b[34m288\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                             mem  0.16 GB tm     12.00us/    21.92ms (    12.29 GFLOPS   49.2|49.2    GB/s) ['uniform']\n",
      "hc     : \u001b[34m1536\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    15.79 us  <  beam2  : \u001b[34m288\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m                     :    16.25 us\n",
      "\u001b[32m*** METAL    256\u001b[0m E_\u001b[34m1536\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm    147.58us/    22.06ms (     2.00 GFLOPS    8.0|8.0     GB/s) ['uniform']\n",
      "hc     : \u001b[34m64\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.33 us  <  beam2  : \u001b[34m64\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m                        :     7.42 us\n",
      "\u001b[32m*** METAL    257\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  0.16 GB tm      8.38us/    22.07ms (     1.96 GFLOPS    7.8|7.8     GB/s) ['uniform']\n",
      "*** METAL    258 E_\u001b[34m1536\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm    449.33us/    22.52ms (     0.66 GFLOPS    2.6|2.6     GB/s) ['uniform']\n",
      "*** METAL    259 E_\u001b[34m1536\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm     15.88us/    22.54ms (    18.58 GFLOPS   74.3|74.3    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m2304\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    23.08 us  <  hc     : \u001b[34m3072\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    24.75 us\n",
      "\u001b[32m*** METAL    260\u001b[0m E_\u001b[34m2304\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm     27.08us/    22.56ms (    21.78 GFLOPS   87.1|87.1    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m4608\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m                      :    41.00 us  <  hc     : \u001b[34m6144\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    52.50 us\n",
      "\u001b[32m*** METAL    261\u001b[0m E_\u001b[34m4608\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm     82.79us/    22.65ms (    14.25 GFLOPS   57.0|57.0    GB/s) ['uniform']\n",
      "hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     8.87 us  <  beam2  : \u001b[34m256\u001b[0m \u001b[36m8\u001b[0m \u001b[36m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m                    :     9.00 us\n",
      "\u001b[32m*** METAL    262\u001b[0m E_\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                             mem  0.16 GB tm     10.00us/    22.66ms (     6.55 GFLOPS   26.2|26.2    GB/s) ['uniform']\n",
      "*** METAL    263 E_\u001b[34m4608\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm     46.83us/    22.70ms (    25.19 GFLOPS  100.8|100.8   GB/s) ['uniform']\n",
      "*** METAL    264 E_\u001b[34m4608\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm    465.00us/    23.17ms (     2.54 GFLOPS   10.1|10.1    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m9216\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    91.21 us  <  hc     : \u001b[34m12288\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                     :    97.63 us\n",
      "\u001b[32m*** METAL    265\u001b[0m E_\u001b[34m9216\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.17 GB tm     80.25us/    23.25ms (    29.40 GFLOPS  117.6|117.6   GB/s) ['uniform']\n",
      "beam2  : \u001b[34m768\u001b[0m \u001b[36m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                    :   155.88 us  <  hc     : \u001b[34m24576\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                     :   176.46 us\n",
      "\u001b[32m*** METAL    266\u001b[0m E_\u001b[34m768\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                            mem  0.17 GB tm    293.25us/    23.54ms (    16.09 GFLOPS   64.4|64.4    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m1024\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                    :    14.13 us  <  hc     : \u001b[34m1024\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    14.21 us\n",
      "\u001b[32m*** METAL    267\u001b[0m E_\u001b[34m1024\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  0.16 GB tm     16.87us/    23.56ms (    15.53 GFLOPS   62.1|62.1    GB/s) ['uniform']\n",
      "*** METAL    268 E_\u001b[34m768\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                            mem  0.17 GB tm    153.29us/    23.71ms (    30.78 GFLOPS  123.1|123.1   GB/s) ['uniform']\n",
      "*** METAL    269 E_\u001b[34m768\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                            mem  0.17 GB tm    148.17us/    23.86ms (    31.85 GFLOPS  127.4|127.4   GB/s) ['uniform']\n",
      "hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     7.71 us  <  beam2  : \u001b[34m128\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m                       :     8.00 us\n",
      "\u001b[32m*** METAL    270\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                             mem  0.16 GB tm     12.63us/    23.87ms (     2.60 GFLOPS   10.4|10.4    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                          :     7.00 us  <  hc     : \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m                            :     7.04 us\n",
      "\u001b[32m*** METAL    271\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                  mem  0.16 GB tm      7.67us/    23.88ms (     0.01 GFLOPS    0.0|0.0     GB/s) ['uniform']\n",
      "hc     : \u001b[34m3\u001b[0m \u001b[36m3\u001b[0m \u001b[33m2\u001b[0m                          :     7.00 us  <  beam2  : \u001b[36m3\u001b[0m \u001b[33m2\u001b[0m \u001b[33m3\u001b[0m                          :     7.25 us\n",
      "\u001b[32m*** METAL    272\u001b[0m E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                  mem  0.16 GB tm      7.75us/    23.89ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['cat']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m8\u001b[0m \u001b[36m8\u001b[0m                         :     7.25 us  <  hc     : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.33 us\n",
      "\u001b[32m*** METAL    273\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m\u001b[0m                                 mem  0.16 GB tm      8.21us/    23.90ms (     1.00 GFLOPS    4.0|4.0     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m32\u001b[0m                           :     6.38 us  <  hc     : \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                           :     7.00 us\n",
      "\u001b[32m*** METAL    274\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m\u001b[0m                                   mem  0.16 GB tm      6.88us/    23.90ms (     0.02 GFLOPS    0.1|0.1     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m256\u001b[0m \u001b[36m16\u001b[0m \u001b[33m8\u001b[0m                       :     8.92 us  <  hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     9.38 us\n",
      "\u001b[32m*** METAL    275\u001b[0m E_\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90m\u001b[0m                               mem  0.16 GB tm      9.87us/    23.91ms (     6.64 GFLOPS   26.5|26.5    GB/s) ['uniform']\n",
      "hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     6.67 us  <  beam2  : \u001b[34m32\u001b[0m \u001b[36m16\u001b[0m                          :     7.29 us\n",
      "\u001b[32m*** METAL    276\u001b[0m E_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.16 GB tm      8.17us/    23.92ms (     0.13 GFLOPS    0.5|0.5     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m40\u001b[0m \u001b[34m2\u001b[0m \u001b[36m8\u001b[0m                         :     7.29 us  <  hc     : \u001b[34m10\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m                        :     7.29 us\n",
      "\u001b[32m*** METAL    277\u001b[0m E_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m\u001b[0m                                 mem  0.16 GB tm      7.83us/    23.93ms (     0.49 GFLOPS    0.7|1.3     GB/s) ['__sub__', '__truediv__', '__mul__']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m                          :     7.00 us  <  hc     : \u001b[34m8\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     7.00 us\n",
      "\u001b[32m*** METAL    278\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m\u001b[0m                                  mem  0.16 GB tm      7.75us/    23.94ms (     0.26 GFLOPS    1.1|1.1     GB/s) ['uniform']\n",
      "hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     7.00 us  <  beam2  : \u001b[34m16\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m                         :     7.12 us\n",
      "\u001b[32m*** METAL    279\u001b[0m E_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                               mem  0.16 GB tm      8.25us/    23.95ms (     0.12 GFLOPS    0.5|0.5     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m                         :     7.67 us  <  hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :    11.63 us\n",
      "\u001b[32m*** METAL    280\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m\u001b[0m                                 mem  0.16 GB tm      9.04us/    23.95ms (     0.51 GFLOPS    0.7|0.7     GB/s) ['randn']\n",
      "hc     : \u001b[34m2048\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    21.79 us  <  beam2  : \u001b[34m2048\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    22.25 us\n",
      "\u001b[32m*** METAL    281\u001b[0m E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm     22.96us/    23.98ms (    22.84 GFLOPS   91.3|91.3    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m                         :     7.37 us  <  hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     7.42 us\n",
      "\u001b[32m*** METAL    282\u001b[0m E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  0.16 GB tm      7.75us/    23.99ms (     0.13 GFLOPS    0.5|0.5     GB/s) ['uniform']\n",
      "hc     : \u001b[34m2000\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    21.50 us  <  beam2  : \u001b[34m4000\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                      :    27.29 us\n",
      "\u001b[32m*** METAL    283\u001b[0m E_\u001b[34m2000\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  0.16 GB tm     25.04us/    24.01ms (    30.67 GFLOPS   81.8|81.8    GB/s) ['glorot_uniform']\n",
      "beam2  : \u001b[34m42\u001b[0m \u001b[34m40\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m3\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m3\u001b[0m \u001b[31m7\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m7\u001b[0m \u001b[33m4\u001b[0m :  3695.25 us  <  hc     : \u001b[34m320\u001b[0m \u001b[34m14\u001b[0m \u001b[34m7\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[36m2\u001b[0m \u001b[31m3\u001b[0m \u001b[31m7\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m7\u001b[0m      :  6041.42 us\n",
      "\u001b[32m*** METAL    284\u001b[0m r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m7\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m        mem  0.31 GB tm \u001b[33m    11.67ms\u001b[0m/    35.69ms (  1057.53 GFLOPS   14.7|114.5   GB/s) ['conv2d']\n",
      "hc     : \u001b[34m10\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m2\u001b[0m                :    14.33 us  <  beam2  : \u001b[34m10\u001b[0m \u001b[34m512\u001b[0m \u001b[36m32\u001b[0m \u001b[31m2\u001b[0m                    :    48.21 us\n",
      "\u001b[32m*** METAL    285\u001b[0m r_\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                        mem  0.31 GB tm      4.50us/    35.69ms (   145.64 GFLOPS  147.6|327.7   GB/s) ['linear']\n",
      "beam2  : \u001b[34m7\u001b[0m \u001b[34m16\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[36m3\u001b[0m \u001b[31m500\u001b[0m                 :   115.96 us  <  hc     : \u001b[34m7\u001b[0m \u001b[34m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m125\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m               :   187.46 us\n",
      "\u001b[32m*** METAL    286\u001b[0m r_\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m500\u001b[0m\u001b[90m\u001b[0m                         mem  0.31 GB tm     39.83us/    35.73ms (   539.85 GFLOPS   26.8|1082.9  GB/s) ['__eq__', 'mul', 'sum']\n",
      "beam2  : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m1764\u001b[0m                   :  1065.00 us  <  hc     : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m1764\u001b[0m                   :  1405.08 us\n",
      "\u001b[32m*** METAL    287\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1764\u001b[0m\u001b[90m\u001b[0m                           mem  0.31 GB tm   1552.12us/    37.28ms (    23.50 GFLOPS   93.1|94.1    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m441\u001b[0m \u001b[35m4\u001b[0m                  :  1056.00 us  <  hc     : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m1764\u001b[0m                   :  1103.67 us\n",
      "\u001b[32m*** METAL    288\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                          mem  0.31 GB tm   1081.33us/    38.36ms (   100.63 GFLOPS  133.6|135.1   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m10\u001b[0m \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m \u001b[33m11\u001b[0m \u001b[33m2\u001b[0m           :  3288.25 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m4\u001b[0m \u001b[34m11\u001b[0m \u001b[34m11\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m           :  3552.13 us\n",
      "\u001b[32m*** METAL    289\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m11\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                   mem  0.47 GB tm   4154.12us/    42.52ms (    37.53 GFLOPS   73.0|76.4    GB/s) ['layernorm', 'relu', 'max_pool2d']\n",
      "beam2  : \u001b[34m640\u001b[0m \u001b[34m7\u001b[0m \u001b[36m32\u001b[0m \u001b[36m3\u001b[0m \u001b[36m3\u001b[0m \u001b[31m3\u001b[0m \u001b[31m3\u001b[0m \u001b[33m7\u001b[0m             :  1592.62 us  <  hc     : \u001b[34m2560\u001b[0m \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m8\u001b[0m \u001b[36m3\u001b[0m \u001b[36m3\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m             :  1675.00 us\n",
      "\u001b[32m*** METAL    290\u001b[0m r_\u001b[34m640\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                     mem  0.50 GB tm   2486.71us/    45.00ms (    32.69 GFLOPS   78.3|145.3   GB/s) ['max_pool2d']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[34m320\u001b[0m \u001b[34m7\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m3\u001b[0m \u001b[31m8\u001b[0m \u001b[31m3\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m3\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m :  3207.67 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[31m64\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m      :  6767.12 us\n",
      "\u001b[32m*** METAL    291\u001b[0m r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m      mem  0.54 GB tm   9485.62us/    54.49ms (  1292.01 GFLOPS    7.6|187.2   GB/s) ['conv2d']\n",
      "hc     : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m441\u001b[0m                    :   374.75 us  <  beam2  : \u001b[34m1280\u001b[0m \u001b[32m8\u001b[0m \u001b[31m882\u001b[0m                     :   469.79 us\n",
      "\u001b[32m*** METAL    292\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90m\u001b[0m                            mem  0.54 GB tm    647.33us/    55.14ms (    14.49 GFLOPS   55.8|58.1    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m1280\u001b[0m \u001b[32m49\u001b[0m \u001b[31m144\u001b[0m                    :   288.25 us  <  hc     : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m441\u001b[0m                    :   320.00 us\n",
      "\u001b[32m*** METAL    293\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m49\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  0.54 GB tm    276.75us/    55.41ms (   110.14 GFLOPS  130.6|177.7   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m20\u001b[0m \u001b[34m64\u001b[0m \u001b[34m3\u001b[0m \u001b[36m16\u001b[0m \u001b[33m7\u001b[0m \u001b[33m7\u001b[0m \u001b[33m3\u001b[0m               :   764.63 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m4\u001b[0m \u001b[34m49\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m3\u001b[0m              :   859.58 us\n",
      "\u001b[32m*** METAL    294\u001b[0m E_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                       mem  0.57 GB tm    630.00us/    56.04ms (    43.11 GFLOPS  114.7|115.5   GB/s) ['layernorm', 'relu']\n",
      "*** METAL    295 r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m      mem  0.61 GB tm   6531.71us/    62.58ms (  1876.32 GFLOPS   11.1|271.8   GB/s) ['conv2d']\n",
      "*** METAL    296 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90m\u001b[0m                            mem  0.61 GB tm    315.25us/    62.89ms (    29.75 GFLOPS  114.6|119.3   GB/s) ['layernorm']\n",
      "*** METAL    297 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m49\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  0.61 GB tm    283.00us/    63.17ms (   107.71 GFLOPS  127.7|173.8   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m441\u001b[0m \u001b[34m160\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m                :   603.04 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m147\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m                 :   607.87 us\n",
      "\u001b[32m*** METAL    298\u001b[0m E_\u001b[34m441\u001b[0m\u001b[90m_\u001b[0m\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                        mem  0.65 GB tm    995.42us/    64.17ms (    22.68 GFLOPS   72.6|108.9   GB/s) ['layernorm']\n",
      "*** METAL    299 r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m      mem  0.68 GB tm   3185.00us/    67.35ms (  3847.90 GFLOPS   22.7|557.4   GB/s) ['conv2d']\n",
      "*** METAL    300 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90m\u001b[0m                            mem  0.68 GB tm    298.83us/    67.65ms (    31.39 GFLOPS  120.9|125.8   GB/s) ['layernorm']\n",
      "*** METAL    301 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m49\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  0.68 GB tm    291.50us/    67.94ms (   104.57 GFLOPS  124.0|168.7   GB/s) ['layernorm']\n",
      "*** METAL    302 E_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                       mem  0.72 GB tm    637.62us/    68.58ms (    42.59 GFLOPS  113.3|114.1   GB/s) ['layernorm', 'relu']\n",
      "*** METAL    303 r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m      mem  0.75 GB tm   3407.42us/    71.99ms (  3596.73 GFLOPS   21.2|521.0   GB/s) ['conv2d']\n",
      "*** METAL    304 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90m\u001b[0m                            mem  0.75 GB tm    340.96us/    72.33ms (    27.51 GFLOPS  106.0|110.3   GB/s) ['layernorm']\n",
      "*** METAL    305 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m49\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  0.75 GB tm    292.12us/    72.62ms (   104.35 GFLOPS  123.7|168.3   GB/s) ['layernorm']\n",
      "*** METAL    306 E_\u001b[34m441\u001b[0m\u001b[90m_\u001b[0m\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                        mem  0.79 GB tm    564.46us/    73.19ms (    40.00 GFLOPS  128.0|192.0   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m11\u001b[0m \u001b[34m10\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m32\u001b[0m \u001b[31m8\u001b[0m \u001b[31m3\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m :  3919.83 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m2\u001b[0m \u001b[34m11\u001b[0m \u001b[34m11\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m64\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m     :  5829.00 us\n",
      "\u001b[32m*** METAL    307\u001b[0m r_\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m     mem  0.81 GB tm   5025.88us/    78.21ms (  1730.21 GFLOPS   11.2|315.6   GB/s) ['conv2d']\n",
      "beam2  : \u001b[34m320\u001b[0m \u001b[34m4\u001b[0m \u001b[34m11\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m8\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m     :   756.92 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m2\u001b[0m \u001b[34m11\u001b[0m \u001b[34m11\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m       :  1538.50 us\n",
      "\u001b[32m*** METAL    308\u001b[0m r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m             mem  0.83 GB tm   1043.25us/    79.26ms (   884.49 GFLOPS   53.7|304.0   GB/s) ['conv2d']\n",
      "beam2  : \u001b[34m160\u001b[0m \u001b[36m16\u001b[0m \u001b[31m1936\u001b[0m                    :   238.67 us  <  hc     : \u001b[34m20\u001b[0m \u001b[36m32\u001b[0m \u001b[31m484\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m                  :   400.21 us\n",
      "\u001b[32m*** METAL    309\u001b[0m r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                            mem  0.83 GB tm    177.71us/    79.43ms (    27.90 GFLOPS  111.6|111.6   GB/s) ['layernorm']\n",
      "*** METAL    310 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                            mem  0.83 GB tm    590.75us/    80.02ms (     8.39 GFLOPS   33.6|33.6    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m2560\u001b[0m \u001b[32m16\u001b[0m \u001b[31m121\u001b[0m                    :   185.87 us  <  hc     : \u001b[34m20\u001b[0m \u001b[36m32\u001b[0m \u001b[31m484\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m                  :   425.25 us\n",
      "\u001b[32m*** METAL    311\u001b[0m r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m121\u001b[0m\u001b[90m\u001b[0m                            mem  0.83 GB tm    628.46us/    80.65ms (    25.03 GFLOPS   31.6|36.5    GB/s) ['layernorm']\n",
      "*** METAL    312 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m121\u001b[0m\u001b[90m\u001b[0m                            mem  0.83 GB tm    160.12us/    80.81ms (    98.23 GFLOPS  123.9|143.2   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m40\u001b[0m \u001b[34m2\u001b[0m \u001b[34m121\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                :  1258.79 us  <  beam2  : \u001b[34m20\u001b[0m \u001b[34m8\u001b[0m \u001b[34m121\u001b[0m \u001b[36m16\u001b[0m \u001b[36m32\u001b[0m                 :  1323.00 us\n",
      "\u001b[32m*** METAL    313\u001b[0m E_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m121\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                        mem  0.85 GB tm    975.88us/    81.79ms (    16.51 GFLOPS   40.7|50.8    GB/s) ['layernorm', 'relu']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[34m320\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m11\u001b[0m \u001b[35m3\u001b[0m :  4701.71 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m2\u001b[0m \u001b[34m11\u001b[0m \u001b[34m11\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m    :  9383.00 us\n",
      "\u001b[32m*** METAL    314\u001b[0m r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m11\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m         mem  0.87 GB tm \u001b[33m    11.38ms\u001b[0m/    93.16ms (  1464.29 GFLOPS    3.5|157.4   GB/s) ['conv2d']\n",
      "*** METAL    315 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                            mem  0.87 GB tm    171.58us/    93.34ms (    28.90 GFLOPS  115.6|115.6   GB/s) ['layernorm']\n",
      "*** METAL    316 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m121\u001b[0m\u001b[90m\u001b[0m                            mem  0.87 GB tm    134.79us/    93.47ms (   116.69 GFLOPS  147.2|170.2   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m40\u001b[0m \u001b[34m8\u001b[0m \u001b[34m121\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m                  :   821.29 us  <  beam2  : \u001b[34m20\u001b[0m \u001b[34m128\u001b[0m \u001b[34m121\u001b[0m \u001b[36m16\u001b[0m                  :  2715.08 us\n",
      "\u001b[32m*** METAL    317\u001b[0m E_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m121\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m\u001b[0m                          mem  0.89 GB tm   1081.25us/    94.55ms (    32.09 GFLOPS   55.0|128.3   GB/s) ['layernorm', '__add__']\n",
      "*** METAL    318 r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m11\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m         mem  0.91 GB tm   4255.58us/    98.81ms (  3914.07 GFLOPS    9.5|420.7   GB/s) ['conv2d']\n",
      "*** METAL    319 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                            mem  0.91 GB tm    295.71us/    99.10ms (    16.77 GFLOPS   67.1|67.1    GB/s) ['layernorm']\n",
      "*** METAL    320 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m121\u001b[0m\u001b[90m\u001b[0m                            mem  0.91 GB tm    163.50us/    99.27ms (    96.20 GFLOPS  121.4|140.3   GB/s) ['layernorm']\n",
      "*** METAL    321 E_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m121\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                        mem  0.93 GB tm    819.00us/   100.09ms (    19.67 GFLOPS   48.4|60.5    GB/s) ['layernorm', 'relu']\n",
      "*** METAL    322 r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m11\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m         mem  0.95 GB tm   4216.13us/   104.30ms (  3950.70 GFLOPS    9.5|424.7   GB/s) ['conv2d']\n",
      "*** METAL    323 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                            mem  0.95 GB tm    446.17us/   104.75ms (    11.11 GFLOPS   44.5|44.5    GB/s) ['layernorm']\n",
      "*** METAL    324 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m121\u001b[0m\u001b[90m\u001b[0m                            mem  0.95 GB tm    138.38us/   104.89ms (   113.67 GFLOPS  143.4|165.8   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m80\u001b[0m \u001b[34m121\u001b[0m \u001b[36m4\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                :   349.75 us  <  hc     : \u001b[34m80\u001b[0m \u001b[34m121\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                  :   367.50 us\n",
      "\u001b[32m*** METAL    325\u001b[0m E_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[34m121\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                        mem  0.97 GB tm    544.87us/   105.43ms (    20.47 GFLOPS   72.8|91.0    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m20\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m6\u001b[0m \u001b[35m3\u001b[0m \u001b[33m2\u001b[0m :  2663.17 us  <  hc     : \u001b[34m40\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[31m3\u001b[0m \u001b[31m3\u001b[0m \u001b[33m6\u001b[0m \u001b[33m6\u001b[0m \u001b[33m4\u001b[0m        : 18791.75 us\n",
      "\u001b[32m*** METAL    326\u001b[0m r_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m6\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m        mem  0.98 GB tm   4620.75us/   110.05ms (  1860.24 GFLOPS    7.1|234.9   GB/s) ['conv2d']\n",
      "hc     : \u001b[34m320\u001b[0m \u001b[34m4\u001b[0m \u001b[34m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m        :  1562.83 us  <  beam2  : \u001b[34m320\u001b[0m \u001b[34m8\u001b[0m \u001b[34m6\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m       :  1948.50 us\n",
      "\u001b[32m*** METAL    327\u001b[0m r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                mem  0.99 GB tm   1299.75us/   111.35ms (   580.86 GFLOPS   24.4|686.7   GB/s) ['conv2d']\n",
      "beam2  : \u001b[34m320\u001b[0m \u001b[36m16\u001b[0m \u001b[32m32\u001b[0m \u001b[31m18\u001b[0m                   :   115.71 us  <  hc     : \u001b[34m40\u001b[0m \u001b[36m32\u001b[0m \u001b[31m144\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m                  :   193.08 us\n",
      "\u001b[32m*** METAL    328\u001b[0m r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  0.99 GB tm     96.50us/   111.45ms (    86.59 GFLOPS  122.5|353.1   GB/s) ['layernorm']\n",
      "*** METAL    329 r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  0.99 GB tm     88.50us/   111.54ms (    94.42 GFLOPS  133.5|385.1   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m2560\u001b[0m \u001b[36m2\u001b[0m \u001b[32m32\u001b[0m \u001b[31m18\u001b[0m                   :   113.21 us  <  hc     : \u001b[34m40\u001b[0m \u001b[36m32\u001b[0m \u001b[31m144\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m                  :   198.25 us\n",
      "\u001b[32m*** METAL    330\u001b[0m r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  0.99 GB tm    189.50us/   111.73ms (    78.68 GFLOPS   62.5|183.3   GB/s) ['layernorm']\n",
      "*** METAL    331 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  0.99 GB tm    275.21us/   112.00ms (    54.18 GFLOPS   43.0|126.2   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m160\u001b[0m \u001b[34m16\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[33m9\u001b[0m                :   201.00 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m16\u001b[0m \u001b[34m3\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m3\u001b[0m              :   230.00 us\n",
      "\u001b[32m*** METAL    332\u001b[0m E_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m9\u001b[0m\u001b[90m\u001b[0m                        mem  1.00 GB tm    169.37us/   112.17ms (    54.17 GFLOPS  139.5|154.8   GB/s) ['layernorm', 'relu']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[34m80\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m6\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m :  4188.63 us  <  hc     : \u001b[34m40\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m256\u001b[0m \u001b[31m3\u001b[0m \u001b[31m3\u001b[0m \u001b[33m6\u001b[0m \u001b[33m6\u001b[0m \u001b[33m4\u001b[0m        : 40903.08 us\n",
      "\u001b[32m*** METAL    333\u001b[0m r_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m6\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m         mem  1.02 GB tm \u001b[33m    11.34ms\u001b[0m/   123.51ms (  1509.24 GFLOPS    2.3|201.1   GB/s) ['conv2d']\n",
      "*** METAL    334 r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  1.02 GB tm    142.54us/   123.66ms (    58.62 GFLOPS   82.9|239.1   GB/s) ['layernorm']\n",
      "*** METAL    335 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  1.02 GB tm    136.63us/   123.79ms (   109.13 GFLOPS   86.6|254.2   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m160\u001b[0m \u001b[34m16\u001b[0m \u001b[34m3\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m4\u001b[0m              :   331.75 us  <  beam2  : \u001b[34m20\u001b[0m \u001b[34m256\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m9\u001b[0m                  :   357.38 us\n",
      "\u001b[32m*** METAL    336\u001b[0m E_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                      mem  1.03 GB tm    283.04us/   124.08ms (    57.31 GFLOPS  125.3|166.7   GB/s) ['layernorm', '__add__']\n",
      "*** METAL    337 r_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m6\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m         mem  1.04 GB tm   4775.50us/   128.85ms (  3584.62 GFLOPS    5.4|477.6   GB/s) ['conv2d']\n",
      "*** METAL    338 r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  1.04 GB tm    154.87us/   129.01ms (    53.95 GFLOPS   76.3|220.0   GB/s) ['layernorm']\n",
      "*** METAL    339 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  1.04 GB tm    314.63us/   129.32ms (    47.39 GFLOPS   37.6|110.4   GB/s) ['layernorm']\n",
      "*** METAL    340 E_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m9\u001b[0m\u001b[90m\u001b[0m                        mem  1.05 GB tm    187.50us/   129.51ms (    48.93 GFLOPS  126.0|139.8   GB/s) ['layernorm', 'relu']\n",
      "*** METAL    341 r_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m6\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m         mem  1.06 GB tm   4785.83us/   134.29ms (  3576.88 GFLOPS    5.4|476.5   GB/s) ['conv2d']\n",
      "*** METAL    342 r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  1.06 GB tm    111.17us/   134.40ms (    75.17 GFLOPS  106.3|306.6   GB/s) ['layernorm']\n",
      "*** METAL    343 r_\u001b[34m2560\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m18\u001b[0m\u001b[90m\u001b[0m                           mem  1.06 GB tm     85.62us/   134.49ms (   174.12 GFLOPS  138.2|405.7   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m640\u001b[0m \u001b[34m12\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m                  :   201.21 us  <  beam2  : \u001b[34m160\u001b[0m \u001b[34m36\u001b[0m \u001b[36m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                  :   208.12 us\n",
      "\u001b[32m*** METAL    344\u001b[0m E_\u001b[34m640\u001b[0m\u001b[90m_\u001b[0m\u001b[34m12\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                          mem  1.08 GB tm    166.63us/   134.66ms (    41.30 GFLOPS  141.8|188.8   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m40\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m256\u001b[0m \u001b[33m3\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m        :  4139.25 us  <  beam2  : \u001b[34m40\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m3\u001b[0m \u001b[31m256\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m3\u001b[0m  :  5251.92 us\n",
      "\u001b[32m*** METAL    345\u001b[0m r_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m                mem  1.08 GB tm \u001b[33m    10.58ms\u001b[0m/   145.24ms (   507.46 GFLOPS    2.1|286.0   GB/s) ['conv2d']\n",
      "beam2  : \u001b[34m320\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m3\u001b[0m        :  1430.42 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[31m64\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m          :  1857.17 us\n",
      "\u001b[32m*** METAL    346\u001b[0m r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                mem  1.09 GB tm    710.21us/   145.95ms (  2834.75 GFLOPS   25.7|967.1   GB/s) ['conv2d']\n",
      "beam2  : \u001b[34m160\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[31m144\u001b[0m                   :    61.25 us  <  hc     : \u001b[34m80\u001b[0m \u001b[36m32\u001b[0m \u001b[31m36\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m                   :   108.92 us\n",
      "\u001b[32m*** METAL    347\u001b[0m r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                           mem  1.09 GB tm    104.92us/   146.05ms (    14.15 GFLOPS   56.6|56.6    GB/s) ['layernorm']\n",
      "*** METAL    348 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                           mem  1.09 GB tm     95.21us/   146.15ms (    15.60 GFLOPS   62.4|62.4    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m160\u001b[0m \u001b[36m16\u001b[0m \u001b[31m144\u001b[0m \u001b[33m4\u001b[0m                   :    88.67 us  <  hc     : \u001b[34m80\u001b[0m \u001b[36m32\u001b[0m \u001b[31m36\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m                   :   111.33 us\n",
      "\u001b[32m*** METAL    349\u001b[0m r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.09 GB tm     99.54us/   146.25ms (    44.95 GFLOPS   60.1|60.1    GB/s) ['layernorm']\n",
      "*** METAL    350 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.09 GB tm     94.54us/   146.34ms (    47.33 GFLOPS   63.3|63.3    GB/s) ['layernorm']\n",
      "hc     : \u001b[34m160\u001b[0m \u001b[34m32\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m3\u001b[0m                :    97.13 us  <  beam2  : \u001b[34m128\u001b[0m \u001b[34m20\u001b[0m \u001b[36m16\u001b[0m \u001b[33m9\u001b[0m \u001b[33m4\u001b[0m                  :   110.50 us\n",
      "\u001b[32m*** METAL    351\u001b[0m E_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                        mem  1.09 GB tm    517.50us/   146.86ms (     9.50 GFLOPS   23.0|30.4    GB/s) ['layernorm', 'relu']\n",
      "beam2  : \u001b[34m20\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m3\u001b[0m \u001b[35m3\u001b[0m \u001b[33m3\u001b[0m :  2607.46 us  <  hc     : \u001b[34m40\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m512\u001b[0m \u001b[33m3\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m        :  3568.29 us\n",
      "\u001b[32m*** METAL    352\u001b[0m r_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m         mem  1.10 GB tm   2962.04us/   149.82ms (  3807.32 GFLOPS    7.2|511.8   GB/s) ['conv2d']\n",
      "*** METAL    353 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                           mem  1.10 GB tm     43.25us/   149.86ms (    34.33 GFLOPS  137.3|137.3   GB/s) ['layernorm']\n",
      "*** METAL    354 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.10 GB tm     24.00us/   149.89ms (   186.45 GFLOPS  249.2|249.2   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m320\u001b[0m \u001b[34m8\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m9\u001b[0m                   :   144.79 us  <  hc     : \u001b[34m160\u001b[0m \u001b[34m32\u001b[0m \u001b[34m3\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m                :   261.17 us\n",
      "\u001b[32m*** METAL    355\u001b[0m E_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m9\u001b[0m\u001b[90m\u001b[0m                           mem  1.10 GB tm    210.21us/   150.10ms (    36.63 GFLOPS   85.0|96.6    GB/s) ['layernorm', '__add__']\n",
      "*** METAL    356 r_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m         mem  1.11 GB tm   2800.54us/   152.90ms (  4026.88 GFLOPS    7.6|541.3   GB/s) ['conv2d']\n",
      "*** METAL    357 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                           mem  1.11 GB tm     42.04us/   152.94ms (    35.32 GFLOPS  141.3|141.3   GB/s) ['layernorm']\n",
      "*** METAL    358 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.11 GB tm     24.00us/   152.96ms (   186.45 GFLOPS  249.2|249.2   GB/s) ['layernorm']\n",
      "*** METAL    359 E_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                        mem  1.12 GB tm     45.46us/   153.01ms (   108.13 GFLOPS  261.3|346.0   GB/s) ['layernorm', 'relu']\n",
      "*** METAL    360 r_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m         mem  1.12 GB tm   3313.96us/   156.32ms (  3403.01 GFLOPS    6.4|457.4   GB/s) ['conv2d']\n",
      "*** METAL    361 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                           mem  1.12 GB tm     39.88us/   156.36ms (    37.24 GFLOPS  148.9|148.9   GB/s) ['layernorm']\n",
      "*** METAL    362 r_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.12 GB tm     22.75us/   156.39ms (   196.70 GFLOPS  262.9|262.9   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m1280\u001b[0m \u001b[34m3\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m                  :    89.75 us  <  beam2  : \u001b[34m32\u001b[0m \u001b[34m9\u001b[0m \u001b[36m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[33m5\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m               :    92.96 us\n",
      "\u001b[32m*** METAL    363\u001b[0m E_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                          mem  1.13 GB tm    125.37us/   156.51ms (    27.44 GFLOPS   94.7|125.5   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m320\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m         :   241.96 us  <  hc     : \u001b[34m80\u001b[0m \u001b[36m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m3\u001b[0m \u001b[31m128\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   299.83 us\n",
      "\u001b[32m*** METAL    364\u001b[0m r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.13 GB tm     62.00us/   156.57ms (  2708.64 GFLOPS  102.1|528.5   GB/s) ['conv2d']\n",
      "hc     : \u001b[34m320\u001b[0m \u001b[36m32\u001b[0m \u001b[35m9\u001b[0m                       :    11.12 us  <  beam2  : \u001b[34m640\u001b[0m \u001b[36m16\u001b[0m \u001b[35m9\u001b[0m                       :    17.33 us\n",
      "\u001b[32m*** METAL    365\u001b[0m r_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m9\u001b[0m\u001b[90m\u001b[0m                               mem  1.13 GB tm      4.25us/   156.58ms (    19.28 GFLOPS   96.4|96.4    GB/s) ['softmax']\n",
      "hc     : \u001b[34m80\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m \u001b[35m9\u001b[0m                      :    11.83 us  <  beam2  : \u001b[34m640\u001b[0m \u001b[36m16\u001b[0m \u001b[31m9\u001b[0m                       :    12.17 us\n",
      "\u001b[32m*** METAL    366\u001b[0m r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m9\u001b[0m\u001b[90m\u001b[0m                              mem  1.13 GB tm      7.00us/   156.59ms (    52.66 GFLOPS   64.4|64.4    GB/s) ['softmax']\n",
      "hc     : \u001b[34m320\u001b[0m \u001b[36m32\u001b[0m \u001b[36m3\u001b[0m \u001b[33m3\u001b[0m                     :     6.54 us  <  beam2  : \u001b[34m160\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m9\u001b[0m                     :    13.67 us\n",
      "\u001b[32m*** METAL    367\u001b[0m E_\u001b[34m320\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.13 GB tm      8.38us/   156.59ms (    51.35 GFLOPS   97.8|117.4   GB/s) ['softmax']\n",
      "beam2  : \u001b[34m20\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m9\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m                :     7.25 us  <  hc     : \u001b[34m80\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m4\u001b[0m \u001b[35m9\u001b[0m                    :    11.87 us\n",
      "\u001b[32m*** METAL    368\u001b[0m r_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m9\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                        mem  1.13 GB tm      7.17us/   156.60ms (    51.44 GFLOPS   62.9|88.6    GB/s) ['__matmul__']\n",
      "beam2  : \u001b[34m40\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m8\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m             :    13.79 us  <  hc     : \u001b[34m10\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m               :    26.00 us  <  tc     : \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m8\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m5\u001b[0m \u001b[33m4\u001b[0m          :    26.37 us\n",
      "\u001b[32m*** METAL    369\u001b[0m r_\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                     mem  1.13 GB tm      5.38us/   156.61ms (   495.33 GFLOPS   33.6|274.3   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[34m40\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[35m8\u001b[0m          :    33.75 us  <  tc     : \u001b[34m8\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m8\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m5\u001b[0m \u001b[33m4\u001b[0m        :    56.29 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :    78.63 us\n",
      "\u001b[32m*** METAL    370\u001b[0m r_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m40\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m\u001b[0m                  mem  1.13 GB tm     13.50us/   156.62ms (  1553.45 GFLOPS   64.5|679.6   GB/s) ['linear']\n",
      "hc     : \u001b[34m2\u001b[0m \u001b[34m7\u001b[0m \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m3\u001b[0m                 :   106.75 us  <  beam2  : \u001b[34m2\u001b[0m \u001b[34m21\u001b[0m \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m                  :   112.38 us\n",
      "\u001b[32m*** METAL    371\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                         mem  1.13 GB tm     39.62us/   156.66ms (    34.73 GFLOPS  102.7|162.1   GB/s) ['stack', 'cat']\n",
      "beam2  : \u001b[34m4\u001b[0m \u001b[34m21\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m       :   943.38 us  <  hc     : \u001b[34m8\u001b[0m \u001b[34m7\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m3\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m         :  2939.46 us\n",
      "\u001b[32m*** METAL    372\u001b[0m r_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.14 GB tm    824.25us/   157.48ms (   859.90 GFLOPS   11.3|330.6   GB/s) ['linear', '__add__', 'dropout']\n",
      "scheduled 22 kernels\n",
      "beam2  : \u001b[34m3\u001b[0m \u001b[34m7\u001b[0m \u001b[36m3\u001b[0m \u001b[33m7\u001b[0m                        :     7.38 us  <  hc     : \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m3\u001b[0m \u001b[36m3\u001b[0m                        :     7.42 us\n",
      "\u001b[32m*** METAL    373\u001b[0m E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.14 GB tm      4.08us/   157.49ms (     0.11 GFLOPS    0.5|0.5     GB/s) ['__eq__']\n",
      "beam2  : \u001b[34m1344\u001b[0m \u001b[32m32\u001b[0m \u001b[31m16\u001b[0m                     :    31.75 us  <  hc     : \u001b[34m1344\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m                     :    44.25 us\n",
      "\u001b[32m*** METAL    374\u001b[0m r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.14 GB tm     30.37us/   157.52ms (    69.38 GFLOPS   90.8|283.2   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   375\u001b[0m custom_random                            mem  1.14 GB \n",
      "\u001b[32m*** CUSTOM   376\u001b[0m custom_random                            mem  1.14 GB \n",
      "\u001b[32m*** CUSTOM   377\u001b[0m custom_random                            mem  1.14 GB \n",
      "\u001b[32m*** CUSTOM   378\u001b[0m custom_random                            mem  1.14 GB \n",
      "\u001b[32m*** CUSTOM   379\u001b[0m custom_random                            mem  1.14 GB \n",
      "\u001b[32m*** CUSTOM   380\u001b[0m custom_random                            mem  1.14 GB \n",
      "beam2  : \u001b[34m672\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m \u001b[33m2\u001b[0m                    :    34.46 us  <  hc     : \u001b[34m1344\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m                     :    35.12 us\n",
      "\u001b[32m*** METAL    381\u001b[0m r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.14 GB tm     39.29us/   157.56ms (    64.03 GFLOPS   70.3|111.6   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m6144\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[33m4\u001b[0m                    :    51.12 us  <  hc     : \u001b[34m8192\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    68.25 us\n",
      "\u001b[32m*** METAL    382\u001b[0m E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.15 GB tm     54.50us/   157.61ms (    28.86 GFLOPS  115.4|115.4   GB/s) ['uniform']\n",
      "beam2  : \u001b[34m48\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m                         :     7.21 us  <  hc     : \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                        :     7.29 us\n",
      "\u001b[32m*** METAL    383\u001b[0m E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.14 GB tm      8.12us/   157.62ms (     0.38 GFLOPS    1.5|1.5     GB/s) ['uniform']\n",
      "hc     : \u001b[34m2048\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    35.88 us  <  beam2  : \u001b[34m256\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m                 :    36.75 us\n",
      "\u001b[32m*** METAL    384\u001b[0m E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.15 GB tm     37.25us/   157.66ms (    70.37 GFLOPS   84.4|84.4    GB/s) ['normal']\n",
      "*** METAL    385 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.14 GB tm      6.71us/   157.66ms (     0.15 GFLOPS    0.6|0.6     GB/s) ['uniform']\n",
      "hc     : \u001b[34m168\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                   :    43.33 us  <  beam2  : \u001b[34m42\u001b[0m \u001b[34m32\u001b[0m \u001b[36m16\u001b[0m \u001b[36m32\u001b[0m                    :    49.50 us\n",
      "\u001b[32m*** METAL    386\u001b[0m E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.15 GB tm    216.75us/   157.88ms (     7.14 GFLOPS   25.4|31.7    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m42\u001b[0m \u001b[34m96\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m       :   760.21 us  <  tc     : \u001b[34m42\u001b[0m \u001b[34m12\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :  1181.58 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m3\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m           :  4712.04 us\n",
      "\u001b[32m*** METAL    387\u001b[0m r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.15 GB tm   2524.71us/   160.41ms (   838.11 GFLOPS    5.6|161.1   GB/s) ['linear']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m3\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m8\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m3\u001b[0m         :   110.46 us  <  hc     : \u001b[34m2\u001b[0m \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m          :   786.21 us\n",
      "\u001b[32m*** METAL    388\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.16 GB tm    129.75us/   160.54ms (   297.75 GFLOPS   70.6|108.3   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "hc     : \u001b[34m336\u001b[0m \u001b[36m32\u001b[0m \u001b[35m21\u001b[0m                      :    16.83 us  <  beam2  : \u001b[34m42\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[31m21\u001b[0m \u001b[33m2\u001b[0m                 :    18.29 us\n",
      "\u001b[32m*** METAL    389\u001b[0m r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.16 GB tm     19.62us/   160.55ms (    10.96 GFLOPS   48.2|48.2    GB/s) ['softmax']\n",
      "beam2  : \u001b[34m672\u001b[0m \u001b[36m16\u001b[0m \u001b[31m3\u001b[0m \u001b[35m7\u001b[0m                     :    16.13 us  <  hc     : \u001b[34m112\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m \u001b[35m21\u001b[0m                    :    16.25 us\n",
      "\u001b[32m*** METAL    390\u001b[0m r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.16 GB tm     74.00us/   160.63ms (    12.35 GFLOPS   13.4|13.4    GB/s) ['softmax']\n",
      "hc     : \u001b[34m336\u001b[0m \u001b[34m7\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                     :    74.00 us  <  beam2  : \u001b[34m168\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m7\u001b[0m \u001b[33m4\u001b[0m                   :   120.29 us\n",
      "\u001b[32m*** METAL    391\u001b[0m E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.16 GB tm     74.96us/   160.70ms (    26.11 GFLOPS   37.3|44.2    GB/s) ['softmax', 'dropout']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[34m3\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m         :   152.37 us  <  hc     : \u001b[34m64\u001b[0m \u001b[34m7\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m21\u001b[0m               :   548.58 us\n",
      "\u001b[32m*** METAL    392\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.16 GB tm    156.08us/   160.86ms (   314.91 GFLOPS   76.3|141.1   GB/s) ['__matmul__']\n",
      "hc     : \u001b[34m64\u001b[0m \u001b[34m21\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                   :    42.88 us  <  beam2  : \u001b[34m16\u001b[0m \u001b[34m7\u001b[0m \u001b[34m4\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m8\u001b[0m \u001b[33m4\u001b[0m                :    45.58 us\n",
      "\u001b[32m*** METAL    393\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.16 GB tm    228.33us/   161.09ms (     0.00 GFLOPS   24.1|24.1    GB/s) \n",
      "beam2  : \u001b[34m168\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m          :   518.87 us  <  tc     : \u001b[34m42\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :   868.58 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m            :  2192.50 us\n",
      "\u001b[32m*** METAL    394\u001b[0m r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.16 GB tm   1198.25us/   162.29ms (   591.51 GFLOPS   10.1|303.2   GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    395 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.16 GB tm    211.17us/   162.50ms (     9.98 GFLOPS   13.1|40.7    GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   396\u001b[0m custom_random                            mem  1.17 GB \n",
      "\u001b[32m*** CUSTOM   397\u001b[0m custom_random                            mem  1.17 GB \n",
      "\u001b[32m*** CUSTOM   398\u001b[0m custom_random                            mem  1.17 GB \n",
      "\u001b[32m*** CUSTOM   399\u001b[0m custom_random                            mem  1.17 GB \n",
      "*** METAL    400 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.17 GB tm     39.42us/   162.54ms (    63.83 GFLOPS   70.1|111.3   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m8192\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    85.33 us  <  beam2  : \u001b[34m4096\u001b[0m \u001b[36m16\u001b[0m \u001b[36m16\u001b[0m                     :    88.00 us\n",
      "\u001b[32m*** METAL    401\u001b[0m E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.18 GB tm     70.12us/   162.61ms (    29.91 GFLOPS  119.6|119.6   GB/s) ['uniform']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m16\u001b[0m \u001b[33m2\u001b[0m                        :     7.38 us  <  hc     : \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.38 us\n",
      "\u001b[32m*** METAL    402\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.17 GB tm      8.08us/   162.62ms (     0.51 GFLOPS    2.0|2.0     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m2048\u001b[0m \u001b[36m32\u001b[0m \u001b[36m8\u001b[0m \u001b[33m2\u001b[0m                    :    69.46 us  <  hc     : \u001b[34m8192\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    83.79 us\n",
      "\u001b[32m*** METAL    403\u001b[0m E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.18 GB tm     61.13us/   162.68ms (    34.31 GFLOPS  137.2|137.2   GB/s) ['uniform']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m4\u001b[0m \u001b[36m4\u001b[0m                         :     7.13 us  <  hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     7.25 us\n",
      "\u001b[32m*** METAL    404\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.17 GB tm      8.00us/   162.68ms (     0.13 GFLOPS    0.5|0.5     GB/s) ['uniform']\n",
      "*** METAL    405 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.18 GB tm     70.88us/   162.76ms (    21.85 GFLOPS   77.8|97.1    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m84\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m       :  1049.00 us  <  tc     : \u001b[34m42\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :  1513.00 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m           :  3167.13 us\n",
      "\u001b[32m*** METAL    406\u001b[0m r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.19 GB tm   3557.46us/   166.31ms (   801.58 GFLOPS    5.0|153.2   GB/s) ['linear', 'gelu']\n",
      "beam2  : \u001b[34m42\u001b[0m \u001b[34m32\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m256\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m      :  1436.50 us  <  tc     : \u001b[34m42\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m256\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :  1710.08 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m512\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m            :  2729.42 us\n",
      "\u001b[32m*** METAL    407\u001b[0m r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.19 GB tm   2845.75us/   169.16ms (   990.69 GFLOPS    6.3|186.9   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   408\u001b[0m custom_random                            mem  1.19 GB \n",
      "hc     : \u001b[34m7168\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    92.79 us  <  beam2  : \u001b[34m1344\u001b[0m \u001b[36m32\u001b[0m \u001b[36m16\u001b[0m                     :   105.75 us\n",
      "\u001b[32m*** METAL    409\u001b[0m E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.20 GB tm    113.96us/   169.27ms (    30.19 GFLOPS   96.6|96.6    GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    410 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.20 GB tm      7.25us/   169.28ms (     0.06 GFLOPS    0.3|0.3     GB/s) ['__eq__']\n",
      "*** METAL    411 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.20 GB tm     31.67us/   169.31ms (    66.55 GFLOPS   87.1|271.6   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   412\u001b[0m custom_random                            mem  1.20 GB \n",
      "\u001b[32m*** CUSTOM   413\u001b[0m custom_random                            mem  1.20 GB \n",
      "\u001b[32m*** CUSTOM   414\u001b[0m custom_random                            mem  1.20 GB \n",
      "\u001b[32m*** CUSTOM   415\u001b[0m custom_random                            mem  1.20 GB \n",
      "\u001b[32m*** CUSTOM   416\u001b[0m custom_random                            mem  1.20 GB \n",
      "\u001b[32m*** CUSTOM   417\u001b[0m custom_random                            mem  1.20 GB \n",
      "*** METAL    418 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.20 GB tm     41.50us/   169.35ms (    60.63 GFLOPS   66.6|105.7   GB/s) ['layernorm']\n",
      "*** METAL    419 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.21 GB tm     61.29us/   169.41ms (    25.66 GFLOPS  102.6|102.6   GB/s) ['uniform']\n",
      "*** METAL    420 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.20 GB tm      6.83us/   169.42ms (     0.45 GFLOPS    1.8|1.8     GB/s) ['uniform']\n",
      "*** METAL    421 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.21 GB tm     40.67us/   169.46ms (    64.46 GFLOPS   77.4|77.4    GB/s) ['normal']\n",
      "*** METAL    422 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.20 GB tm    562.71us/   170.02ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['uniform']\n",
      "*** METAL    423 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.21 GB tm     50.58us/   170.07ms (    30.61 GFLOPS  109.0|136.0   GB/s) ['layernorm']\n",
      "*** METAL    424 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.21 GB tm   2315.71us/   172.39ms (   913.76 GFLOPS    6.1|175.6   GB/s) ['linear']\n",
      "*** METAL    425 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.22 GB tm    127.92us/   172.52ms (   302.02 GFLOPS   71.6|109.9   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    426 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.22 GB tm    149.63us/   172.67ms (     1.44 GFLOPS    6.3|6.3     GB/s) ['softmax']\n",
      "*** METAL    427 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.22 GB tm     19.17us/   172.69ms (    47.68 GFLOPS   51.6|51.6    GB/s) ['softmax']\n",
      "*** METAL    428 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.22 GB tm     78.25us/   172.77ms (    25.01 GFLOPS   35.7|42.3    GB/s) ['softmax', 'dropout']\n",
      "*** METAL    429 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.22 GB tm    152.38us/   172.92ms (   322.57 GFLOPS   78.2|144.5   GB/s) ['__matmul__']\n",
      "*** METAL    430 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.22 GB tm     41.58us/   172.96ms (     0.00 GFLOPS  132.4|132.4   GB/s) \n",
      "*** METAL    431 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.22 GB tm   1193.37us/   174.15ms (   593.92 GFLOPS   10.1|304.5   GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    432 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.22 GB tm     33.08us/   174.19ms (    63.70 GFLOPS   83.4|260.0   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   433\u001b[0m custom_random                            mem  1.23 GB \n",
      "\u001b[32m*** CUSTOM   434\u001b[0m custom_random                            mem  1.23 GB \n",
      "\u001b[32m*** CUSTOM   435\u001b[0m custom_random                            mem  1.23 GB \n",
      "\u001b[32m*** CUSTOM   436\u001b[0m custom_random                            mem  1.23 GB \n",
      "*** METAL    437 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.23 GB tm     28.00us/   174.21ms (    89.86 GFLOPS   98.7|156.7   GB/s) ['layernorm']\n",
      "*** METAL    438 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.24 GB tm    387.63us/   174.60ms (     5.41 GFLOPS   21.6|21.6    GB/s) ['uniform']\n",
      "*** METAL    439 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.23 GB tm      2.37us/   174.60ms (     1.72 GFLOPS    6.9|6.9     GB/s) ['uniform']\n",
      "*** METAL    440 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.24 GB tm     58.50us/   174.66ms (    35.85 GFLOPS  143.4|143.4   GB/s) ['uniform']\n",
      "*** METAL    441 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.23 GB tm      2.25us/   174.66ms (     0.46 GFLOPS    1.8|1.8     GB/s) ['uniform']\n",
      "*** METAL    442 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.24 GB tm     41.12us/   174.71ms (    37.65 GFLOPS  134.1|167.3   GB/s) ['layernorm']\n",
      "*** METAL    443 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.25 GB tm    692.46us/   175.40ms (  4118.09 GFLOPS   25.9|787.0   GB/s) ['linear', 'gelu']\n",
      "*** METAL    444 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.25 GB tm    795.92us/   176.19ms (  3542.16 GFLOPS   22.6|668.3   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   445\u001b[0m custom_random                            mem  1.25 GB \n",
      "*** METAL    446 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.25 GB tm    100.12us/   176.29ms (    34.36 GFLOPS  110.0|110.0   GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    447 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.25 GB tm      2.29us/   176.30ms (     0.19 GFLOPS    1.0|1.0     GB/s) ['__eq__']\n",
      "*** METAL    448 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.25 GB tm     27.50us/   176.32ms (    76.63 GFLOPS  100.3|312.8   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   449\u001b[0m custom_random                            mem  1.26 GB \n",
      "\u001b[32m*** CUSTOM   450\u001b[0m custom_random                            mem  1.26 GB \n",
      "\u001b[32m*** CUSTOM   451\u001b[0m custom_random                            mem  1.26 GB \n",
      "\u001b[32m*** CUSTOM   452\u001b[0m custom_random                            mem  1.26 GB \n",
      "\u001b[32m*** CUSTOM   453\u001b[0m custom_random                            mem  1.26 GB \n",
      "\u001b[32m*** CUSTOM   454\u001b[0m custom_random                            mem  1.26 GB \n",
      "*** METAL    455 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.26 GB tm    148.42us/   176.47ms (    16.95 GFLOPS   18.6|29.6    GB/s) ['layernorm']\n",
      "*** METAL    456 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.27 GB tm     45.42us/   176.52ms (    34.63 GFLOPS  138.5|138.5   GB/s) ['uniform']\n",
      "*** METAL    457 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.26 GB tm      2.71us/   176.52ms (     1.13 GFLOPS    4.5|4.5     GB/s) ['uniform']\n",
      "*** METAL    458 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.26 GB tm     24.42us/   176.55ms (   107.36 GFLOPS  128.8|128.8   GB/s) ['normal']\n",
      "*** METAL    459 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.26 GB tm      2.67us/   176.55ms (     0.38 GFLOPS    1.5|1.5     GB/s) ['uniform']\n",
      "*** METAL    460 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.27 GB tm     41.33us/   176.59ms (    37.46 GFLOPS  133.4|166.5   GB/s) ['layernorm']\n",
      "*** METAL    461 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.27 GB tm    509.12us/   177.10ms (  4156.14 GFLOPS   27.8|798.8   GB/s) ['linear']\n",
      "*** METAL    462 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.27 GB tm    427.00us/   177.53ms (    90.48 GFLOPS   21.4|32.9    GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    463 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.27 GB tm      5.00us/   177.53ms (    43.01 GFLOPS  189.2|189.2   GB/s) ['softmax']\n",
      "*** METAL    464 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.27 GB tm      5.25us/   177.54ms (   174.08 GFLOPS  188.4|188.4   GB/s) ['softmax']\n",
      "*** METAL    465 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.28 GB tm     18.75us/   177.55ms (   104.37 GFLOPS  149.1|176.6   GB/s) ['softmax', 'dropout']\n",
      "*** METAL    466 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.28 GB tm     46.33us/   177.60ms (  1060.84 GFLOPS  257.1|475.3   GB/s) ['__matmul__']\n",
      "*** METAL    467 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.28 GB tm     14.37us/   177.61ms (     0.00 GFLOPS  383.0|383.0   GB/s) \n",
      "*** METAL    468 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.28 GB tm    461.00us/   178.08ms (  1537.47 GFLOPS   26.2|788.1   GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    469 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.28 GB tm     53.00us/   178.13ms (    39.76 GFLOPS   52.0|162.3   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   470\u001b[0m custom_random                            mem  1.29 GB \n",
      "\u001b[32m*** CUSTOM   471\u001b[0m custom_random                            mem  1.29 GB \n",
      "\u001b[32m*** CUSTOM   472\u001b[0m custom_random                            mem  1.29 GB \n",
      "\u001b[32m*** CUSTOM   473\u001b[0m custom_random                            mem  1.29 GB \n",
      "*** METAL    474 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.29 GB tm     54.25us/   178.18ms (    46.38 GFLOPS   50.9|80.9    GB/s) ['layernorm']\n",
      "*** METAL    475 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.30 GB tm     85.92us/   178.27ms (    24.41 GFLOPS   97.6|97.6    GB/s) ['uniform']\n",
      "*** METAL    476 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.29 GB tm      2.71us/   178.27ms (     1.51 GFLOPS    6.0|6.0     GB/s) ['uniform']\n",
      "*** METAL    477 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.30 GB tm     56.00us/   178.33ms (    37.45 GFLOPS  149.8|149.8   GB/s) ['uniform']\n",
      "*** METAL    478 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.29 GB tm      2.42us/   178.33ms (     0.42 GFLOPS    1.7|1.7     GB/s) ['uniform']\n",
      "*** METAL    479 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.29 GB tm     44.50us/   178.37ms (    34.79 GFLOPS  124.0|154.6   GB/s) ['layernorm']\n",
      "*** METAL    480 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.31 GB tm   2473.00us/   180.85ms (  1153.09 GFLOPS    7.3|220.4   GB/s) ['linear', 'gelu']\n",
      "*** METAL    481 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.31 GB tm   3677.25us/   184.52ms (   766.68 GFLOPS    4.9|144.7   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   482\u001b[0m custom_random                            mem  1.31 GB \n",
      "*** METAL    483 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.31 GB tm    112.00us/   184.64ms (    30.72 GFLOPS   98.3|98.3    GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    484 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.31 GB tm    384.37us/   185.02ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__']\n",
      "*** METAL    485 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.31 GB tm     32.50us/   185.05ms (    64.84 GFLOPS   84.9|264.7   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   486\u001b[0m custom_random                            mem  1.32 GB \n",
      "\u001b[32m*** CUSTOM   487\u001b[0m custom_random                            mem  1.32 GB \n",
      "\u001b[32m*** CUSTOM   488\u001b[0m custom_random                            mem  1.32 GB \n",
      "\u001b[32m*** CUSTOM   489\u001b[0m custom_random                            mem  1.32 GB \n",
      "\u001b[32m*** CUSTOM   490\u001b[0m custom_random                            mem  1.32 GB \n",
      "\u001b[32m*** CUSTOM   491\u001b[0m custom_random                            mem  1.32 GB \n",
      "*** METAL    492 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.32 GB tm     35.79us/   185.09ms (    70.29 GFLOPS   77.2|122.6   GB/s) ['layernorm']\n",
      "*** METAL    493 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.33 GB tm    140.87us/   185.23ms (    11.16 GFLOPS   44.7|44.7    GB/s) ['uniform']\n",
      "*** METAL    494 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.32 GB tm      5.12us/   185.24ms (     0.60 GFLOPS    2.4|2.4     GB/s) ['uniform']\n",
      "*** METAL    495 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.32 GB tm     50.50us/   185.29ms (    51.91 GFLOPS   62.3|62.3    GB/s) ['normal']\n",
      "*** METAL    496 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.32 GB tm      4.12us/   185.29ms (     0.25 GFLOPS    1.0|1.0     GB/s) ['uniform']\n",
      "*** METAL    497 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.32 GB tm     47.75us/   185.34ms (    32.42 GFLOPS  115.5|144.1   GB/s) ['layernorm']\n",
      "*** METAL    498 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.33 GB tm   1515.17us/   186.85ms (  1396.54 GFLOPS    9.3|268.4   GB/s) ['linear']\n",
      "*** METAL    499 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.33 GB tm     68.50us/   186.92ms (   563.99 GFLOPS  133.7|205.2   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    500 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.33 GB tm     10.71us/   186.93ms (    20.08 GFLOPS   88.4|88.4    GB/s) ['softmax']\n",
      "*** METAL    501 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.33 GB tm    168.29us/   187.10ms (     5.43 GFLOPS    5.9|5.9     GB/s) ['softmax']\n",
      "*** METAL    502 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.33 GB tm    117.46us/   187.22ms (    16.66 GFLOPS   23.8|28.2    GB/s) ['softmax', 'dropout']\n",
      "*** METAL    503 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.34 GB tm     84.75us/   187.30ms (   579.96 GFLOPS  140.6|259.8   GB/s) ['__matmul__']\n",
      "*** METAL    504 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.34 GB tm     23.83us/   187.33ms (     0.00 GFLOPS  231.0|231.0   GB/s) \n",
      "*** METAL    505 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.34 GB tm    630.33us/   187.96ms (  1124.44 GFLOPS   19.1|576.4   GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    506 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.34 GB tm     31.50us/   187.99ms (    66.90 GFLOPS   87.6|273.1   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   507\u001b[0m custom_random                            mem  1.35 GB \n",
      "\u001b[32m*** CUSTOM   508\u001b[0m custom_random                            mem  1.35 GB \n",
      "\u001b[32m*** CUSTOM   509\u001b[0m custom_random                            mem  1.35 GB \n",
      "\u001b[32m*** CUSTOM   510\u001b[0m custom_random                            mem  1.35 GB \n",
      "*** METAL    511 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.35 GB tm    112.08us/   188.10ms (    22.45 GFLOPS   24.7|39.1    GB/s) ['layernorm']\n",
      "*** METAL    512 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.36 GB tm     86.75us/   188.19ms (    24.17 GFLOPS   96.7|96.7    GB/s) ['uniform']\n",
      "*** METAL    513 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.35 GB tm      3.54us/   188.19ms (     1.16 GFLOPS    4.6|4.6     GB/s) ['uniform']\n",
      "*** METAL    514 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.36 GB tm     58.08us/   188.25ms (    36.11 GFLOPS  144.4|144.4   GB/s) ['uniform']\n",
      "*** METAL    515 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.35 GB tm      3.62us/   188.25ms (     0.28 GFLOPS    1.1|1.1     GB/s) ['uniform']\n",
      "*** METAL    516 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.35 GB tm     48.38us/   188.30ms (    32.01 GFLOPS  114.0|142.2   GB/s) ['layernorm']\n",
      "*** METAL    517 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.37 GB tm   1031.04us/   189.33ms (  2765.75 GFLOPS   17.4|528.6   GB/s) ['linear', 'gelu']\n",
      "*** METAL    518 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.37 GB tm   1207.75us/   190.54ms (  2334.31 GFLOPS   14.9|440.4   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   519\u001b[0m custom_random                            mem  1.37 GB \n",
      "*** METAL    520 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.37 GB tm    192.50us/   190.73ms (    17.87 GFLOPS   57.2|57.2    GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    521 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.37 GB tm      3.17us/   190.74ms (     0.14 GFLOPS    0.7|0.7     GB/s) ['__eq__']\n",
      "*** METAL    522 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.37 GB tm     42.00us/   190.78ms (    50.18 GFLOPS   65.7|204.8   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   523\u001b[0m custom_random                            mem  1.38 GB \n",
      "\u001b[32m*** CUSTOM   524\u001b[0m custom_random                            mem  1.38 GB \n",
      "\u001b[32m*** CUSTOM   525\u001b[0m custom_random                            mem  1.38 GB \n",
      "\u001b[32m*** CUSTOM   526\u001b[0m custom_random                            mem  1.38 GB \n",
      "\u001b[32m*** CUSTOM   527\u001b[0m custom_random                            mem  1.38 GB \n",
      "\u001b[32m*** CUSTOM   528\u001b[0m custom_random                            mem  1.38 GB \n",
      "*** METAL    529 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.38 GB tm     36.25us/   190.81ms (    69.41 GFLOPS   76.2|121.0   GB/s) ['layernorm']\n",
      "*** METAL    530 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.39 GB tm     47.50us/   190.86ms (    33.11 GFLOPS  132.5|132.5   GB/s) ['uniform']\n",
      "*** METAL    531 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.38 GB tm      3.37us/   190.86ms (     0.91 GFLOPS    3.6|3.6     GB/s) ['uniform']\n",
      "*** METAL    532 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.38 GB tm     31.00us/   190.90ms (    84.56 GFLOPS  101.5|101.5   GB/s) ['normal']\n",
      "*** METAL    533 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.38 GB tm      2.96us/   190.90ms (     0.35 GFLOPS    1.4|1.4     GB/s) ['uniform']\n",
      "*** METAL    534 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.38 GB tm     39.12us/   190.94ms (    39.57 GFLOPS  141.0|175.9   GB/s) ['layernorm']\n",
      "*** METAL    535 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.39 GB tm    591.79us/   191.53ms (  3575.57 GFLOPS   23.9|687.2   GB/s) ['linear']\n",
      "*** METAL    536 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.39 GB tm     55.00us/   191.58ms (   702.43 GFLOPS  166.5|255.6   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    537 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.39 GB tm      5.75us/   191.59ms (    37.40 GFLOPS  164.6|164.6   GB/s) ['softmax']\n",
      "*** METAL    538 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.39 GB tm      5.46us/   191.60ms (   167.43 GFLOPS  181.2|181.2   GB/s) ['softmax']\n",
      "*** METAL    539 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.39 GB tm     28.87us/   191.62ms (    67.77 GFLOPS   96.8|114.7   GB/s) ['softmax', 'dropout']\n",
      "*** METAL    540 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.40 GB tm     49.04us/   191.67ms (  1002.25 GFLOPS  242.9|449.0   GB/s) ['__matmul__']\n",
      "*** METAL    541 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.40 GB tm     14.83us/   191.69ms (     0.00 GFLOPS  371.1|371.1   GB/s) \n",
      "*** METAL    542 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.40 GB tm    342.00us/   192.03ms (  2072.43 GFLOPS   35.3|1062.4  GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    543 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.40 GB tm     11.63us/   192.04ms (   181.28 GFLOPS  237.2|739.9   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   544\u001b[0m custom_random                            mem  1.41 GB \n",
      "\u001b[32m*** CUSTOM   545\u001b[0m custom_random                            mem  1.41 GB \n",
      "\u001b[32m*** CUSTOM   546\u001b[0m custom_random                            mem  1.41 GB \n",
      "\u001b[32m*** CUSTOM   547\u001b[0m custom_random                            mem  1.41 GB \n",
      "*** METAL    548 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.41 GB tm     42.63us/   192.08ms (    59.03 GFLOPS   64.8|102.9   GB/s) ['layernorm']\n",
      "*** METAL    549 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.42 GB tm     56.75us/   192.14ms (    36.95 GFLOPS  147.8|147.8   GB/s) ['uniform']\n",
      "*** METAL    550 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.41 GB tm      2.50us/   192.14ms (     1.64 GFLOPS    6.6|6.6     GB/s) ['uniform']\n",
      "*** METAL    551 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.42 GB tm     54.46us/   192.20ms (    38.51 GFLOPS  154.0|154.0   GB/s) ['uniform']\n",
      "*** METAL    552 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.41 GB tm      2.17us/   192.20ms (     0.47 GFLOPS    1.9|1.9     GB/s) ['uniform']\n",
      "*** METAL    553 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.41 GB tm     40.13us/   192.24ms (    38.59 GFLOPS  137.5|171.5   GB/s) ['layernorm']\n",
      "*** METAL    554 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.42 GB tm    729.00us/   192.97ms (  3911.66 GFLOPS   24.6|747.6   GB/s) ['linear', 'gelu']\n",
      "*** METAL    555 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.43 GB tm    723.62us/   193.69ms (  3896.02 GFLOPS   24.8|735.1   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   556\u001b[0m custom_random                            mem  1.43 GB \n",
      "*** METAL    557 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.43 GB tm    108.79us/   193.80ms (    31.63 GFLOPS  101.2|101.2   GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    558 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.43 GB tm      3.13us/   193.81ms (     0.14 GFLOPS    0.7|0.7     GB/s) ['__eq__']\n",
      "*** METAL    559 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.43 GB tm     13.71us/   193.82ms (   153.73 GFLOPS  201.2|627.5   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   560\u001b[0m custom_random                            mem  1.44 GB \n",
      "\u001b[32m*** CUSTOM   561\u001b[0m custom_random                            mem  1.44 GB \n",
      "\u001b[32m*** CUSTOM   562\u001b[0m custom_random                            mem  1.44 GB \n",
      "\u001b[32m*** CUSTOM   563\u001b[0m custom_random                            mem  1.44 GB \n",
      "\u001b[32m*** CUSTOM   564\u001b[0m custom_random                            mem  1.44 GB \n",
      "\u001b[32m*** CUSTOM   565\u001b[0m custom_random                            mem  1.44 GB \n",
      "*** METAL    566 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.44 GB tm     46.46us/   193.87ms (    54.16 GFLOPS   59.5|94.4    GB/s) ['layernorm']\n",
      "*** METAL    567 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.45 GB tm     60.87us/   193.93ms (    25.84 GFLOPS  103.4|103.4   GB/s) ['uniform']\n",
      "*** METAL    568 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.44 GB tm      7.71us/   193.93ms (     0.40 GFLOPS    1.6|1.6     GB/s) ['uniform']\n",
      "*** METAL    569 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.44 GB tm     40.96us/   193.97ms (    64.00 GFLOPS   76.8|76.8    GB/s) ['normal']\n",
      "*** METAL    570 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.44 GB tm      4.21us/   193.98ms (     0.24 GFLOPS    1.0|1.0     GB/s) ['uniform']\n",
      "*** METAL    571 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.44 GB tm     45.00us/   194.02ms (    34.41 GFLOPS  122.6|152.9   GB/s) ['layernorm']\n",
      "*** METAL    572 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.45 GB tm   1129.25us/   195.15ms (  1873.80 GFLOPS   12.5|360.1   GB/s) ['linear']\n",
      "*** METAL    573 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.45 GB tm     65.96us/   195.22ms (   585.73 GFLOPS  138.9|213.1   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    574 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.45 GB tm      8.00us/   195.23ms (    26.88 GFLOPS  118.3|118.3   GB/s) ['softmax']\n",
      "*** METAL    575 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.45 GB tm      8.37us/   195.24ms (   109.13 GFLOPS  118.1|118.1   GB/s) ['softmax']\n",
      "*** METAL    576 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.45 GB tm     40.04us/   195.28ms (    48.87 GFLOPS   69.8|82.7    GB/s) ['softmax', 'dropout']\n",
      "*** METAL    577 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.46 GB tm     84.17us/   195.36ms (   583.98 GFLOPS  141.5|261.6   GB/s) ['__matmul__']\n",
      "*** METAL    578 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.46 GB tm     23.04us/   195.38ms (     0.00 GFLOPS  238.9|238.9   GB/s) \n",
      "*** METAL    579 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.46 GB tm    654.38us/   196.04ms (  1083.13 GFLOPS   18.4|555.2   GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    580 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.46 GB tm     36.62us/   196.07ms (    57.54 GFLOPS   75.3|234.9   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   581\u001b[0m custom_random                            mem  1.47 GB \n",
      "\u001b[32m*** CUSTOM   582\u001b[0m custom_random                            mem  1.47 GB \n",
      "\u001b[32m*** CUSTOM   583\u001b[0m custom_random                            mem  1.47 GB \n",
      "\u001b[32m*** CUSTOM   584\u001b[0m custom_random                            mem  1.47 GB \n",
      "*** METAL    585 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.47 GB tm     38.71us/   196.11ms (    65.00 GFLOPS   71.4|113.3   GB/s) ['layernorm']\n",
      "*** METAL    586 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.47 GB tm     63.21us/   196.18ms (    33.18 GFLOPS  132.7|132.7   GB/s) ['uniform']\n",
      "*** METAL    587 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.47 GB tm      4.79us/   196.18ms (     0.85 GFLOPS    3.4|3.4     GB/s) ['uniform']\n",
      "*** METAL    588 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.47 GB tm     62.67us/   196.24ms (    33.47 GFLOPS  133.9|133.9   GB/s) ['uniform']\n",
      "*** METAL    589 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.47 GB tm      6.87us/   196.25ms (     0.15 GFLOPS    0.6|0.6     GB/s) ['uniform']\n",
      "*** METAL    590 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.47 GB tm     53.46us/   196.30ms (    28.96 GFLOPS  103.2|128.7   GB/s) ['layernorm']\n",
      "*** METAL    591 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.48 GB tm   2172.17us/   198.48ms (  1312.79 GFLOPS    8.3|250.9   GB/s) ['linear', 'gelu']\n",
      "*** METAL    592 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.49 GB tm   1527.46us/   200.00ms (  1845.72 GFLOPS   11.8|348.2   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   593\u001b[0m custom_random                            mem  1.49 GB \n",
      "*** METAL    594 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.49 GB tm    104.75us/   200.11ms (    32.85 GFLOPS  105.1|105.1   GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    595 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.49 GB tm    549.83us/   200.66ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__eq__']\n",
      "*** METAL    596 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.49 GB tm     31.50us/   200.69ms (    66.90 GFLOPS   87.6|273.1   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   597\u001b[0m custom_random                            mem  1.50 GB \n",
      "\u001b[32m*** CUSTOM   598\u001b[0m custom_random                            mem  1.50 GB \n",
      "\u001b[32m*** CUSTOM   599\u001b[0m custom_random                            mem  1.50 GB \n",
      "\u001b[32m*** CUSTOM   600\u001b[0m custom_random                            mem  1.50 GB \n",
      "\u001b[32m*** CUSTOM   601\u001b[0m custom_random                            mem  1.50 GB \n",
      "\u001b[32m*** CUSTOM   602\u001b[0m custom_random                            mem  1.50 GB \n",
      "*** METAL    603 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.50 GB tm     76.25us/   200.77ms (    33.00 GFLOPS   36.2|57.5    GB/s) ['layernorm']\n",
      "*** METAL    604 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.50 GB tm     49.08us/   200.81ms (    32.04 GFLOPS  128.2|128.2   GB/s) ['uniform']\n",
      "*** METAL    605 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.50 GB tm      6.04us/   200.82ms (     0.51 GFLOPS    2.0|2.0     GB/s) ['uniform']\n",
      "*** METAL    606 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.50 GB tm    517.37us/   201.34ms (     5.07 GFLOPS    6.1|6.1     GB/s) ['normal']\n",
      "*** METAL    607 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.50 GB tm     39.00us/   201.38ms (     0.03 GFLOPS    0.1|0.1     GB/s) ['uniform']\n",
      "*** METAL    608 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.50 GB tm     54.00us/   201.43ms (    28.67 GFLOPS  102.1|127.4   GB/s) ['layernorm']\n",
      "*** METAL    609 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.51 GB tm    588.50us/   202.02ms (  3595.57 GFLOPS   24.1|691.1   GB/s) ['linear']\n",
      "*** METAL    610 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.51 GB tm     96.00us/   202.12ms (   402.43 GFLOPS   95.4|146.4   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    611 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.51 GB tm      5.38us/   202.12ms (    40.01 GFLOPS  176.0|176.0   GB/s) ['softmax']\n",
      "*** METAL    612 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.51 GB tm      5.00us/   202.13ms (   182.78 GFLOPS  197.8|197.8   GB/s) ['softmax']\n",
      "*** METAL    613 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.51 GB tm     19.29us/   202.15ms (   101.44 GFLOPS  144.9|171.7   GB/s) ['softmax', 'dropout']\n",
      "*** METAL    614 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.52 GB tm     48.04us/   202.19ms (  1023.11 GFLOPS  248.0|458.4   GB/s) ['__matmul__']\n",
      "*** METAL    615 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.52 GB tm     14.25us/   202.21ms (     0.00 GFLOPS  386.3|386.3   GB/s) \n",
      "*** METAL    616 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.52 GB tm    296.50us/   202.50ms (  2390.46 GFLOPS   40.7|1225.4  GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    617 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.52 GB tm     15.75us/   202.52ms (   133.80 GFLOPS  175.1|546.1   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   618\u001b[0m custom_random                            mem  1.53 GB \n",
      "\u001b[32m*** CUSTOM   619\u001b[0m custom_random                            mem  1.53 GB \n",
      "\u001b[32m*** CUSTOM   620\u001b[0m custom_random                            mem  1.53 GB \n",
      "\u001b[32m*** CUSTOM   621\u001b[0m custom_random                            mem  1.53 GB \n",
      "*** METAL    622 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.53 GB tm     46.50us/   202.57ms (    54.11 GFLOPS   59.4|94.3    GB/s) ['layernorm']\n",
      "*** METAL    623 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.53 GB tm     71.50us/   202.64ms (    29.33 GFLOPS  117.3|117.3   GB/s) ['uniform']\n",
      "*** METAL    624 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.53 GB tm      2.75us/   202.64ms (     1.49 GFLOPS    6.0|6.0     GB/s) ['uniform']\n",
      "*** METAL    625 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.53 GB tm     64.17us/   202.70ms (    32.68 GFLOPS  130.7|130.7   GB/s) ['uniform']\n",
      "*** METAL    626 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.53 GB tm      2.13us/   202.71ms (     0.48 GFLOPS    1.9|1.9     GB/s) ['uniform']\n",
      "*** METAL    627 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.53 GB tm     38.63us/   202.75ms (    40.09 GFLOPS  142.8|178.2   GB/s) ['layernorm']\n",
      "*** METAL    628 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.54 GB tm   1176.33us/   203.92ms (  2424.14 GFLOPS   15.3|463.3   GB/s) ['linear', 'gelu']\n",
      "*** METAL    629 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.55 GB tm    685.87us/   204.61ms (  4110.46 GFLOPS   26.2|775.5   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   630\u001b[0m custom_random                            mem  1.55 GB \n",
      "*** METAL    631 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.55 GB tm    116.71us/   204.72ms (    29.48 GFLOPS   94.3|94.3    GB/s) ['dropout', '__add__']\n",
      "scheduled 22 kernels\n",
      "*** METAL    632 E_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                mem  1.55 GB tm      2.67us/   204.73ms (     0.17 GFLOPS    0.8|0.8     GB/s) ['__eq__']\n",
      "*** METAL    633 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.55 GB tm     68.50us/   204.80ms (    30.76 GFLOPS   40.3|125.6   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   634\u001b[0m custom_random                            mem  1.56 GB \n",
      "\u001b[32m*** CUSTOM   635\u001b[0m custom_random                            mem  1.56 GB \n",
      "\u001b[32m*** CUSTOM   636\u001b[0m custom_random                            mem  1.56 GB \n",
      "\u001b[32m*** CUSTOM   637\u001b[0m custom_random                            mem  1.56 GB \n",
      "\u001b[32m*** CUSTOM   638\u001b[0m custom_random                            mem  1.56 GB \n",
      "\u001b[32m*** CUSTOM   639\u001b[0m custom_random                            mem  1.56 GB \n",
      "*** METAL    640 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.56 GB tm     39.38us/   204.83ms (    63.90 GFLOPS   70.2|111.4   GB/s) ['layernorm']\n",
      "*** METAL    641 E_\u001b[34m6144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  1.56 GB tm     45.13us/   204.88ms (    34.86 GFLOPS  139.4|139.4   GB/s) ['uniform']\n",
      "*** METAL    642 E_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.56 GB tm      2.46us/   204.88ms (     1.25 GFLOPS    5.0|5.0     GB/s) ['uniform']\n",
      "*** METAL    643 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  1.56 GB tm     29.75us/   204.91ms (    88.12 GFLOPS  105.7|105.7   GB/s) ['normal']\n",
      "*** METAL    644 E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.56 GB tm      3.00us/   204.92ms (     0.34 GFLOPS    1.4|1.4     GB/s) ['uniform']\n",
      "*** METAL    645 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.56 GB tm     37.33us/   204.95ms (    41.47 GFLOPS  147.7|184.3   GB/s) ['layernorm']\n",
      "*** METAL    646 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  1.57 GB tm    507.33us/   205.46ms (  4170.82 GFLOPS   27.9|801.6   GB/s) ['linear']\n",
      "*** METAL    647 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  1.57 GB tm     54.83us/   205.51ms (   704.56 GFLOPS  167.0|256.4   GB/s) ['__matmul__', '__mul__', 'masked_fill']\n",
      "*** METAL    648 r_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                              mem  1.57 GB tm      4.92us/   205.52ms (    43.74 GFLOPS  192.4|192.4   GB/s) ['softmax']\n",
      "*** METAL    649 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m7\u001b[0m\u001b[90m\u001b[0m                             mem  1.57 GB tm      5.00us/   205.52ms (   182.79 GFLOPS  197.8|197.8   GB/s) ['softmax']\n",
      "*** METAL    650 E_\u001b[34m336\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                             mem  1.57 GB tm     24.50us/   205.55ms (    79.87 GFLOPS  114.1|135.2   GB/s) ['softmax', 'dropout']\n",
      "*** METAL    651 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                 mem  1.58 GB tm     45.79us/   205.59ms (  1073.38 GFLOPS  260.2|480.9   GB/s) ['__matmul__']\n",
      "*** METAL    652 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.58 GB tm     14.75us/   205.61ms (     0.00 GFLOPS  373.2|373.2   GB/s) \n",
      "*** METAL    653 r_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.58 GB tm    294.13us/   205.90ms (  2409.76 GFLOPS   41.0|1235.3  GB/s) ['linear', 'dropout', '__add__']\n",
      "scheduled 13 kernels\n",
      "*** METAL    654 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.58 GB tm     26.87us/   205.93ms (    78.41 GFLOPS  102.6|320.1   GB/s) ['layernorm']\n",
      "\u001b[32m*** CUSTOM   655\u001b[0m custom_random                            mem  1.59 GB \n",
      "\u001b[32m*** CUSTOM   656\u001b[0m custom_random                            mem  1.59 GB \n",
      "\u001b[32m*** CUSTOM   657\u001b[0m custom_random                            mem  1.59 GB \n",
      "\u001b[32m*** CUSTOM   658\u001b[0m custom_random                            mem  1.59 GB \n",
      "*** METAL    659 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.59 GB tm     91.00us/   206.02ms (    27.65 GFLOPS   30.4|48.2    GB/s) ['layernorm']\n",
      "*** METAL    660 E_\u001b[34m8192\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.59 GB tm     54.67us/   206.08ms (    38.36 GFLOPS  153.5|153.5   GB/s) ['uniform']\n",
      "*** METAL    661 E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.59 GB tm      2.37us/   206.08ms (     1.72 GFLOPS    6.9|6.9     GB/s) ['uniform']\n",
      "*** METAL    662 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.59 GB tm     55.12us/   206.13ms (    38.04 GFLOPS  152.2|152.2   GB/s) ['uniform']\n",
      "*** METAL    663 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.59 GB tm      2.12us/   206.14ms (     0.48 GFLOPS    1.9|1.9     GB/s) ['uniform']\n",
      "*** METAL    664 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.59 GB tm     42.13us/   206.18ms (    36.75 GFLOPS  130.9|163.4   GB/s) ['layernorm']\n",
      "*** METAL    665 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  1.60 GB tm    689.88us/   206.87ms (  4133.51 GFLOPS   26.0|790.0   GB/s) ['linear', 'gelu']\n",
      "*** METAL    666 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  1.61 GB tm    685.25us/   207.55ms (  4114.21 GFLOPS   26.2|776.2   GB/s) ['linear']\n",
      "\u001b[32m*** CUSTOM   667\u001b[0m custom_random                            mem  1.61 GB \n",
      "*** METAL    668 E_\u001b[34m7168\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                              mem  1.61 GB tm    117.33us/   207.67ms (    29.32 GFLOPS   93.8|93.8    GB/s) ['dropout', '__add__']\n",
      "*** METAL    669 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.61 GB tm    457.42us/   208.13ms (     4.61 GFLOPS    6.0|18.8    GB/s) ['layernorm']\n",
      "*** METAL    670 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.61 GB tm     32.42us/   208.16ms (    77.61 GFLOPS   85.2|135.3   GB/s) ['layernorm']\n",
      "*** METAL    671 E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  1.61 GB tm     41.88us/   208.20ms (    36.97 GFLOPS  131.7|164.3   GB/s) ['layernorm']\n",
      "scheduled 44 kernels\n",
      "memory reduced from 2.26 MB -> 1.45 MB, 16 -> 8 bufs\n",
      "*** METAL    672 E_\u001b[36m2\u001b[0m\u001b[90m\u001b[0m                                      mem  1.61 GB tm      3.75us/   208.21ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__sub__', '__add__']\n",
      "\u001b[32m*** METAL    673\u001b[0m \u001b[33mcopy      440,   METAL <- NPY    \u001b[0m        mem  1.61 GB tm     35.25us/   208.24ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[36m3\u001b[0m \u001b[31m15\u001b[0m \u001b[33m5\u001b[0m                         :     6.00 us  <  hc     : \u001b[34m5\u001b[0m \u001b[36m3\u001b[0m \u001b[35m15\u001b[0m                         :     6.92 us\n",
      "\u001b[32m*** METAL    674\u001b[0m r_\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m15\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                                 mem  1.61 GB tm      3.25us/   208.24ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__getitem__']\n",
      "*** METAL    675 r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                               mem  1.61 GB tm      4.00us/   208.25ms (     7.17 GFLOPS    4.1|12.8    GB/s) ['square', 'sum']\n",
      "*** METAL    676 r_\u001b[31m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                                 mem  1.61 GB tm      2.75us/   208.25ms (     0.51 GFLOPS    0.0|0.0     GB/s) ['argmax']\n",
      "*** METAL    677 r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                                 mem  1.61 GB tm      2.58us/   208.25ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__getitem__']\n",
      "*** METAL    678 r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                               mem  1.61 GB tm      4.17us/   208.26ms (     6.88 GFLOPS    3.9|12.3    GB/s) ['square', 'sum']\n",
      "\u001b[32m*** CUSTOM   679\u001b[0m custom_random                            mem  1.62 GB \n",
      "beam2  : \u001b[31m5\u001b[0m \u001b[33m5\u001b[0m                            :     6.13 us  <  hc     : \u001b[33m5\u001b[0m \u001b[35m5\u001b[0m                            :     6.17 us\n",
      "\u001b[32m*** METAL    680\u001b[0m r_\u001b[31m5\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                                    mem  1.62 GB tm      3.21us/   208.26ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['arange', '__mul__', '__add__', '__getitem__']\n",
      "beam2  : \u001b[34m3\u001b[0m \u001b[31m21\u001b[0m \u001b[33m7\u001b[0m                         :     5.75 us  <  hc     : \u001b[34m7\u001b[0m \u001b[36m3\u001b[0m \u001b[35m21\u001b[0m                         :     6.62 us\n",
      "\u001b[32m*** METAL    681\u001b[0m r_\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                                 mem  1.62 GB tm      4.00us/   208.27ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__getitem__']\n",
      "\u001b[32m*** CUSTOM   682\u001b[0m custom_random                            mem  1.62 GB \n",
      "\u001b[32m*** CUSTOM   683\u001b[0m custom_random                            mem  1.62 GB \n",
      "*** METAL    684 E_\u001b[34m30\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m\u001b[0m                                mem  1.62 GB tm      4.38us/   208.27ms (     2.63 GFLOPS    3.5|7.0     GB/s) ['__sub__', '__truediv__', '__mul__']\n",
      "hc     : \u001b[34m55\u001b[0m                             :     7.25 us  <  beam2  : \u001b[34m11\u001b[0m \u001b[33m5\u001b[0m                           :     7.96 us\n",
      "\u001b[32m*** METAL    685\u001b[0m E_\u001b[34m55\u001b[0m\u001b[90m\u001b[0m                                     mem  1.62 GB tm      4.87us/   208.27ms (     0.03 GFLOPS    0.2|0.2     GB/s) ['__getitem__']\n",
      "*** METAL    686 E_\u001b[34m2048\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm     19.04us/   208.29ms (    27.53 GFLOPS  110.1|110.1   GB/s) ['uniform']\n",
      "hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     8.00 us  <  beam2  : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m                       :     8.12 us\n",
      "\u001b[32m*** METAL    687\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn2\u001b[0m                             mem  1.62 GB tm      9.12us/   208.30ms (     3.59 GFLOPS   14.4|14.4    GB/s) ['uniform']\n",
      "*** METAL    688 E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                  mem  1.62 GB tm      6.75us/   208.31ms (     0.01 GFLOPS    0.0|0.0     GB/s) ['uniform']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[34m11\u001b[0m \u001b[36m32\u001b[0m \u001b[31m15\u001b[0m \u001b[33m2\u001b[0m \u001b[33m5\u001b[0m                 :    12.88 us  <  hc     : \u001b[34m55\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[33m4\u001b[0m \u001b[35m15\u001b[0m                   :    15.08 us\n",
      "\u001b[32m*** METAL    689\u001b[0m r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m15\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                         mem  1.62 GB tm     14.04us/   208.32ms (    22.56 GFLOPS    2.6|16.0    GB/s) ['__getitem__']\n",
      "tc     : \u001b[34m42\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :   344.50 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m            :   779.50 us  <  beam2  : \u001b[34m42\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m         :   824.17 us\n",
      "\u001b[32m*** METAL    690\u001b[0m r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m              mem  1.62 GB tm    786.21us/   209.11ms (   896.25 GFLOPS    8.3|115.5   GB/s) ['linear']\n",
      "*** METAL    691 r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m10\u001b[0m\u001b[90m\u001b[0m                       mem  1.62 GB tm     16.79us/   209.13ms (   112.70 GFLOPS   23.5|134.2   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[34m16\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m21\u001b[0m \u001b[33m5\u001b[0m              :    40.54 us  <  hc     : \u001b[34m8\u001b[0m \u001b[34m5\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[35m21\u001b[0m                :   112.21 us\n",
      "\u001b[32m*** METAL    692\u001b[0m r_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                      mem  1.62 GB tm     38.04us/   209.16ms (   361.78 GFLOPS   89.6|179.2   GB/s) ['__getitem__']\n",
      "*** METAL    693 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm     41.21us/   209.21ms (   564.18 GFLOPS   19.1|297.4   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m2\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[36m4\u001b[0m \u001b[33m11\u001b[0m                 :    26.83 us  <  hc     : \u001b[34m8\u001b[0m \u001b[34m11\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                  :    27.46 us\n",
      "\u001b[32m*** METAL    694\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m11\u001b[0m\u001b[90m\u001b[0m                         mem  1.62 GB tm     62.42us/   209.27ms (     0.00 GFLOPS   46.2|46.2    GB/s) ['cat']\n",
      "*** METAL    695 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm    340.83us/   209.61ms (   135.37 GFLOPS    3.6|46.5    GB/s) ['linear']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m          :    67.42 us  <  tc     : \u001b[34m22\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m          :    84.62 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m              :   168.71 us\n",
      "\u001b[32m*** METAL    696\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm     61.58us/   209.67ms (   374.96 GFLOPS   25.9|120.0   GB/s) ['linear']\n",
      "*** METAL    697 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.62 GB tm     56.04us/   209.73ms (     9.45 GFLOPS   12.9|27.3    GB/s) ['square', 'sum']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[31m4\u001b[0m \u001b[35m4\u001b[0m                      :     8.00 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.50 us\n",
      "\u001b[32m*** METAL    698\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm      9.00us/   209.74ms (     1.25 GFLOPS    5.3|5.3     GB/s) ['log_softmax']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m16\u001b[0m \u001b[35m16\u001b[0m                       :     7.75 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.63 us\n",
      "\u001b[32m*** METAL    699\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                               mem  1.62 GB tm    541.63us/   210.28ms (     0.03 GFLOPS    0.1|0.1     GB/s) ['log_softmax']\n",
      "*** METAL    700 r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm     43.75us/   210.32ms (   133.11 GFLOPS   17.9|26.3    GB/s) ['__add__', 'einsum', '__mul__', 'sqrt', '__neg__']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m                       :     8.38 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.67 us\n",
      "\u001b[32m*** METAL    701\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90mn2\u001b[0m                             mem  1.62 GB tm      9.50us/   210.33ms (     4.97 GFLOPS    5.3|5.3     GB/s) ['log_softmax']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m                       :     8.25 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.92 us\n",
      "\u001b[32m*** METAL    702\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90mn3\u001b[0m                             mem  1.62 GB tm      9.54us/   210.34ms (     4.94 GFLOPS    5.3|5.3     GB/s) ['log_softmax']\n",
      "*** METAL    703 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                               mem  1.62 GB tm      7.17us/   210.35ms (     1.47 GFLOPS    6.7|6.7     GB/s) ['argmax']\n",
      "*** METAL    704 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90mn1\u001b[0m                             mem  1.62 GB tm      7.88us/   210.35ms (     5.90 GFLOPS    6.4|12.2    GB/s) ['argmax']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[34m64\u001b[0m \u001b[36m32\u001b[0m \u001b[31m16\u001b[0m \u001b[33m4\u001b[0m                  :    46.37 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m16\u001b[0m               :   104.25 us\n",
      "\u001b[32m*** METAL    705\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                          mem  1.62 GB tm     46.50us/   210.40ms (   162.78 GFLOPS   31.4|344.9   GB/s) ['gather', '__sub__']\n",
      "*** METAL    706 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                             mem  1.62 GB tm    469.00us/   210.87ms (     1.13 GFLOPS    1.5|3.3     GB/s) ['square', 'sum']\n",
      "*** METAL    707 r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm     42.83us/   210.91ms (   135.96 GFLOPS   18.3|26.8    GB/s) ['__add__', 'einsum', '__mul__', 'sqrt', '__neg__']\n",
      "*** METAL    708 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                               mem  1.62 GB tm      6.75us/   210.92ms (     1.56 GFLOPS    7.1|7.1     GB/s) ['argmax']\n",
      "*** METAL    709 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90mn1\u001b[0m                             mem  1.62 GB tm      7.96us/   210.93ms (     5.84 GFLOPS    6.4|12.0    GB/s) ['argmax']\n",
      "*** METAL    710 E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  1.62 GB tm      6.75us/   210.93ms (     0.00 GFLOPS    1.7|1.7     GB/s) ['stack']\n",
      "beam2  : \u001b[34m176\u001b[0m \u001b[36m4\u001b[0m \u001b[32m4\u001b[0m \u001b[31m4\u001b[0m                      :     8.88 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    13.13 us\n",
      "\u001b[32m*** METAL    711\u001b[0m r_\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[32m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm     68.83us/   211.00ms (     1.23 GFLOPS    1.5|2.8     GB/s) ['gather', 'log_softmax']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m8\u001b[0m \u001b[31m16\u001b[0m                        :     8.79 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.87 us\n",
      "\u001b[32m*** METAL    712\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                                mem  1.62 GB tm     11.33us/   211.01ms (     6.09 GFLOPS    8.9|8.9     GB/s) ['gather', 'log_softmax']\n",
      "beam2  : \u001b[32m64\u001b[0m \u001b[31m11\u001b[0m                          :     9.54 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                          :    10.58 us\n",
      "\u001b[32m*** METAL    713\u001b[0m r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m11\u001b[0m\u001b[90m\u001b[0m                                  mem  1.62 GB tm     10.46us/   211.03ms (     0.93 GFLOPS    0.3|1.9     GB/s) ['exp', '__rsub__', 'pow', '__rmul__', '__mul__', 'mean']\n",
      "beam2  : \u001b[32m64\u001b[0m \u001b[31m11\u001b[0m                          :     9.29 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                          :    11.25 us\n",
      "\u001b[32m*** METAL    714\u001b[0m r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m11\u001b[0m\u001b[90mn1\u001b[0m                                mem  1.62 GB tm    646.79us/   211.67ms (     0.02 GFLOPS    0.0|0.0     GB/s) ['exp', '__rsub__', 'pow', '__rmul__', '__mul__', 'mean', '__add__']\n",
      "\u001b[32m*** CLANG    715\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     32.33us/   211.70ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "scheduled 28 kernels\n",
      "memory reduced from 4.00 MB -> 3.61 MB, 6 -> 4 bufs\n",
      "\u001b[32m*** CUSTOM   716\u001b[0m custom_random                            mem  1.62 GB \n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[31m16\u001b[0m                      :     8.58 us  <  hc     : \u001b[34m1408\u001b[0m \u001b[32m16\u001b[0m                        :    18.50 us\n",
      "\u001b[32m*** METAL    717\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm      9.62us/   211.71ms (     4.68 GFLOPS    9.9|9.9     GB/s) ['__truediv__', 'softmax']\n",
      "hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m2\u001b[0m                       :     7.62 us  <  beam2  : \u001b[34m1024\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m                       :    11.21 us\n",
      "\u001b[32m*** METAL    718\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                               mem  1.62 GB tm      8.58us/   211.72ms (     0.00 GFLOPS    7.6|7.6     GB/s) ['stack']\n",
      "\u001b[32m*** CUSTOM   719\u001b[0m custom_random                            mem  1.62 GB \n",
      "\u001b[32m*** CUSTOM   720\u001b[0m custom_random                            mem  1.62 GB \n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m16\u001b[0m \u001b[31m704\u001b[0m \u001b[33m2\u001b[0m                    :     5.75 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                      :    11.67 us\n",
      "\u001b[32m*** METAL    721\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m704\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.62 GB tm      6.92us/   211.73ms (     0.00 GFLOPS    0.4|0.4     GB/s) ['arange']\n",
      "beam2  : \u001b[36m2\u001b[0m \u001b[35m2\u001b[0m                            :     6.13 us  <  hc     : \u001b[33m2\u001b[0m \u001b[35m2\u001b[0m                            :     6.13 us\n",
      "\u001b[32m*** METAL    722\u001b[0m r_\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m2\u001b[0m\u001b[90mn1\u001b[0m                                  mem  1.62 GB tm      6.63us/   211.74ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['arange']\n",
      "beam2  : \u001b[34m176\u001b[0m \u001b[36m8\u001b[0m \u001b[32m8\u001b[0m \u001b[31m2\u001b[0m                      :     8.25 us  <  hc     : \u001b[34m1408\u001b[0m \u001b[32m16\u001b[0m                        :    18.71 us\n",
      "\u001b[32m*** METAL    723\u001b[0m r_\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm     10.00us/   211.75ms (    21.40 GFLOPS   10.1|58.6    GB/s) ['__truediv__', 'softmax']\n",
      "hc     : \u001b[34m1280\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    16.58 us  <  beam2  : \u001b[34m80\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m                    :    16.79 us\n",
      "\u001b[32m*** METAL    724\u001b[0m E_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm     17.50us/   211.76ms (    18.72 GFLOPS   74.9|74.9    GB/s) ['uniform']\n",
      "beam2  : \u001b[34m20\u001b[0m \u001b[36m32\u001b[0m                          :     7.21 us  <  hc     : \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                         :     7.37 us\n",
      "\u001b[32m*** METAL    725\u001b[0m E_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m\u001b[0m                                  mem  1.62 GB tm      8.04us/   211.77ms (     0.32 GFLOPS    0.3|0.6     GB/s) ['uniform']\n",
      "hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                        :     7.25 us  <  beam2  : \u001b[34m22\u001b[0m \u001b[36m32\u001b[0m                          :     7.33 us\n",
      "\u001b[32m*** METAL    726\u001b[0m E_\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  1.62 GB tm      8.29us/   211.78ms (     0.25 GFLOPS    0.7|0.7     GB/s) ['__getitem__']\n",
      "beam2  : \u001b[36m2\u001b[0m                              :     7.12 us  <  hc     : \u001b[36m2\u001b[0m                              :     7.12 us\n",
      "\u001b[32m*** METAL    727\u001b[0m E_\u001b[36m2\u001b[0m\u001b[90mn1\u001b[0m                                    mem  1.62 GB tm      7.33us/   211.79ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__getitem__']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m                    :     5.38 us  <  hc     : \u001b[34m44\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                      :     6.46 us\n",
      "\u001b[32m*** METAL    728\u001b[0m E_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.62 GB tm     10.50us/   211.80ms (    15.02 GFLOPS   18.2|34.3    GB/s) ['__truediv__', 'softmax']\n",
      "tc     : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m5\u001b[0m      :   357.38 us  <  beam2  : \u001b[34m88\u001b[0m \u001b[34m40\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m           :   405.33 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m5\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :   595.37 us\n",
      "\u001b[32m*** METAL    729\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m              mem  1.62 GB tm    326.46us/   212.12ms (   707.32 GFLOPS    9.2|82.9    GB/s) ['linear']\n",
      "beam2  : \u001b[34m11\u001b[0m \u001b[36m4\u001b[0m \u001b[36m32\u001b[0m \u001b[35m16\u001b[0m \u001b[33m2\u001b[0m \u001b[33m8\u001b[0m                 :    10.58 us  <  hc     : \u001b[34m176\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[35m16\u001b[0m                    :    10.88 us\n",
      "\u001b[32m*** METAL    730\u001b[0m r_\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90m\u001b[0m                         mem  1.62 GB tm     11.54us/   212.14ms (     3.54 GFLOPS   15.6|22.9    GB/s) ['multinomial']\n",
      "beam2  : \u001b[34m176\u001b[0m \u001b[36m8\u001b[0m \u001b[35m16\u001b[0m                       :     8.62 us  <  hc     : \u001b[34m1408\u001b[0m \u001b[32m16\u001b[0m                        :    18.58 us\n",
      "\u001b[32m*** METAL    731\u001b[0m r_\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                               mem  1.62 GB tm      9.83us/   212.15ms (     9.16 GFLOPS   10.3|10.9    GB/s) ['multinomial']\n",
      "hc     : \u001b[34m11\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.17 us  <  beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[33m2\u001b[0m                        :     7.50 us\n",
      "\u001b[32m*** METAL    732\u001b[0m E_\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  1.62 GB tm      8.62us/   212.15ms (     0.00 GFLOPS    1.3|1.3     GB/s) ['__eq__']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m2\u001b[0m                      :     4.13 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     4.33 us\n",
      "\u001b[32m*** METAL    733\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                              mem  1.62 GB tm      8.46us/   212.16ms (     1.69 GFLOPS    1.3|3.9     GB/s) ['__getitem__']\n",
      "hc     : \u001b[34m2\u001b[0m \u001b[34m88\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                  :    18.75 us  <  beam2  : \u001b[34m2\u001b[0m \u001b[34m704\u001b[0m \u001b[36m16\u001b[0m \u001b[36m16\u001b[0m                    :    44.54 us\n",
      "\u001b[32m*** METAL    734\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                          mem  1.62 GB tm     19.38us/   212.18ms (    13.95 GFLOPS   75.0|111.6   GB/s) ['__eq__', 'masked_fill']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[34m5\u001b[0m \u001b[36m32\u001b[0m \u001b[31m704\u001b[0m \u001b[31m2\u001b[0m \u001b[35m16\u001b[0m \u001b[33m2\u001b[0m           :  1729.25 us  <  hc     : \u001b[34m11\u001b[0m \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m704\u001b[0m \u001b[31m2\u001b[0m \u001b[33m4\u001b[0m \u001b[35m16\u001b[0m         :  3378.13 us\n",
      "\u001b[32m*** METAL    735\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m704\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                   mem  1.62 GB tm   1646.92us/   213.83ms (   505.71 GFLOPS    0.6|806.9   GB/s) ['__getitem__']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[34m176\u001b[0m \u001b[34m16\u001b[0m \u001b[36m16\u001b[0m \u001b[31m16\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m             :    85.88 us  <  hc     : \u001b[34m2\u001b[0m \u001b[34m22\u001b[0m \u001b[34m4\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m16\u001b[0m             :    86.75 us\n",
      "\u001b[32m*** METAL    736\u001b[0m r_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                     mem  1.63 GB tm     85.62us/   213.91ms (   282.04 GFLOPS   34.1|185.2   GB/s) ['__eq__', 'gather', 'masked_fill']\n",
      "hc     : \u001b[34m22\u001b[0m \u001b[34m5\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[35m2\u001b[0m                    :     3.92 us  <  beam2  : \u001b[34m44\u001b[0m \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[31m2\u001b[0m \u001b[33m2\u001b[0m                    :     4.71 us\n",
      "\u001b[32m*** METAL    737\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m5\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.62 GB tm    208.71us/   214.12ms (     0.03 GFLOPS    0.4|0.4     GB/s) ['sum']\n",
      "beam2  : \u001b[34m352\u001b[0m \u001b[36m32\u001b[0m \u001b[36m16\u001b[0m \u001b[31m2\u001b[0m                    :    23.50 us  <  hc     : \u001b[34m5632\u001b[0m \u001b[36m32\u001b[0m \u001b[35m2\u001b[0m                      :    32.25 us\n",
      "\u001b[32m*** METAL    738\u001b[0m r_\u001b[34m352\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m\u001b[0m                            mem  1.62 GB tm     24.92us/   214.15ms (    14.47 GFLOPS   86.8|86.8    GB/s) ['sum']\n",
      "*** METAL    739 r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m            mem  1.62 GB tm     68.63us/   214.22ms (   674.94 GFLOPS   17.7|220.6   GB/s) ['linear', 'relu']\n",
      "*** METAL    740 r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm     40.04us/   214.26ms (   580.62 GFLOPS   19.7|306.1   GB/s) ['linear', 'relu']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[35m16\u001b[0m            :    16.54 us  <  hc     : \u001b[34m11\u001b[0m \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[31m32\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m               :    50.25 us\n",
      "\u001b[32m*** METAL    741\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m16\u001b[0m\u001b[90m\u001b[0m                    mem  1.62 GB tm     18.00us/   214.28ms (   182.73 GFLOPS   25.0|90.1    GB/s) ['linear', '__add__', '__sub__']\n",
      "beam2  : \u001b[32m64\u001b[0m \u001b[31m110\u001b[0m                         :     8.71 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m440\u001b[0m                         :    16.13 us\n",
      "\u001b[32m*** METAL    742\u001b[0m r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m110\u001b[0m\u001b[90mn1\u001b[0m                               mem  1.62 GB tm     16.58us/   214.29ms (     2.37 GFLOPS    1.7|2.7     GB/s) ['abs', 'mean']\n",
      "\u001b[32m*** CLANG    743\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     24.46us/   214.32ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                          :    11.63 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                          :    11.63 us\n",
      "\u001b[32m*** METAL    744\u001b[0m r_\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m44\u001b[0m\u001b[90m\u001b[0m                                  mem  1.62 GB tm     12.42us/   214.33ms (     0.19 GFLOPS    0.5|0.5     GB/s) ['__eq__', 'int', 'sum', '__truediv__']\n",
      "\u001b[32m*** CLANG    745\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     26.12us/   214.35ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                          :    11.21 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                          :    12.58 us\n",
      "\u001b[32m*** METAL    746\u001b[0m r_\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m44\u001b[0m\u001b[90mn1\u001b[0m                                mem  1.62 GB tm      6.00us/   214.36ms (     0.40 GFLOPS    1.1|1.1     GB/s) ['__eq__', 'int', 'sum', '__truediv__']\n",
      "\u001b[32m*** CLANG    747\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     25.37us/   214.39ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "*** METAL    748 r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                  mem  1.62 GB tm     10.54us/   214.40ms (   277.82 GFLOPS   40.0|149.6   GB/s) ['linear', '__sub__']\n",
      "*** METAL    749 r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m110\u001b[0m\u001b[90mn1\u001b[0m                               mem  1.62 GB tm      5.67us/   214.40ms (     6.95 GFLOPS    5.0|8.0     GB/s) ['abs', 'mean']\n",
      "\u001b[32m*** CLANG    750\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     12.00us/   214.41ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "\u001b[32m*** CLANG    751\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     10.67us/   214.42ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[32m64\u001b[0m \u001b[31m110\u001b[0m                         :    16.04 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m440\u001b[0m                         :    31.50 us\n",
      "\u001b[32m*** METAL    752\u001b[0m r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m110\u001b[0m\u001b[90mn2\u001b[0m                               mem  1.62 GB tm      8.62us/   214.43ms (     4.56 GFLOPS    3.3|5.2     GB/s) ['abs', 'max']\n",
      "\u001b[32m*** CLANG    753\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     20.96us/   214.45ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[32m64\u001b[0m \u001b[31m110\u001b[0m                         :    13.58 us  <  hc     : \u001b[32m16\u001b[0m \u001b[31m440\u001b[0m                         :    25.42 us\n",
      "\u001b[32m*** METAL    754\u001b[0m r_\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m110\u001b[0m\u001b[90mn3\u001b[0m                               mem  1.62 GB tm      7.50us/   214.46ms (     2.43 GFLOPS    3.8|6.0     GB/s) ['square', 'mean']\n",
      "\u001b[32m*** CLANG    755\u001b[0m \u001b[33mcopy        4,   CLANG <- METAL  \u001b[0m        mem  1.62 GB tm     18.79us/   214.48ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "scheduled 494 kernels\n",
      "memory reduced from 3963.32 MB -> 590.57 MB, 279 -> 84 bufs\n",
      "beam2  :                                :     8.00 us  <  hc     :                                :     8.00 us\n",
      "\u001b[32m*** METAL    756\u001b[0m E_\u001b[90mn4\u001b[0m                                     mem  1.54 GB tm      7.58us/   214.49ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__imul__']\n",
      "*** METAL    757 E_\u001b[90mn1\u001b[0m                                     mem  1.54 GB tm      6.62us/   214.49ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__imul__']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m32\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                 :    28.75 us  <  hc     : \u001b[34m8192\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    61.37 us\n",
      "\u001b[32m*** METAL    758\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                         mem  1.54 GB tm     27.38us/   214.52ms (     0.00 GFLOPS  114.9|114.9   GB/s) \n",
      "hc     : \u001b[34m55\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     7.25 us  <  beam2  : \u001b[34m64\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                      :     7.50 us\n",
      "\u001b[32m*** METAL    759\u001b[0m E_\u001b[34m55\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                              mem  1.86 GB tm      8.96us/   214.53ms (     3.14 GFLOPS    6.3|6.3     GB/s) ['abs', 'abs bw', '__sub__ bw']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m8\u001b[0m \u001b[33m2\u001b[0m                         :     7.38 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                        :     7.46 us\n",
      "\u001b[32m*** METAL    760\u001b[0m E_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                 mem  2.02 GB tm      8.42us/   214.54ms (     1.09 GFLOPS    0.7|0.7     GB/s) ['exp', '__rsub__', 'pow', '__rmul__', '__mul__ bw', '__rmul__ bw', 'pow bw', '__add__', '__rsub__ bw', 'exp bw']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m                     :     8.25 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.58 us\n",
      "\u001b[32m*** METAL    761\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                             mem  2.02 GB tm      9.58us/   214.55ms (     4.70 GFLOPS    5.3|5.3     GB/s) ['log_softmax bw']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m                         :     6.62 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                        :     6.63 us\n",
      "\u001b[32m*** METAL    762\u001b[0m E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90mn1\u001b[0m                               mem  2.02 GB tm      8.00us/   214.56ms (     1.14 GFLOPS    0.7|0.7     GB/s) ['exp', '__rsub__', 'pow', '__rmul__', '__mul__ bw', '__rmul__ bw', 'pow bw', '__add__', '__rsub__ bw', 'exp bw']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m8\u001b[0m \u001b[32m16\u001b[0m                        :     8.21 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.67 us\n",
      "\u001b[32m*** METAL    763\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90mn1\u001b[0m                              mem  2.02 GB tm     90.46us/   214.65ms (     2.37 GFLOPS    1.1|10.0    GB/s) ['log_softmax bw']\n",
      "beam2  : \u001b[34m672\u001b[0m \u001b[36m16\u001b[0m \u001b[31m21\u001b[0m                      :    16.00 us  <  hc     : \u001b[34m112\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m \u001b[35m21\u001b[0m                    :    16.37 us\n",
      "\u001b[32m*** METAL    764\u001b[0m r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     18.63us/   214.67ms (    48.49 GFLOPS   53.1|53.1    GB/s) ['softmax bw']\n",
      "*** METAL    765 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     18.17us/   214.68ms (    49.72 GFLOPS   54.5|54.5    GB/s) ['softmax bw']\n",
      "*** METAL    766 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     18.00us/   214.70ms (    50.18 GFLOPS   55.0|55.0    GB/s) ['softmax bw']\n",
      "*** METAL    767 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     16.83us/   214.72ms (    53.65 GFLOPS   58.8|58.8    GB/s) ['softmax bw']\n",
      "*** METAL    768 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm    245.58us/   214.96ms (     3.68 GFLOPS    4.0|4.0     GB/s) ['softmax bw']\n",
      "*** METAL    769 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     17.67us/   214.98ms (    51.12 GFLOPS   56.0|56.0    GB/s) ['softmax bw']\n",
      "*** METAL    770 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     16.87us/   215.00ms (    53.52 GFLOPS   58.6|58.6    GB/s) ['softmax bw']\n",
      "*** METAL    771 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m21\u001b[0m\u001b[90m\u001b[0m                              mem  2.02 GB tm     17.04us/   215.02ms (    53.00 GFLOPS   58.0|58.0    GB/s) ['softmax bw']\n",
      "beam2  : \u001b[34m256\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                   :    13.42 us  <  hc     : \u001b[34m2048\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    19.87 us\n",
      "\u001b[32m*** METAL    772\u001b[0m E_\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                           mem  2.02 GB tm    362.75us/   215.38ms (     0.00 GFLOPS    2.9|2.9     GB/s) \n",
      "beam2  : \u001b[34m1024\u001b[0m \u001b[36m32\u001b[0m \u001b[36m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                  :    33.58 us  <  hc     : \u001b[34m8192\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    69.37 us\n",
      "\u001b[32m*** METAL    773\u001b[0m E_\u001b[34m1024\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                          mem  2.03 GB tm     32.13us/   215.41ms (     0.00 GFLOPS  130.6|130.6   GB/s) \n",
      "beam2  : \u001b[34m1024\u001b[0m \u001b[36m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                  :    31.21 us  <  hc     : \u001b[34m8192\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    68.04 us\n",
      "\u001b[32m*** METAL    774\u001b[0m E_\u001b[34m1024\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                          mem  2.03 GB tm     31.46us/   215.44ms (     0.00 GFLOPS  133.3|266.7   GB/s) \n",
      "hc     : \u001b[34m98\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                        :     6.67 us  <  beam2  : \u001b[34m98\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                        :     6.75 us\n",
      "\u001b[32m*** METAL    775\u001b[0m E_\u001b[34m98\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                                mem  2.03 GB tm      7.13us/   215.45ms (     0.00 GFLOPS    5.3|5.3     GB/s) \n",
      "beam2  : \u001b[34m1280\u001b[0m \u001b[34m21\u001b[0m \u001b[34m7\u001b[0m \u001b[36m16\u001b[0m \u001b[32m3\u001b[0m \u001b[33m3\u001b[0m \u001b[35m3\u001b[0m             :  1957.46 us  <  hc     : \u001b[34m640\u001b[0m \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m32\u001b[0m \u001b[36m3\u001b[0m \u001b[33m3\u001b[0m \u001b[35m3\u001b[0m \u001b[35m3\u001b[0m             :  2020.67 us\n",
      "\u001b[32m*** METAL    776\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m3\u001b[0m\u001b[90m\u001b[0m                     mem  2.07 GB tm   5108.75us/   220.56ms (    63.64 GFLOPS   45.2|176.8   GB/s) ['max_pool2d bw']\n",
      "hc     : \u001b[34m80\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m \u001b[35m9\u001b[0m                      :    11.63 us  <  beam2  : \u001b[34m320\u001b[0m \u001b[36m32\u001b[0m \u001b[31m32\u001b[0m                      :    14.88 us\n",
      "\u001b[32m*** METAL    777\u001b[0m r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m9\u001b[0m\u001b[90mn1\u001b[0m                            mem  2.07 GB tm     16.92us/   220.58ms (    16.34 GFLOPS   26.6|26.6    GB/s) ['softmax bw']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m4\u001b[0m \u001b[33m8\u001b[0m                          :     6.00 us  <  hc     : \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                           :     6.13 us\n",
      "\u001b[32m*** METAL    778\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90m\u001b[0m                                  mem  2.07 GB tm     54.75us/   220.63ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[34m96\u001b[0m \u001b[36m32\u001b[0m \u001b[36m3\u001b[0m \u001b[33m4\u001b[0m                      :     7.37 us  <  hc     : \u001b[34m384\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                       :     8.38 us\n",
      "\u001b[32m*** METAL    779\u001b[0m E_\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                              mem  2.07 GB tm      7.67us/   220.64ms (     0.00 GFLOPS   19.2|19.2    GB/s) \n",
      "beam2  : \u001b[34m96\u001b[0m \u001b[36m4\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m                    :     8.13 us  <  hc     : \u001b[34m768\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                       :    10.63 us\n",
      "\u001b[32m*** METAL    780\u001b[0m E_\u001b[34m96\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  2.07 GB tm      8.50us/   220.65ms (     0.00 GFLOPS   34.7|34.7    GB/s) \n",
      "beam2  : \u001b[34m144\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                   :    10.25 us  <  hc     : \u001b[34m1536\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    15.29 us\n",
      "\u001b[32m*** METAL    781\u001b[0m E_\u001b[34m144\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                           mem  2.07 GB tm     10.58us/   220.66ms (     0.00 GFLOPS   55.7|55.7    GB/s) \n",
      "hc     : \u001b[34m64\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     6.62 us  <  beam2  : \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                      :     7.38 us\n",
      "\u001b[32m*** METAL    782\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                              mem  2.07 GB tm      6.92us/   220.66ms (     0.00 GFLOPS    4.7|4.7     GB/s) \n",
      "beam2  : \u001b[34m288\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                   :    13.88 us  <  hc     : \u001b[34m3072\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    23.75 us\n",
      "\u001b[32m*** METAL    783\u001b[0m E_\u001b[34m288\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                           mem  2.07 GB tm     14.08us/   220.68ms (     0.00 GFLOPS   83.8|83.8    GB/s) \n",
      "beam2  : \u001b[34m1152\u001b[0m \u001b[36m32\u001b[0m \u001b[33m16\u001b[0m                     :    23.92 us  <  hc     : \u001b[34m6144\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                      :    54.54 us\n",
      "\u001b[32m*** METAL    784\u001b[0m E_\u001b[34m1152\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m16\u001b[0m\u001b[90m\u001b[0m                             mem  2.07 GB tm     24.12us/   220.70ms (     0.00 GFLOPS   97.8|97.8    GB/s) \n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m4\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                     :     7.08 us  <  hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     7.54 us\n",
      "\u001b[32m*** METAL    785\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                             mem  2.07 GB tm      7.54us/   220.71ms (     0.00 GFLOPS   17.4|17.4    GB/s) \n",
      "beam2  : \u001b[34m1152\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                  :    35.25 us  <  hc     : \u001b[34m12288\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                     :    85.83 us\n",
      "\u001b[32m*** METAL    786\u001b[0m E_\u001b[34m1152\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                          mem  2.08 GB tm     33.92us/   220.74ms (     0.00 GFLOPS  139.1|139.1   GB/s) \n",
      "hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     6.50 us  <  beam2  : \u001b[34m4\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[33m4\u001b[0m                       :     6.67 us\n",
      "\u001b[32m*** METAL    787\u001b[0m E_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn2\u001b[0m                               mem  2.08 GB tm      6.67us/   220.75ms (     0.00 GFLOPS    0.3|0.3     GB/s) \n",
      "beam2  : \u001b[34m3072\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[33m3\u001b[0m \u001b[33m2\u001b[0m                  :    78.75 us  <  hc     : \u001b[34m24576\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m                     :   121.88 us\n",
      "\u001b[32m*** METAL    788\u001b[0m E_\u001b[34m3072\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                          mem  2.09 GB tm     72.63us/   220.82ms (     0.00 GFLOPS  129.9|129.9   GB/s) \n",
      "beam2  : \u001b[34m128\u001b[0m \u001b[36m4\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                   :    10.33 us  <  hc     : \u001b[34m1024\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    12.08 us\n",
      "\u001b[32m*** METAL    789\u001b[0m E_\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                           mem  2.09 GB tm     10.00us/   220.83ms (     0.00 GFLOPS   52.4|52.4    GB/s) \n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m                      :     6.83 us  <  hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     6.88 us\n",
      "\u001b[32m*** METAL    790\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  2.09 GB tm      7.21us/   220.84ms (     0.00 GFLOPS    9.1|9.1     GB/s) \n",
      "beam2  : \u001b[36m16\u001b[0m \u001b[33m2\u001b[0m                           :     6.29 us  <  hc     : \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m                            :     6.58 us\n",
      "\u001b[32m*** METAL    791\u001b[0m E_\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                   mem  2.09 GB tm      6.62us/   220.85ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "hc     : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     6.00 us  <  beam2  : \u001b[34m32\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                        :     6.38 us\n",
      "\u001b[32m*** METAL    792\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                              mem  2.09 GB tm      6.75us/   220.85ms (     0.00 GFLOPS    2.4|2.4     GB/s) \n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m \u001b[33m2\u001b[0m                        :     6.29 us  <  hc     : \u001b[34m8\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     6.33 us\n",
      "\u001b[32m*** METAL    793\u001b[0m E_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                                mem  2.09 GB tm      6.88us/   220.86ms (     0.00 GFLOPS    0.6|0.6     GB/s) \n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m8\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                     :     7.12 us  <  hc     : \u001b[34m256\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     7.75 us\n",
      "\u001b[32m*** METAL    794\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                             mem  2.09 GB tm      7.58us/   220.87ms (     0.00 GFLOPS   17.3|17.3    GB/s) \n",
      "beam2  : \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m                        :     6.25 us  <  hc     : \u001b[34m4\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                         :     6.38 us\n",
      "\u001b[32m*** METAL    795\u001b[0m E_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                                mem  2.09 GB tm      6.75us/   220.87ms (     0.00 GFLOPS    0.3|0.3     GB/s) \n",
      "beam2  : \u001b[34m160\u001b[0m \u001b[36m8\u001b[0m \u001b[36m8\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                    :    10.50 us  <  hc     : \u001b[34m1280\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                      :    14.00 us\n",
      "\u001b[32m*** METAL    796\u001b[0m E_\u001b[34m160\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm     11.04us/   220.89ms (     0.00 GFLOPS   59.4|59.4    GB/s) \n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m32\u001b[0m \u001b[33m5\u001b[0m                         :     6.25 us  <  hc     : \u001b[34m5\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                         :     6.38 us\n",
      "\u001b[32m*** METAL    797\u001b[0m E_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m5\u001b[0m\u001b[90m\u001b[0m                                 mem  2.13 GB tm      6.75us/   220.89ms (     0.00 GFLOPS    0.2|0.2     GB/s) \n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m                      :     6.79 us  <  hc     : \u001b[34m128\u001b[0m \u001b[36m32\u001b[0m \u001b[33m4\u001b[0m                       :     6.88 us\n",
      "\u001b[32m*** METAL    798\u001b[0m E_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                              mem  2.13 GB tm      7.50us/   220.90ms (     0.00 GFLOPS    8.7|8.7     GB/s) \n",
      "\u001b[32m*** METAL    799\u001b[0m \u001b[33mcopy        4,   METAL <- NPY    \u001b[0m        mem  2.13 GB tm     43.83us/   220.94ms (     0.00 GFLOPS    0.0|0.0     GB/s) \n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[31m4\u001b[0m \u001b[35m4\u001b[0m                    :     8.42 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.71 us\n",
      "\u001b[32m*** METAL    800\u001b[0m r_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm      9.25us/   220.95ms (     7.31 GFLOPS    8.0|8.0     GB/s) ['log_softmax']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[36m16\u001b[0m \u001b[31m4\u001b[0m \u001b[35m4\u001b[0m                      :     4.13 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :     8.29 us\n",
      "\u001b[32m*** METAL    801\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m4\u001b[0m\u001b[90mn1\u001b[0m                            mem  2.13 GB tm      8.71us/   220.96ms (     5.25 GFLOPS    5.8|5.8     GB/s) ['log_softmax']\n",
      "hc     : \u001b[34m1344\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m                     :    33.17 us  <  beam2  : \u001b[34m1344\u001b[0m \u001b[32m32\u001b[0m \u001b[31m16\u001b[0m                     :    44.63 us\n",
      "\u001b[32m*** METAL    802\u001b[0m r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     74.67us/   221.04ms (    32.54 GFLOPS   37.0|58.8    GB/s) ['layernorm']\n",
      "*** METAL    803 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm    137.21us/   221.17ms (    17.71 GFLOPS   20.1|32.0    GB/s) ['layernorm']\n",
      "*** METAL    804 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     51.00us/   221.22ms (    47.65 GFLOPS   54.2|86.0    GB/s) ['layernorm']\n",
      "*** METAL    805 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     43.17us/   221.27ms (    56.29 GFLOPS   64.0|101.6   GB/s) ['layernorm']\n",
      "*** METAL    806 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     42.00us/   221.31ms (    57.86 GFLOPS   65.8|104.4   GB/s) ['layernorm']\n",
      "*** METAL    807 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     36.38us/   221.35ms (    66.80 GFLOPS   76.0|120.6   GB/s) ['layernorm']\n",
      "*** METAL    808 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     34.25us/   221.38ms (    70.95 GFLOPS   80.7|128.1   GB/s) ['layernorm']\n",
      "*** METAL    809 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     33.62us/   221.41ms (    72.27 GFLOPS   82.2|130.5   GB/s) ['layernorm']\n",
      "*** METAL    810 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     34.13us/   221.45ms (    71.21 GFLOPS   81.0|128.6   GB/s) ['layernorm']\n",
      "*** METAL    811 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm    505.87us/   221.95ms (     4.80 GFLOPS    5.5|8.7     GB/s) ['layernorm']\n",
      "*** METAL    812 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     33.04us/   221.99ms (    73.54 GFLOPS   83.6|132.8   GB/s) ['layernorm']\n",
      "*** METAL    813 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     32.17us/   222.02ms (    75.54 GFLOPS   85.9|136.4   GB/s) ['layernorm']\n",
      "*** METAL    814 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     32.17us/   222.05ms (    75.54 GFLOPS   85.9|136.4   GB/s) ['layernorm']\n",
      "*** METAL    815 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     32.29us/   222.08ms (    75.25 GFLOPS   85.6|135.8   GB/s) ['layernorm']\n",
      "*** METAL    816 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     32.75us/   222.12ms (    74.20 GFLOPS   84.4|133.9   GB/s) ['layernorm']\n",
      "*** METAL    817 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     32.75us/   222.15ms (    74.20 GFLOPS   84.4|133.9   GB/s) ['layernorm']\n",
      "*** METAL    818 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     32.08us/   222.18ms (    75.74 GFLOPS   86.1|136.7   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m80\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[31m144\u001b[0m                    :    65.08 us  <  hc     : \u001b[34m320\u001b[0m \u001b[36m32\u001b[0m \u001b[31m36\u001b[0m \u001b[35m4\u001b[0m                    :    78.25 us\n",
      "\u001b[32m*** METAL    819\u001b[0m r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    208.42us/   222.39ms (    21.27 GFLOPS   28.7|28.7    GB/s) ['layernorm']\n",
      "*** METAL    820 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    101.92us/   222.49ms (    43.51 GFLOPS   58.7|58.7    GB/s) ['layernorm']\n",
      "*** METAL    821 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    200.33us/   222.69ms (    22.13 GFLOPS   29.9|29.9    GB/s) ['layernorm']\n",
      "*** METAL    822 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    308.71us/   223.00ms (    14.36 GFLOPS   19.4|19.4    GB/s) ['layernorm']\n",
      "*** METAL    823 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m144\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm     93.12us/   223.09ms (    47.61 GFLOPS   64.2|64.2    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m1280\u001b[0m \u001b[36m4\u001b[0m \u001b[32m64\u001b[0m \u001b[31m9\u001b[0m                    :   114.58 us  <  hc     : \u001b[34m160\u001b[0m \u001b[36m32\u001b[0m \u001b[31m144\u001b[0m \u001b[35m4\u001b[0m                   :   158.62 us\n",
      "\u001b[32m*** METAL    824\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m9\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    122.17us/   223.22ms (   246.77 GFLOPS   96.9|815.4   GB/s) ['layernorm']\n",
      "*** METAL    825 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m9\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    112.13us/   223.33ms (   268.87 GFLOPS  105.6|888.4   GB/s) ['layernorm']\n",
      "*** METAL    826 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m9\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    127.75us/   223.46ms (   235.98 GFLOPS   92.7|779.8   GB/s) ['layernorm']\n",
      "*** METAL    827 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m9\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    126.87us/   223.58ms (   237.61 GFLOPS   93.3|785.1   GB/s) ['layernorm']\n",
      "*** METAL    828 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[32m64\u001b[0m\u001b[90m_\u001b[0m\u001b[31m9\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm    111.17us/   223.69ms (   271.18 GFLOPS  106.5|896.1   GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m80\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m \u001b[31m1936\u001b[0m                   :   237.79 us  <  hc     : \u001b[34m80\u001b[0m \u001b[36m32\u001b[0m \u001b[31m484\u001b[0m \u001b[35m4\u001b[0m                    :   249.46 us\n",
      "\u001b[32m*** METAL    829\u001b[0m r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm    294.79us/   223.99ms (    50.45 GFLOPS   67.3|67.3    GB/s) ['layernorm']\n",
      "*** METAL    830 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm    415.63us/   224.40ms (    35.78 GFLOPS   47.7|47.7    GB/s) ['layernorm']\n",
      "*** METAL    831 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm    188.46us/   224.59ms (    78.91 GFLOPS  105.3|105.3   GB/s) ['layernorm']\n",
      "*** METAL    832 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm    186.75us/   224.78ms (    79.63 GFLOPS  106.3|106.3   GB/s) ['layernorm']\n",
      "*** METAL    833 r_\u001b[34m80\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m1936\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm    195.83us/   224.98ms (    75.94 GFLOPS  101.3|101.3   GB/s) ['layernorm']\n",
      "hc     : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m441\u001b[0m                    :   330.58 us  <  beam2  : \u001b[34m80\u001b[0m \u001b[36m16\u001b[0m \u001b[31m7056\u001b[0m                     :   333.88 us\n",
      "\u001b[32m*** METAL    834\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90mn1\u001b[0m                          mem  2.13 GB tm    349.33us/   225.32ms (    78.56 GFLOPS  103.4|107.9   GB/s) ['layernorm']\n",
      "*** METAL    835 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90mn1\u001b[0m                          mem  2.13 GB tm    346.50us/   225.67ms (    79.20 GFLOPS  104.3|108.8   GB/s) ['layernorm']\n",
      "*** METAL    836 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90mn1\u001b[0m                          mem  2.13 GB tm    341.79us/   226.01ms (    80.29 GFLOPS  105.7|110.3   GB/s) ['layernorm']\n",
      "*** METAL    837 r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m441\u001b[0m\u001b[90mn1\u001b[0m                          mem  2.13 GB tm    664.38us/   226.68ms (    41.31 GFLOPS   54.4|56.7    GB/s) ['layernorm']\n",
      "beam2  : \u001b[34m1280\u001b[0m \u001b[32m8\u001b[0m \u001b[32m4\u001b[0m \u001b[31m882\u001b[0m                   :  1065.08 us  <  hc     : \u001b[34m1280\u001b[0m \u001b[32m16\u001b[0m \u001b[31m1764\u001b[0m                   :  1150.75 us\n",
      "\u001b[32m*** METAL    838\u001b[0m r_\u001b[34m1280\u001b[0m\u001b[90m_\u001b[0m\u001b[32m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m882\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm   2013.08us/   228.69ms (    54.51 GFLOPS   71.8|74.6    GB/s) ['layernorm']\n",
      "*** METAL    839 E_\u001b[90mn2\u001b[0m                                     mem  2.13 GB tm      7.00us/   228.70ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__rsub__']\n",
      "*** METAL    840 E_\u001b[90mn2\u001b[0m                                     mem  2.13 GB tm      7.38us/   228.70ms (     0.00 GFLOPS    0.0|0.0     GB/s) ['__rsub__']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[34m2\u001b[0m \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m \u001b[31m704\u001b[0m \u001b[31m2\u001b[0m \u001b[33m10\u001b[0m            :   631.00 us  <  hc     : \u001b[34m176\u001b[0m \u001b[34m5\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[31m704\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m2\u001b[0m          :  2219.37 us\n",
      "\u001b[32m*** METAL    841\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m704\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m10\u001b[0m\u001b[90m\u001b[0m                    mem  2.13 GB tm   1558.92us/   230.26ms (   549.37 GFLOPS    0.6|611.2   GB/s) ['__getitem__', '__getitem__ bw']\n",
      "beam2  : \u001b[34m352\u001b[0m \u001b[31m16\u001b[0m \u001b[33m2\u001b[0m                       :    11.96 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.83 us\n",
      "\u001b[32m*** METAL    842\u001b[0m r_\u001b[34m352\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                               mem  2.13 GB tm     12.54us/   230.28ms (     4.60 GFLOPS    1.1|2.7     GB/s) ['gather', 'gather bw', 'log_softmax bw']\n",
      "beam2  : \u001b[34m352\u001b[0m \u001b[36m2\u001b[0m \u001b[31m16\u001b[0m                       :    10.88 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    12.87 us\n",
      "\u001b[32m*** METAL    843\u001b[0m r_\u001b[34m352\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m16\u001b[0m\u001b[90m\u001b[0m                               mem  2.13 GB tm     12.08us/   230.29ms (     4.78 GFLOPS    1.2|4.7     GB/s) ['gather', 'gather bw', 'log_softmax bw']\n",
      "beam2  : \u001b[34m20\u001b[0m \u001b[36m16\u001b[0m \u001b[32m32\u001b[0m \u001b[31m22\u001b[0m                    :    17.75 us  <  hc     : \u001b[34m320\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                      :    69.87 us\n",
      "\u001b[32m*** METAL    844\u001b[0m r3_\u001b[34m20\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m\u001b[0m                           mem  2.13 GB tm     18.37us/   230.31ms (    40.68 GFLOPS   49.4|140.4   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m8\u001b[0m \u001b[32m8\u001b[0m \u001b[31m2\u001b[0m                       :     8.71 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    13.00 us\n",
      "\u001b[32m*** METAL    845\u001b[0m r_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m8\u001b[0m\u001b[90m_\u001b[0m\u001b[31m2\u001b[0m\u001b[90m\u001b[0m                               mem  2.13 GB tm     10.00us/   230.32ms (    17.46 GFLOPS   10.7|40.6    GB/s) ['gather', 'gather bw', 'log_softmax', 'log_softmax bw', '__add__']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[36m4\u001b[0m \u001b[36m8\u001b[0m \u001b[32m4\u001b[0m \u001b[31m4\u001b[0m                     :     9.17 us  <  hc     : \u001b[34m704\u001b[0m \u001b[32m16\u001b[0m                         :    13.13 us\n",
      "\u001b[32m*** METAL    846\u001b[0m r_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[32m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m4\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     10.00us/   230.33ms (    13.80 GFLOPS   10.7|20.3    GB/s) ['gather', 'gather bw', 'log_softmax', 'log_softmax bw', '__add__']\n",
      "beam2  : \u001b[34m88\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                      :     4.21 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                    :     5.58 us\n",
      "\u001b[32m*** METAL    847\u001b[0m E_\u001b[34m88\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                              mem  2.13 GB tm     10.00us/   230.34ms (    18.02 GFLOPS   15.5|38.3    GB/s) ['gather', 'gather bw', 'log_softmax', 'log_softmax bw', '__add__']\n",
      "beam2  : \u001b[34m22\u001b[0m \u001b[34m16\u001b[0m \u001b[36m32\u001b[0m                       :    11.17 us  <  hc     : \u001b[34m11\u001b[0m \u001b[36m16\u001b[0m \u001b[36m4\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m                    :    11.50 us\n",
      "\u001b[32m*** METAL    848\u001b[0m E_\u001b[34m22\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m\u001b[0m                               mem  2.13 GB tm     51.13us/   230.39ms (     3.53 GFLOPS    3.0|7.9     GB/s) ['gather', 'gather bw', 'log_softmax', 'log_softmax bw', '__add__']\n",
      "beam2  : \u001b[34m176\u001b[0m \u001b[36m32\u001b[0m \u001b[36m4\u001b[0m \u001b[33m2\u001b[0m                     :     8.63 us  <  hc     : \u001b[34m88\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m2\u001b[0m                      :     8.63 us\n",
      "\u001b[32m*** METAL    849\u001b[0m E_\u001b[34m176\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                             mem  2.13 GB tm     10.17us/   230.40ms (     0.00 GFLOPS   17.7|35.5    GB/s) ['__add__']\n",
      "beam2  : \u001b[34m2\u001b[0m \u001b[36m16\u001b[0m \u001b[32m32\u001b[0m \u001b[31m22\u001b[0m                     :    12.13 us  <  hc     : \u001b[34m32\u001b[0m \u001b[32m16\u001b[0m \u001b[31m44\u001b[0m                       :    16.87 us\n",
      "\u001b[32m*** METAL    850\u001b[0m r3_\u001b[34m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[32m32\u001b[0m\u001b[90m_\u001b[0m\u001b[31m22\u001b[0m\u001b[90m\u001b[0m                            mem  2.13 GB tm     12.88us/   230.41ms (     5.81 GFLOPS    7.0|20.0    GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m44\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m4\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m      :    39.29 us  <  tc     : \u001b[34m11\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m4\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m      :    42.00 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m32\u001b[0m               :   123.42 us\n",
      "\u001b[32m*** METAL    851\u001b[0m r_\u001b[34m44\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m4\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  2.13 GB tm     47.71us/   230.46ms (   483.53 GFLOPS   33.5|120.9   GB/s) ['linear bw']\n",
      "tc     : \u001b[34m11\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m40\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :   139.87 us  <  beam2  : \u001b[34m88\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m40\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m           :   209.25 us  <  hc     : \u001b[34m22\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m80\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m             :   291.33 us\n",
      "\u001b[32m*** METAL    852\u001b[0m r_\u001b[34m11\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m40\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m             mem  2.13 GB tm    281.62us/   230.74ms (   820.41 GFLOPS   15.8|112.6   GB/s) ['linear bw', '__add__']\n",
      "beam2  : \u001b[34m4\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m            :    90.33 us  <  tc     : \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m          :   246.12 us  <  hc     : \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m176\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m               :   267.96 us\n",
      "\u001b[32m*** METAL    853\u001b[0m r3_\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                   mem  2.13 GB tm     70.83us/   230.81ms (   329.95 GFLOPS   26.3|168.9   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "hc     : \u001b[34m8\u001b[0m \u001b[34m7\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m5\u001b[0m               :    57.17 us  <  beam2  : \u001b[34m2\u001b[0m \u001b[34m21\u001b[0m \u001b[34m8\u001b[0m \u001b[36m16\u001b[0m \u001b[36m32\u001b[0m \u001b[31m5\u001b[0m \u001b[33m4\u001b[0m               :    81.50 us\n",
      "\u001b[32m*** METAL    854\u001b[0m r_\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m5\u001b[0m\u001b[90m\u001b[0m                       mem  2.13 GB tm     97.13us/   230.91ms (    89.74 GFLOPS   57.9|122.8   GB/s) ['__getitem__', '__getitem__ bw', '__add__']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[34m10\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m       :   331.54 us  <  tc     : \u001b[34m2\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m88\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m5\u001b[0m      :   448.75 us  <  hc     : \u001b[34m10\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m176\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :   520.96 us\n",
      "\u001b[32m*** METAL    855\u001b[0m r3_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m10\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m88\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m              mem  2.13 GB tm    278.08us/   231.19ms (   840.24 GFLOPS   20.2|271.9   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "tc     : \u001b[34m21\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :   782.00 us  <  beam2  : \u001b[34m24\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m16\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m7\u001b[0m       :   807.75 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :  2125.75 us\n",
      "\u001b[32m*** METAL    856\u001b[0m r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m             mem  2.13 GB tm    736.75us/   231.92ms (   956.42 GFLOPS    8.9|123.3   GB/s) ['linear bw']\n",
      "hc     : \u001b[34m1344\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m                     :    60.79 us  <  beam2  : \u001b[34m28\u001b[0m \u001b[36m16\u001b[0m \u001b[31m512\u001b[0m \u001b[33m3\u001b[0m                    :   146.12 us\n",
      "\u001b[32m*** METAL    857\u001b[0m r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90mn1\u001b[0m                           mem  2.13 GB tm     65.00us/   231.99ms (    40.69 GFLOPS   85.0|112.5   GB/s) ['layernorm', 'layernorm bw']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m        :   605.87 us  <  tc     : \u001b[34m4\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :   916.38 us  <  hc     : \u001b[34m16\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m336\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :  1515.21 us\n",
      "\u001b[32m*** METAL    858\u001b[0m r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  2.14 GB tm   1603.21us/   233.59ms (   442.48 GFLOPS    6.7|140.9   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m672\u001b[0m \u001b[36m2\u001b[0m \u001b[32m16\u001b[0m \u001b[35m32\u001b[0m                    :    58.38 us  <  hc     : \u001b[34m1344\u001b[0m \u001b[32m16\u001b[0m \u001b[31m32\u001b[0m                     :    61.13 us\n",
      "\u001b[32m*** METAL    859\u001b[0m r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m32\u001b[0m\u001b[90m\u001b[0m                            mem  2.14 GB tm     62.00us/   233.65ms (    94.69 GFLOPS   89.1|117.9   GB/s) ['layernorm bw', 'layernorm', '__add__']\n",
      "beam2  : \u001b[34m42\u001b[0m \u001b[34m16\u001b[0m \u001b[36m16\u001b[0m \u001b[36m32\u001b[0m \u001b[33m2\u001b[0m                  :    92.87 us  <  hc     : \u001b[34m168\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                   :   108.00 us\n",
      "\u001b[32m*** METAL    860\u001b[0m E_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m                          mem  2.14 GB tm    111.00us/   233.76ms (    65.09 GFLOPS   99.4|148.8   GB/s) ['dropout', 'layernorm bw', 'layernorm', '__add__', 'dropout bw']\n",
      "beam2  : \u001b[34m42\u001b[0m \u001b[34m128\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m      :  1550.75 us  <  tc     : \u001b[34m21\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :  1894.42 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m           :  2921.50 us\n",
      "\u001b[32m*** METAL    861\u001b[0m r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  2.14 GB tm   2869.62us/   236.63ms (   982.21 GFLOPS    6.3|188.0   GB/s) ['linear bw']\n",
      "beam2  : \u001b[34m84\u001b[0m \u001b[34m128\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m      :  1636.75 us  <  tc     : \u001b[34m42\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m64\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :  2058.13 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m128\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m           :  2151.79 us\n",
      "\u001b[32m*** METAL    862\u001b[0m r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  2.15 GB tm   3753.79us/   240.39ms (   771.39 GFLOPS    7.7|195.0   GB/s) ['linear', 'gelu', 'gelu bw', '__add__']\n",
      "beam2  : \u001b[34m16\u001b[0m \u001b[34m256\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m       :  1886.13 us  <  tc     : \u001b[34m4\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m     :  3010.50 us  <  hc     : \u001b[34m16\u001b[0m \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m336\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m           :  4075.29 us\n",
      "\u001b[32m*** METAL    863\u001b[0m r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m              mem  2.16 GB tm   5425.50us/   245.81ms (   523.01 GFLOPS    6.4|166.5   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "tc     : \u001b[34m21\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m256\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m    :  1597.38 us  <  beam2  : \u001b[34m2\u001b[0m \u001b[34m168\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m8\u001b[0m \u001b[31m256\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m      :  2117.17 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m512\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :  3618.79 us\n",
      "\u001b[32m*** METAL    864\u001b[0m r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m            mem  2.16 GB tm   2960.25us/   248.77ms (   952.14 GFLOPS    6.1|119.9   GB/s) ['linear bw']\n",
      "*** METAL    865 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90mn1\u001b[0m                           mem  2.16 GB tm     57.88us/   248.83ms (    45.70 GFLOPS   95.5|126.3   GB/s) ['layernorm', 'layernorm bw']\n",
      "tc     : \u001b[34m16\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m    :  1142.25 us  <  beam2  : \u001b[34m64\u001b[0m \u001b[34m32\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m      :  1237.54 us  <  hc     : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m336\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :  2258.12 us\n",
      "\u001b[32m*** METAL    866\u001b[0m r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m           mem  2.17 GB tm   2987.17us/   251.82ms (   949.89 GFLOPS   11.6|125.1   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    867 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m32\u001b[0m\u001b[90m\u001b[0m                            mem  2.17 GB tm     70.04us/   251.89ms (    83.82 GFLOPS   78.9|104.4   GB/s) ['layernorm bw', 'layernorm', '__add__']\n",
      "hc     : \u001b[34m168\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                   :   144.50 us  <  beam2  : \u001b[34m672\u001b[0m \u001b[34m32\u001b[0m \u001b[36m16\u001b[0m \u001b[36m2\u001b[0m                    :   260.87 us\n",
      "\u001b[32m*** METAL    868\u001b[0m E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn1\u001b[0m                         mem  2.17 GB tm    154.50us/   252.04ms (    77.94 GFLOPS  107.2|142.5   GB/s) ['dropout', 'layernorm bw', 'layernorm', '__add__', 'dropout bw']\n",
      "*** METAL    869 r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m             mem  2.17 GB tm   1366.00us/   253.41ms (   515.84 GFLOPS    4.8|66.5    GB/s) ['linear bw']\n",
      "*** METAL    870 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  2.17 GB tm    903.25us/   254.31ms (   785.38 GFLOPS   11.9|250.0   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m8\u001b[0m \u001b[33m3\u001b[0m         :    98.62 us  <  hc     : \u001b[34m64\u001b[0m \u001b[34m7\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m \u001b[33m3\u001b[0m \u001b[35m21\u001b[0m               :   348.87 us\n",
      "\u001b[32m*** METAL    871\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  2.17 GB tm     90.83us/   254.40ms (   421.17 GFLOPS   70.5|82.3    GB/s) ['__matmul__ bw']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[34m3\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[35m8\u001b[0m         :    94.87 us  <  hc     : \u001b[34m4\u001b[0m \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m16\u001b[0m \u001b[36m8\u001b[0m \u001b[31m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m3\u001b[0m \u001b[35m4\u001b[0m            :   642.25 us\n",
      "\u001b[32m*** METAL    872\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m\u001b[0m                 mem  2.17 GB tm    245.67us/   254.65ms (   170.46 GFLOPS   52.2|86.4    GB/s) ['dropout', '__matmul__ bw', 'dropout bw']\n",
      "hc     : \u001b[34m112\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m \u001b[35m21\u001b[0m                    :    39.75 us  <  beam2  : \u001b[34m84\u001b[0m \u001b[36m32\u001b[0m \u001b[31m21\u001b[0m \u001b[33m4\u001b[0m                     :    74.25 us\n",
      "\u001b[32m*** METAL    873\u001b[0m r_\u001b[34m112\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                            mem  2.17 GB tm     34.04us/   254.68ms (    47.06 GFLOPS   56.9|56.9    GB/s) ['softmax bw', 'softmax']\n",
      "beam2  : \u001b[34m672\u001b[0m \u001b[36m16\u001b[0m \u001b[35m21\u001b[0m                      :    26.25 us  <  hc     : \u001b[34m112\u001b[0m \u001b[36m32\u001b[0m \u001b[33m3\u001b[0m \u001b[35m21\u001b[0m                    :    51.17 us\n",
      "\u001b[32m*** METAL    874\u001b[0m r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90mn1\u001b[0m                            mem  2.17 GB tm     29.83us/   254.71ms (    60.91 GFLOPS   66.3|66.3    GB/s) ['softmax', 'softmax bw', '__add__']\n",
      "beam2  : \u001b[34m32\u001b[0m \u001b[34m21\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m7\u001b[0m                   :    56.00 us  <  hc     : \u001b[34m4\u001b[0m \u001b[34m7\u001b[0m \u001b[34m7\u001b[0m \u001b[36m32\u001b[0m \u001b[36m3\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m                 :   164.50 us\n",
      "\u001b[32m*** METAL    875\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                           mem  2.17 GB tm     57.87us/   254.77ms (    56.29 GFLOPS   50.5|61.9    GB/s) ['softmax', 'softmax bw', '__add__', 'masked_fill bw', '__mul__ bw']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[34m3\u001b[0m \u001b[34m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m2\u001b[0m       :    56.87 us  <  hc     : \u001b[34m64\u001b[0m \u001b[34m7\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m21\u001b[0m               :   133.79 us\n",
      "\u001b[32m*** METAL    876\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  2.17 GB tm     94.62us/   254.87ms (   465.42 GFLOPS  125.9|182.8   GB/s) ['__matmul__ bw']\n",
      "beam2  : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[34m3\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m3\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m8\u001b[0m         :    73.67 us  <  hc     : \u001b[34m64\u001b[0m \u001b[34m7\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m21\u001b[0m               :   535.00 us\n",
      "\u001b[32m*** METAL    877\u001b[0m r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90mn1\u001b[0m               mem  2.17 GB tm    165.75us/   255.03ms (   236.34 GFLOPS   71.9|83.0    GB/s) ['__matmul__ bw']\n",
      "hc     : \u001b[34m32\u001b[0m \u001b[34m7\u001b[0m \u001b[34m24\u001b[0m \u001b[36m2\u001b[0m \u001b[36m3\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m               :   139.12 us  <  beam2  : \u001b[34m4\u001b[0m \u001b[34m7\u001b[0m \u001b[34m192\u001b[0m \u001b[36m16\u001b[0m \u001b[36m3\u001b[0m \u001b[33m4\u001b[0m \u001b[33m2\u001b[0m               :   141.17 us\n",
      "\u001b[32m*** METAL    878\u001b[0m E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[34m24\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                       mem  2.17 GB tm    396.50us/   255.43ms (    18.22 GFLOPS   41.7|83.3    GB/s) ['__add__']\n",
      "tc     : \u001b[34m21\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m192\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m    :   813.62 us  <  beam2  : \u001b[34m168\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m192\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m       :  1094.67 us  <  hc     : \u001b[34m56\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m384\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :  2720.87 us\n",
      "\u001b[32m*** METAL    879\u001b[0m r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m192\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m            mem  2.17 GB tm   2938.25us/   258.37ms (   719.45 GFLOPS    4.8|90.9    GB/s) ['linear bw']\n",
      "*** METAL    880 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90mn1\u001b[0m                           mem  2.17 GB tm     63.37us/   258.43ms (    41.74 GFLOPS   87.2|115.4   GB/s) ['layernorm', 'layernorm bw']\n",
      "beam2  : \u001b[34m48\u001b[0m \u001b[34m64\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m        :  1783.37 us  <  hc     : \u001b[34m64\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[31m336\u001b[0m \u001b[33m3\u001b[0m \u001b[33m4\u001b[0m \u001b[35m4\u001b[0m            :  1974.67 us  <  tc     : \u001b[34m12\u001b[0m \u001b[34m16\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[36m2\u001b[0m \u001b[36m2\u001b[0m \u001b[36m4\u001b[0m \u001b[31m168\u001b[0m \u001b[35m8\u001b[0m \u001b[33m2\u001b[0m \u001b[33m4\u001b[0m \u001b[33m4\u001b[0m    :  2301.08 us\n",
      "\u001b[32m*** METAL    881\u001b[0m r3_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  2.18 GB tm   3827.50us/   262.26ms (   556.02 GFLOPS    7.0|177.0   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    882 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m32\u001b[0m\u001b[90m\u001b[0m                            mem  2.18 GB tm     88.96us/   262.35ms (    65.99 GFLOPS   62.1|82.2    GB/s) ['layernorm bw', 'layernorm', '__add__']\n",
      "hc     : \u001b[34m168\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                   :   191.71 us  <  beam2  : \u001b[34m168\u001b[0m \u001b[34m32\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m                    :   210.75 us\n",
      "\u001b[32m*** METAL    883\u001b[0m E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn2\u001b[0m                         mem  2.18 GB tm    201.46us/   262.55ms (    84.54 GFLOPS  109.6|150.3   GB/s) ['dropout', 'layernorm bw', 'layernorm', '__add__', 'dropout bw']\n",
      "*** METAL    884 r_\u001b[34m42\u001b[0m\u001b[90m_\u001b[0m\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  2.18 GB tm   3555.46us/   266.10ms (   792.75 GFLOPS    5.1|151.7   GB/s) ['linear bw']\n",
      "*** METAL    885 r_\u001b[34m84\u001b[0m\u001b[90m_\u001b[0m\u001b[34m128\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m              mem  2.19 GB tm   2918.79us/   269.02ms (   992.07 GFLOPS    9.9|250.8   GB/s) ['linear', 'gelu', 'gelu bw', '__add__']\n",
      "*** METAL    886 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m256\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m              mem  2.20 GB tm   5318.63us/   274.34ms (   533.52 GFLOPS    6.5|169.8   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    887 r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m256\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m            mem  2.20 GB tm   3728.92us/   278.07ms (   755.87 GFLOPS    4.8|95.2    GB/s) ['linear bw']\n",
      "*** METAL    888 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90mn1\u001b[0m                           mem  2.20 GB tm     80.38us/   278.15ms (    32.91 GFLOPS   68.8|91.0    GB/s) ['layernorm', 'layernorm bw']\n",
      "*** METAL    889 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m           mem  2.21 GB tm   1660.58us/   279.81ms (  1708.72 GFLOPS   20.9|225.0   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    890 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m32\u001b[0m\u001b[90m\u001b[0m                            mem  2.21 GB tm     59.00us/   279.87ms (    99.50 GFLOPS   93.7|123.9   GB/s) ['layernorm bw', 'layernorm', '__add__']\n",
      "hc     : \u001b[34m168\u001b[0m \u001b[34m8\u001b[0m \u001b[36m8\u001b[0m \u001b[36m16\u001b[0m \u001b[33m4\u001b[0m                   :   249.12 us  <  beam2  : \u001b[34m1344\u001b[0m \u001b[34m16\u001b[0m \u001b[36m16\u001b[0m \u001b[33m2\u001b[0m                   :   283.88 us\n",
      "\u001b[32m*** METAL    891\u001b[0m E_\u001b[34m168\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90mn3\u001b[0m                         mem  2.21 GB tm    234.33us/   280.10ms (    93.97 GFLOPS  117.8|164.4   GB/s) ['dropout', 'layernorm bw', 'layernorm', '__add__', 'dropout bw']\n",
      "*** METAL    892 r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m64\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m             mem  2.21 GB tm    737.88us/   280.84ms (   954.96 GFLOPS    8.9|123.1   GB/s) ['linear bw']\n",
      "*** METAL    893 r3_\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  2.21 GB tm    867.38us/   281.71ms (   817.86 GFLOPS   12.4|260.4   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    894 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m\u001b[0m                 mem  2.21 GB tm    183.00us/   281.89ms (   209.05 GFLOPS   35.0|40.8    GB/s) ['__matmul__ bw']\n",
      "*** METAL    895 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m\u001b[0m                 mem  2.21 GB tm     68.79us/   281.96ms (   608.76 GFLOPS  186.3|308.7   GB/s) ['dropout', '__matmul__ bw', 'dropout bw']\n",
      "*** METAL    896 r_\u001b[34m112\u001b[0m\u001b[90m_\u001b[0m\u001b[36m32\u001b[0m\u001b[90m_\u001b[0m\u001b[33m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90m\u001b[0m                            mem  2.21 GB tm     18.42us/   281.98ms (    86.99 GFLOPS  105.1|105.1   GB/s) ['softmax bw', 'softmax']\n",
      "*** METAL    897 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m21\u001b[0m\u001b[90mn1\u001b[0m                            mem  2.21 GB tm     13.50us/   281.99ms (   134.60 GFLOPS  146.5|146.5   GB/s) ['softmax', 'softmax bw', '__add__']\n",
      "*** METAL    898 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[33m7\u001b[0m\u001b[90m\u001b[0m                           mem  2.21 GB tm     30.96us/   282.02ms (   105.23 GFLOPS   94.5|115.7   GB/s) ['softmax', 'softmax bw', '__add__', 'masked_fill bw', '__mul__ bw']\n",
      "*** METAL    899 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[34m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m\u001b[0m               mem  2.21 GB tm     53.12us/   282.08ms (   828.99 GFLOPS  224.2|325.7   GB/s) ['__matmul__ bw']\n",
      "*** METAL    900 r_\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[34m8\u001b[0m\u001b[90m_\u001b[0m\u001b[34m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m3\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m8\u001b[0m\u001b[90mn1\u001b[0m               mem  2.21 GB tm     45.00us/   282.12ms (   870.54 GFLOPS  264.7|305.8   GB/s) ['__matmul__ bw']\n",
      "*** METAL    901 E_\u001b[34m32\u001b[0m\u001b[90m_\u001b[0m\u001b[34m7\u001b[0m\u001b[90m_\u001b[0m\u001b[34m24\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m3\u001b[0m\u001b[90m_\u001b[0m\u001b[36m16\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m                       mem  2.21 GB tm    102.75us/   282.22ms (    70.32 GFLOPS  160.7|321.5   GB/s) ['__add__']\n",
      "*** METAL    902 r_\u001b[34m21\u001b[0m\u001b[90m_\u001b[0m\u001b[34m16\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m192\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m            mem  2.21 GB tm   1189.04us/   283.41ms (  1777.84 GFLOPS   11.9|224.5   GB/s) ['linear bw']\n",
      "*** METAL    903 r_\u001b[34m1344\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[31m32\u001b[0m\u001b[90mn1\u001b[0m                           mem  2.21 GB tm     40.38us/   283.45ms (    65.51 GFLOPS  136.9|181.1   GB/s) ['layernorm', 'layernorm bw']\n",
      "*** METAL    904 r3_\u001b[34m48\u001b[0m\u001b[90m_\u001b[0m\u001b[34m64\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m4\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[31m168\u001b[0m\u001b[90m_\u001b[0m\u001b[35m8\u001b[0m\u001b[90m_\u001b[0m\u001b[33m2\u001b[0m\u001b[90m_\u001b[0m\u001b[33m4\u001b[0m\u001b[90m\u001b[0m               mem  2.21 GB tm   3293.17us/   286.75ms (   646.24 GFLOPS    8.1|205.7   GB/s) ['__rmul__', 'linear bw', '__add__', '__mul__', '__truediv__', 'sqrt', '__sub__']\n",
      "*** METAL    905 r_\u001b[34m672\u001b[0m\u001b[90m_\u001b[0m\u001b[36m2\u001b[0m\u001b[90m_\u001b[0m\u001b[32m16\u001b[0m\u001b[90m_\u001b[0m\u001b[35m32\u001b[0m\u001b[90m\u001b[0m                            mem  2.21 GB tm     65.00us/   286.81ms (    90.32 GFLOPS   85.0|112.5   GB/s) ['layernorm bw', 'layernorm', '__add__']\n",
      "error lowering MetaOps.KERNEL\n",
      "tensor operations:\n",
      "[dropout, layernorm bw, layernorm, __add__, dropout bw]\n"
     ]
    },
    {
     "ename": "CompileError",
     "evalue": "Error Domain=MTLLibraryErrorDomain Code=3 \"program_source:3:896: error: no 'buffer' resource location available for 'data31'\nkernel void E_1344_512(device float* data0, const device float* data1, const device float* data2, const device float* data3, const device float* data4, const device float* data5, const device float* data6, const device float* data7, const device float* data8, const device float* data9, const device float* data10, const device float* data11, const device float* data12, const device float* data13, const device float* data14, const device float* data15, const device float* data16, const device float* data17, const device float* data18, const device float* data19, const device float* data20, const device float* data21, const device float* data22, const device float* data23, const device float* data24, const device float* data25, const device float* data26, const device float* data27, const device float* data28, const device float* data29, const device float* data30, const device float* data31, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {\n                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               ^\n\" UserInfo={NSLocalizedDescription=program_source:3:896: error: no 'buffer' resource location available for 'data31'\nkernel void E_1344_512(device float* data0, const device float* data1, const device float* data2, const device float* data3, const device float* data4, const device float* data5, const device float* data6, const device float* data7, const device float* data8, const device float* data9, const device float* data10, const device float* data11, const device float* data12, const device float* data13, const device float* data14, const device float* data15, const device float* data16, const device float* data17, const device float* data18, const device float* data19, const device float* data20, const device float* data21, const device float* data22, const device float* data23, const device float* data24, const device float* data25, const device float* data26, const device float* data27, const device float* data28, const device float* data29, const device float* data30, const device float* data31, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {\n                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               ^\n}",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/runtime/ops_metal.py:25\u001b[0m, in \u001b[0;36mMetalCompiler.compile\u001b[0;34m(self, src)\u001b[0m\n\u001b[1;32m     24\u001b[0m options\u001b[38;5;241m.\u001b[39msetFastMathEnabled_(getenv(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMETAL_FAST_MATH\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m: library \u001b[38;5;241m=\u001b[39m \u001b[43munwrap2\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnewLibraryWithSource_options_error_\u001b[49m\u001b[43m(\u001b[49m\u001b[43msrc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAssertionError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e: \u001b[38;5;28;01mraise\u001b[39;00m CompileError(e) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/helpers.py:54\u001b[0m, in \u001b[0;36munwrap2\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m     53\u001b[0m ret, err \u001b[38;5;241m=\u001b[39m x\n\u001b[0;32m---> 54\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m err \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28mstr\u001b[39m(err)\n\u001b[1;32m     55\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ret\n",
      "\u001b[0;31mAssertionError\u001b[0m: Error Domain=MTLLibraryErrorDomain Code=3 \"program_source:3:896: error: no 'buffer' resource location available for 'data31'\nkernel void E_1344_512(device float* data0, const device float* data1, const device float* data2, const device float* data3, const device float* data4, const device float* data5, const device float* data6, const device float* data7, const device float* data8, const device float* data9, const device float* data10, const device float* data11, const device float* data12, const device float* data13, const device float* data14, const device float* data15, const device float* data16, const device float* data17, const device float* data18, const device float* data19, const device float* data20, const device float* data21, const device float* data22, const device float* data23, const device float* data24, const device float* data25, const device float* data26, const device float* data27, const device float* data28, const device float* data29, const device float* data30, const device float* data31, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {\n                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               ^\n\" UserInfo={NSLocalizedDescription=program_source:3:896: error: no 'buffer' resource location available for 'data31'\nkernel void E_1344_512(device float* data0, const device float* data1, const device float* data2, const device float* data3, const device float* data4, const device float* data5, const device float* data6, const device float* data7, const device float* data8, const device float* data9, const device float* data10, const device float* data11, const device float* data12, const device float* data13, const device float* data14, const device float* data15, const device float* data16, const device float* data17, const device float* data18, const device float* data19, const device float* data20, const device float* data21, const device float* data22, const device float* data23, const device float* data24, const device float* data25, const device float* data26, const device float* data27, const device float* data28, const device float* data29, const device float* data30, const device float* data31, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {\n                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               ^\n}",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mCompileError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[22], line 152\u001b[0m\n\u001b[1;32m    150\u001b[0m     start \u001b[38;5;241m=\u001b[39m timeit\u001b[38;5;241m.\u001b[39mdefault_timer()\n\u001b[1;32m    151\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe start time is :\u001b[39m\u001b[38;5;124m\"\u001b[39m, start)\n\u001b[0;32m--> 152\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_step_main\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    153\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mobservation.image\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    154\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mobservation.state\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    155\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43maction\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    156\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mepisode_index\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    157\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mframe_index\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    158\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mtimestamp\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    159\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mnext.reward\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    160\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mnext.done\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    161\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mnext.success\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    162\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mindex\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    163\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mobservation.image_is_pad\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    164\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mobservation.state_is_pad\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    165\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43maction_is_pad\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    166\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    167\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe difference of time is :\u001b[39m\u001b[38;5;124m\"\u001b[39m, timeit\u001b[38;5;241m.\u001b[39mdefault_timer() \u001b[38;5;241m-\u001b[39m start)\n\u001b[1;32m    168\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "File \u001b[0;32m/opt/homebrew/Cellar/python@3.12/3.12.6/Frameworks/Python.framework/Versions/3.12/lib/python3.12/contextlib.py:81\u001b[0m, in \u001b[0;36mContextDecorator.__call__.<locals>.inner\u001b[0;34m(*args, **kwds)\u001b[0m\n\u001b[1;32m     78\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[1;32m     79\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwds):\n\u001b[1;32m     80\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_recreate_cm():\n\u001b[0;32m---> 81\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[22], line 121\u001b[0m, in \u001b[0;36mtrain_step_main\u001b[0;34m(observation_image, observation_state, action, episode_index, frame_index, timestamp, next_reward, next_done, next_success, index, observation_image_is_pad, observation_state_is_pad, action_is_pad)\u001b[0m\n\u001b[1;32m    119\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[1;32m    120\u001b[0m \u001b[38;5;66;03m# grad_norm = clip_grad_norm_(vqvae_params, 10.0) if not policy.vqbet.action_head.vqvae_model.discretized else None\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[43mopt_decay\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    122\u001b[0m opt_no_decay\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m    123\u001b[0m lr_scheduler\u001b[38;5;241m.\u001b[39mstep()\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/nn/optim.py:34\u001b[0m, in \u001b[0;36mOptimizer.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m     31\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;124;03m  Performs a single optimization step.\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[38;5;124;03m  \"\"\"\u001b[39;00m\n\u001b[0;32m---> 34\u001b[0m   \u001b[43mTensor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mschedule_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/tensor.py:3256\u001b[0m, in \u001b[0;36m_metadata_wrapper.<locals>._wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   3253\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m: caller \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   3255\u001b[0m token \u001b[38;5;241m=\u001b[39m _METADATA\u001b[38;5;241m.\u001b[39mset(Metadata(name\u001b[38;5;241m=\u001b[39mfn\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, caller\u001b[38;5;241m=\u001b[39mcaller))\n\u001b[0;32m-> 3256\u001b[0m ret \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   3257\u001b[0m _METADATA\u001b[38;5;241m.\u001b[39mreset(token)\n\u001b[1;32m   3258\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ret\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/tensor.py:204\u001b[0m, in \u001b[0;36mTensor.realize\u001b[0;34m(self, do_update_stats, *lst)\u001b[0m\n\u001b[1;32m    202\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrealize\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39mlst:Tensor, do_update_stats\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[1;32m    203\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Triggers the computation needed to create these Tensor(s).\"\"\"\u001b[39;00m\n\u001b[0;32m--> 204\u001b[0m   \u001b[43mrun_schedule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mschedule_with_vars\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mlst\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdo_update_stats\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdo_update_stats\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    205\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/realize.py:221\u001b[0m, in \u001b[0;36mrun_schedule\u001b[0;34m(schedule, var_vals, do_update_stats)\u001b[0m\n\u001b[1;32m    220\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrun_schedule\u001b[39m(schedule:List[ScheduleItem], var_vals:Optional[Dict[Variable, \u001b[38;5;28mint\u001b[39m]]\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, do_update_stats\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[0;32m--> 221\u001b[0m   \u001b[38;5;28;01mfor\u001b[39;00m ei \u001b[38;5;129;01min\u001b[39;00m lower_schedule(schedule):\n\u001b[1;32m    222\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(capturing) \u001b[38;5;129;01mand\u001b[39;00m CAPTURING: capturing[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39madd(ei)\n\u001b[1;32m    223\u001b[0m     ei\u001b[38;5;241m.\u001b[39mrun(var_vals, do_update_stats\u001b[38;5;241m=\u001b[39mdo_update_stats)\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/realize.py:214\u001b[0m, in \u001b[0;36mlower_schedule\u001b[0;34m(schedule)\u001b[0m\n\u001b[1;32m    212\u001b[0m   \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtensor operations:\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    213\u001b[0m   pprint\u001b[38;5;241m.\u001b[39mpprint(si\u001b[38;5;241m.\u001b[39mmetadata, indent\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m)\n\u001b[0;32m--> 214\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m e\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/realize.py:208\u001b[0m, in \u001b[0;36mlower_schedule\u001b[0;34m(schedule)\u001b[0m\n\u001b[1;32m    206\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(schedule):\n\u001b[1;32m    207\u001b[0m   si \u001b[38;5;241m=\u001b[39m schedule\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m--> 208\u001b[0m   \u001b[38;5;28;01mtry\u001b[39;00m: \u001b[38;5;28;01myield\u001b[39;00m \u001b[43mlower_schedule_item\u001b[49m\u001b[43m(\u001b[49m\u001b[43msi\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    209\u001b[0m   \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    210\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m DEBUG \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m2\u001b[39m:\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/realize.py:192\u001b[0m, in \u001b[0;36mlower_schedule_item\u001b[0;34m(si)\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mset\u001b[39m(x\u001b[38;5;241m.\u001b[39mdevice \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m si\u001b[38;5;241m.\u001b[39mbufs)) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m (si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39mop \u001b[38;5;129;01mis\u001b[39;00m MetaOps\u001b[38;5;241m.\u001b[39mEXT \u001b[38;5;129;01mand\u001b[39;00m si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39marg[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m MetaOps\u001b[38;5;241m.\u001b[39mCOPY) \u001b[38;5;129;01mor\u001b[39;00m getenv(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUSE_COPY_KERNEL\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    191\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39mop \u001b[38;5;129;01mis\u001b[39;00m MetaOps\u001b[38;5;241m.\u001b[39mKERNEL:\n\u001b[0;32m--> 192\u001b[0m   runner \u001b[38;5;241m=\u001b[39m \u001b[43mget_runner\u001b[49m\u001b[43m(\u001b[49m\u001b[43msi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutputs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mast\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    193\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m ExecItem(runner, [si\u001b[38;5;241m.\u001b[39mbufs[x] \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m runner\u001b[38;5;241m.\u001b[39mp\u001b[38;5;241m.\u001b[39mglobals], si\u001b[38;5;241m.\u001b[39mmetadata)\n\u001b[1;32m    194\u001b[0m out, (op, arg) \u001b[38;5;241m=\u001b[39m si\u001b[38;5;241m.\u001b[39moutputs[\u001b[38;5;241m0\u001b[39m], si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39marg\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/realize.py:157\u001b[0m, in \u001b[0;36mget_runner\u001b[0;34m(dname, ast)\u001b[0m\n\u001b[1;32m    155\u001b[0m   method_cache[ckey] \u001b[38;5;241m=\u001b[39m ret \u001b[38;5;241m=\u001b[39m CompiledRunner(replace(bret\u001b[38;5;241m.\u001b[39mp, dname\u001b[38;5;241m=\u001b[39mdname), bret\u001b[38;5;241m.\u001b[39mlib)\n\u001b[1;32m    156\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 157\u001b[0m   prg: Program \u001b[38;5;241m=\u001b[39m \u001b[43mget_kernel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mDevice\u001b[49m\u001b[43m[\u001b[49m\u001b[43mdname\u001b[49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrenderer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mast\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mto_program()\n\u001b[1;32m    158\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m getenv(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFUZZ_UOPS\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m    159\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtest\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mexternal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfuzz_uops\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m UOpsFuzzerRunner\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/realize.py:36\u001b[0m, in \u001b[0;36mget_kernel\u001b[0;34m(renderer, ast)\u001b[0m\n\u001b[1;32m     34\u001b[0m lins: List[Tuple[\u001b[38;5;28mstr\u001b[39m, Kernel]] \u001b[38;5;241m=\u001b[39m [(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbeam\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mBEAM\u001b[38;5;241m.\u001b[39mvalue\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, k), ((\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtc\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m used_tensor_cores \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhc\u001b[39m\u001b[38;5;124m\"\u001b[39m), k_opt)]\n\u001b[1;32m     35\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m used_tensor_cores: lins\u001b[38;5;241m.\u001b[39mappend((\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhc\u001b[39m\u001b[38;5;124m\"\u001b[39m, Kernel(ast, opts\u001b[38;5;241m=\u001b[39mrenderer)\u001b[38;5;241m.\u001b[39mhand_coded_optimizations()))\n\u001b[0;32m---> 36\u001b[0m timed \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msorted\u001b[39m([(nm, tk, \u001b[43mtime_linearizer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrawbufs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mallow_test_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclear_l2\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m) \u001b[38;5;28;01mfor\u001b[39;00m nm, tk \u001b[38;5;129;01min\u001b[39;00m lins], key\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mlambda\u001b[39;00m x: x[\u001b[38;5;241m2\u001b[39m])\n\u001b[1;32m     37\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m DEBUG \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m: \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m  <  \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnm\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m6s\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mlin\u001b[38;5;241m.\u001b[39mcolored_shape(\u001b[38;5;241m30\u001b[39m,\u001b[38;5;250m \u001b[39mdense\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtm\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m1e6\u001b[39m\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m8.2f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m us\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m nm, lin, tm \u001b[38;5;129;01min\u001b[39;00m timed))\n\u001b[1;32m     38\u001b[0m k \u001b[38;5;241m=\u001b[39m timed[\u001b[38;5;241m0\u001b[39m][\u001b[38;5;241m1\u001b[39m]\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/engine/search.py:202\u001b[0m, in \u001b[0;36mtime_linearizer\u001b[0;34m(lin, rawbufs, allow_test_size, max_global_size, cnt, disable_cache, clear_l2)\u001b[0m\n\u001b[1;32m    200\u001b[0m var_vals \u001b[38;5;241m=\u001b[39m {k:(k\u001b[38;5;241m.\u001b[39mmax\u001b[38;5;241m+\u001b[39mk\u001b[38;5;241m.\u001b[39mmin)\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m2\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m lin\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39mvars()}\n\u001b[1;32m    201\u001b[0m p \u001b[38;5;241m=\u001b[39m lin\u001b[38;5;241m.\u001b[39mto_program()\n\u001b[0;32m--> 202\u001b[0m tms \u001b[38;5;241m=\u001b[39m _time_program(p, \u001b[43mdev\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompiler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompile\u001b[49m\u001b[43m(\u001b[49m\u001b[43mp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msrc\u001b[49m\u001b[43m)\u001b[49m, var_vals, rawbufs,\n\u001b[1;32m    203\u001b[0m                     max_global_size\u001b[38;5;241m=\u001b[39mmax_global_size \u001b[38;5;28;01mif\u001b[39;00m allow_test_size \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m, clear_l2\u001b[38;5;241m=\u001b[39mclear_l2, cnt\u001b[38;5;241m=\u001b[39mcnt, name\u001b[38;5;241m=\u001b[39mto_function_name(lin\u001b[38;5;241m.\u001b[39mname))\n\u001b[1;32m    205\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m CACHELEVEL \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m2\u001b[39m: diskcache_put(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtime_linearizer\u001b[39m\u001b[38;5;124m\"\u001b[39m, key, tms)\n\u001b[1;32m    206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmin\u001b[39m(tms)\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.12/site-packages/tinygrad/runtime/ops_metal.py:26\u001b[0m, in \u001b[0;36mMetalCompiler.compile\u001b[0;34m(self, src)\u001b[0m\n\u001b[1;32m     24\u001b[0m options\u001b[38;5;241m.\u001b[39msetFastMathEnabled_(getenv(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMETAL_FAST_MATH\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m: library \u001b[38;5;241m=\u001b[39m unwrap2(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice\u001b[38;5;241m.\u001b[39mdevice\u001b[38;5;241m.\u001b[39mnewLibraryWithSource_options_error_(src, options, \u001b[38;5;28;01mNone\u001b[39;00m))\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAssertionError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e: \u001b[38;5;28;01mraise\u001b[39;00m CompileError(e) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m     27\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m library\u001b[38;5;241m.\u001b[39mlibraryDataContents()\u001b[38;5;241m.\u001b[39mbytes()\u001b[38;5;241m.\u001b[39mtobytes()\n",
      "\u001b[0;31mCompileError\u001b[0m: Error Domain=MTLLibraryErrorDomain Code=3 \"program_source:3:896: error: no 'buffer' resource location available for 'data31'\nkernel void E_1344_512(device float* data0, const device float* data1, const device float* data2, const device float* data3, const device float* data4, const device float* data5, const device float* data6, const device float* data7, const device float* data8, const device float* data9, const device float* data10, const device float* data11, const device float* data12, const device float* data13, const device float* data14, const device float* data15, const device float* data16, const device float* data17, const device float* data18, const device float* data19, const device float* data20, const device float* data21, const device float* data22, const device float* data23, const device float* data24, const device float* data25, const device float* data26, const device float* data27, const device float* data28, const device float* data29, const device float* data30, const device float* data31, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {\n                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               ^\n\" UserInfo={NSLocalizedDescription=program_source:3:896: error: no 'buffer' resource location available for 'data31'\nkernel void E_1344_512(device float* data0, const device float* data1, const device float* data2, const device float* data3, const device float* data4, const device float* data5, const device float* data6, const device float* data7, const device float* data8, const device float* data9, const device float* data10, const device float* data11, const device float* data12, const device float* data13, const device float* data14, const device float* data15, const device float* data16, const device float* data17, const device float* data18, const device float* data19, const device float* data20, const device float* data21, const device float* data22, const device float* data23, const device float* data24, const device float* data25, const device float* data26, const device float* data27, const device float* data28, const device float* data29, const device float* data30, const device float* data31, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {\n                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               ^\n}"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from lerobot.common.datasets.lerobot_dataset import LeRobotDataset\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import tinygrad\n",
    "from tinygrad import Tensor, nn, dtypes, TinyJit\n",
    "\n",
    "import timeit\n",
    "\n",
    "from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict\n",
    "\n",
    "# Start of training code\n",
    "\n",
    "# Create a directory to store the training checkpoint.\n",
    "output_directory = Path(\"outputs/train/example_pusht_vqbet\")\n",
    "output_directory.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Number of offline training steps (we'll only do offline training for this example.)\n",
    "# Adjust as you prefer. 5000 steps are needed to get something worth evaluating.\n",
    "training_steps = 250000\n",
    "log_freq = 2\n",
    "\n",
    "# Set up the dataset.\n",
    "delta_timestamps = {\n",
    "    # Load the previous image and state at -0.1 seconds before current frame,\n",
    "    # then load current image and state corresponding to 0.0 second.\n",
    "    \"observation.image\": [-0.2, -0.15, -0.1, -0.05, 0.0],\n",
    "    \"observation.state\": [-0.2, -0.15, -0.1, -0.05, 0.0],\n",
    "    # Load the previous action (-0.1), the next action to be executed (0.0),\n",
    "    # and 14 future actions with a 0.1 seconds spacing. All these actions will be\n",
    "    # used to supervise the policy.\n",
    "    \"action\": [-0.2, -0.15, -0.1, -0.05, 0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5],\n",
    "}\n",
    "dataset = LeRobotDataset(\"lerobot/pusht\", delta_timestamps=delta_timestamps)\n",
    "\n",
    "# Set up the the policy.\n",
    "# Policies are initialized with a configuration class, in this case `VqBeTConfig`.\n",
    "# For this example, no arguments need to be passed because the defaults are set up for PushT.\n",
    "# If you're doing something different, you will likely need to change at least some of the defaults.\n",
    "cfg = VQBeTConfig()\n",
    "policy = VQBeTPolicy(cfg, dataset_stats=dataset.stats)\n",
    "\n",
    "opt_decay, opt_vq_vae, opt_no_decay, decay_params, vqvae_params, no_decay_params = VQBeTOptimizer.create(policy, cfg)\n",
    "opts = OptimizerGroup(opt_decay, opt_no_decay)\n",
    "lr_scheduler = VQBeTScheduler(opts, cfg)\n",
    "\n",
    "@Tensor.train()\n",
    "def train_step_vqvae(\n",
    "    observation_image: Tensor,\n",
    "    observation_state: Tensor,\n",
    "    action: Tensor,\n",
    "    episode_index: Tensor,\n",
    "    frame_index: Tensor,\n",
    "    timestamp: Tensor,\n",
    "    next_reward: Tensor,\n",
    "    next_done: Tensor,\n",
    "    next_success: Tensor,\n",
    "    index: Tensor,\n",
    "    observation_image_is_pad: Tensor,\n",
    "    observation_state_is_pad: Tensor,\n",
    "    action_is_pad: Tensor\n",
    "):\n",
    "    batch = {\n",
    "        'observation.image': observation_image,\n",
    "        'observation.state': observation_state,\n",
    "        'action': action,\n",
    "        'episode_index': episode_index,\n",
    "        'frame_index': frame_index,\n",
    "        'timestamp': timestamp,\n",
    "        'next.reward': next_reward,\n",
    "        'index': index,\n",
    "        'observation.image_is_pad': observation_image_is_pad,\n",
    "        'observation.state_is_pad': observation_state_is_pad,\n",
    "        'action_is_pad': action_is_pad\n",
    "    }\n",
    "    output_dict = policy(batch)\n",
    "    loss = output_dict[\"loss\"]\n",
    "    opt_vq_vae.zero_grad()\n",
    "    loss.backward()    \n",
    "    opt_vq_vae.step()\n",
    "    return loss.realize()\n",
    "\n",
    "@TinyJit\n",
    "@Tensor.train()\n",
    "def train_step_main(\n",
    "    observation_image: Tensor,\n",
    "    observation_state: Tensor,\n",
    "    action: Tensor,\n",
    "    episode_index: Tensor,\n",
    "    frame_index: Tensor,\n",
    "    timestamp: Tensor,\n",
    "    next_reward: Tensor,\n",
    "    next_done: Tensor,\n",
    "    next_success: Tensor,\n",
    "    index: Tensor,\n",
    "    observation_image_is_pad: Tensor,\n",
    "    observation_state_is_pad: Tensor,\n",
    "    action_is_pad: Tensor\n",
    "):\n",
    "    batch = {\n",
    "        'observation.image': observation_image,\n",
    "        'observation.state': observation_state,\n",
    "        'action': action,\n",
    "        'episode_index': episode_index,\n",
    "        'frame_index': frame_index,\n",
    "        'timestamp': timestamp,\n",
    "        'next.reward': next_reward,\n",
    "        'index': index,\n",
    "        'observation.image_is_pad': observation_image_is_pad,\n",
    "        'observation.state_is_pad': observation_state_is_pad,\n",
    "        'action_is_pad': action_is_pad\n",
    "    }\n",
    "    output_dict = policy(batch)\n",
    "    loss = output_dict[\"loss\"]\n",
    "    opts.zero_grad()\n",
    "    opt_decay.zero_grad()\n",
    "    opt_no_decay.zero_grad()\n",
    "    loss.backward()\n",
    "    # grad_norm = clip_grad_norm_(vqvae_params, 10.0) if not policy.vqbet.action_head.vqvae_model.discretized else None\n",
    "    opt_decay.step()\n",
    "    opt_no_decay.step()\n",
    "    lr_scheduler.step()\n",
    "    return loss\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # Run training loop.\n",
    "    print(f'Starting training loop')\n",
    "    # Create dataloader for offline training.\n",
    "    dataloader = DataLoader(\n",
    "        dataset,\n",
    "        num_workers=0,\n",
    "        batch_size=64,\n",
    "        shuffle=True,\n",
    "        pin_memory=True,\n",
    "        drop_last=True,\n",
    "    )\n",
    "\n",
    "    step = 0\n",
    "    if step > 0:\n",
    "        lr_scheduler.num_warmup_steps = max(0, step - 500)\n",
    "        lr_scheduler.epoch_counter = step\n",
    "\n",
    "    done = False\n",
    "    with Tensor.train():\n",
    "        while not done:\n",
    "            for batch in dataloader:\n",
    "                batch = {k: Tensor(v.numpy(), requires_grad=False) for k, v in batch.items()}\n",
    "                if policy.vqbet.action_head.vqvae_model.discretized:\n",
    "                    start = timeit.default_timer()\n",
    "                    print(\"The start time is :\", start)\n",
    "                    loss = train_step_main(\n",
    "                        batch['observation.image'].realize(),\n",
    "                        batch['observation.state'].realize(),\n",
    "                        batch['action'].realize(),\n",
    "                        batch['episode_index'].realize(),\n",
    "                        batch['frame_index'].realize(),\n",
    "                        batch['timestamp'].realize(),\n",
    "                        batch['next.reward'].realize(),\n",
    "                        batch['next.done'].realize(),\n",
    "                        batch['next.success'].realize(),\n",
    "                        batch['index'].realize(),\n",
    "                        batch['observation.image_is_pad'].realize(),\n",
    "                        batch['observation.state_is_pad'].realize(),\n",
    "                        batch['action_is_pad'].realize()\n",
    "                    )\n",
    "                    print(\"The difference of time is :\", timeit.default_timer() - start)\n",
    "                else:\n",
    "                    loss = train_step_vqvae(\n",
    "                        batch['observation.image'].realize(),\n",
    "                        batch['observation.state'].realize(),\n",
    "                        batch['action'].realize(),\n",
    "                        batch['episode_index'].realize(),\n",
    "                        batch['frame_index'].realize(),\n",
    "                        batch['timestamp'].realize(),\n",
    "                        batch['next.reward'].realize(),\n",
    "                        batch['next.done'].realize(),\n",
    "                        batch['next.success'].realize(),\n",
    "                        batch['index'].realize(),\n",
    "                        batch['observation.image_is_pad'].realize(),\n",
    "                        batch['observation.state_is_pad'].realize(),\n",
    "                        batch['action_is_pad'].realize()\n",
    "                    )\n",
    "\n",
    "                #loss = loss_gradnorm_tuple[0]\n",
    "                #gradnorm = loss_gradnorm_tuple[1]\n",
    "            \n",
    "                if step % log_freq == 0:\n",
    "                    print(f\"step: {step} loss: {loss.numpy():.3f}\")\n",
    "                step += 1\n",
    "\n",
    "                if step % 5000 == 0 or step == 19999:\n",
    "                    try:\n",
    "                        if step < cfg.n_vqvae_training_steps:\n",
    "                            state_dict_vae_encoder = get_state_dict(policy.vqbet.action_head.vqvae_model.encoder)\n",
    "                            safe_save(state_dict_vae_encoder, f'{output_directory}/model_vae_encoder_{step}.safetensors')\n",
    "                            state_dict_vae_decoder = get_state_dict(policy.vqbet.action_head.vqvae_model.decoder)\n",
    "                            safe_save(state_dict_vae_decoder, f'{output_directory}/model_vae_decoder_{step}.safetensors')\n",
    "                            state_dict_vae_vq_layer = get_state_dict(policy.vqbet.action_head.vqvae_model.vq_layer)\n",
    "                            safe_save(state_dict_vae_vq_layer, f'{output_directory}/model_vae_vq_layer_{step}.safetensors')\n",
    "                        else:\n",
    "                            state_dict = get_state_dict(policy)\n",
    "                            safe_save(state_dict, f'{output_directory}/model_{step}.safetensors')\n",
    "                    except:\n",
    "                        print(f'Exception with safe save occured')\n",
    "                if step >= training_steps:\n",
    "                    done = True\n",
    "                    break\n",
    "\n",
    "    # Save a policy checkpoint.\n",
    "    state_dict = get_state_dict(policy)\n",
    "    safe_save(state_dict, f'{output_directory}/model_final.safetensors')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0e51120-2052-4149-a2f7-78e24599cd1e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ee8466e-4f44-49b9-a396-5ea1f3f78e96",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}