flax tpu tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-07-29 20:50:15.522351: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-07-29 20:50:15.537523: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-07-29 20:50:15.542099: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-07-29 20:50:16.461237: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import flax\n",
    "from flax import linen as nn\n",
    "import orbax.checkpoint\n",
    "from flax.metrics import tensorboard\n",
    "from flax.training import train_state\n",
    "import tensorflow_datasets as tfds\n",
    "import optax\n",
    "import orbax\n",
    "from flax.training import orbax_utils\n",
    "from typing import Any, Tuple, Mapping,Callable,List,Dict\n",
    "import os\n",
    "import tqdm\n",
    "import time\n",
    "import tensorflow as tf # For dataset\n",
    "from clu import metrics\n",
    "from flax import struct                # Flax dataclasses\n",
    "import augmax\n",
    "import grain.python as pygrain\n",
    "import numpy as np\n",
    "from jax.experimental import mesh_utils\n",
    "\n",
    "from jax.sharding import PositionalSharding, NamedSharding\n",
    "from jax.sharding import PartitionSpec\n",
    "from functools import partial\n",
    "import functools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "jax.distributed.initialize() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.experimental.set_visible_devices([], 'GPU')\n",
    "tf.config.experimental.set_visible_devices([], 'TPU')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
       " TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
       " TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
       " TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "P = jax.sharding.PartitionSpec\n",
    "device_mesh = mesh_utils.create_device_mesh((4,))\n",
    "pos_sharding = PositionalSharding(device_mesh)\n",
    "mesh = jax.sharding.Mesh(mesh_utils.create_device_mesh((4,1)), ('data', 'model'))\n",
    "named_sharding = NamedSharding(mesh, P('data', 'model'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(data_name=\"mnist\", batch_size=64, image_scale=160, splits=[\"train\", \"test\"]):\n",
    "    def augmenter():\n",
    "        # @tf.function()\n",
    "        def augment(sample):\n",
    "            image = tf.cast(sample['image'], tf.float32) / 255.\n",
    "            \n",
    "            image = (\n",
    "                tf.cast(sample[\"image\"], tf.float32) - 127.5\n",
    "            ) / 127.5\n",
    "            image = tf.image.resize(\n",
    "                image, [image_scale, image_scale], method=\"area\", antialias=True\n",
    "            )\n",
    "            image = tf.image.random_flip_left_right(image)\n",
    "            image = tf.image.random_contrast(image, 0.997, 1.05)\n",
    "            image = tf.image.random_brightness(image, 0.2)\n",
    "\n",
    "            image = tf.clip_by_value(image, -1.0, 1.0)\n",
    "            label = sample['label']\n",
    "            return {\"image\": image, \"label\": label}\n",
    "        return augment\n",
    "\n",
    "    # Load CelebA Dataset\n",
    "    (train_ds, test_ds) = tfds.load(data_name, split=splits, shuffle_files=True)\n",
    "    train_len = len(train_ds)\n",
    "    test_len = len(test_ds)\n",
    "    \n",
    "    train_ds: tf.data.Dataset = train_ds\n",
    "    test_ds: tf.data.Dataset = test_ds\n",
    "    \n",
    "    train_ds = (\n",
    "        train_ds\n",
    "        .map(\n",
    "            augmenter(),\n",
    "            num_parallel_calls=tf.data.AUTOTUNE,\n",
    "        )\n",
    "        .cache()  # Cache after augmenting to avoid recomputation\n",
    "        .repeat()  # Repeats the dataset indefinitely\n",
    "        .shuffle(4096)  # Ensure this is adequate for your dataset size\n",
    "        .batch(batch_size, drop_remainder=True)\n",
    "        .prefetch(tf.data.experimental.AUTOTUNE)\n",
    "    )\n",
    "    \n",
    "    test_ds = (\n",
    "        test_ds\n",
    "        .map(\n",
    "            augmenter(),\n",
    "            num_parallel_calls=tf.data.AUTOTUNE,\n",
    "        )\n",
    "        .cache()  # Cache after augmenting to avoid recomputation\n",
    "        # .repeat()  # Repeats the dataset indefinitely\n",
    "        .batch(batch_size, drop_remainder=True)\n",
    "        .prefetch(tf.data.experimental.AUTOTUNE)\n",
    "    )\n",
    "    \n",
    "    def get_trainset():\n",
    "        return train_ds.as_numpy_iterator()\n",
    "    \n",
    "    def get_testset():\n",
    "        return test_ds.as_numpy_iterator()\n",
    "    return {\n",
    "        \"train\": get_trainset,\n",
    "        \"test\": get_testset,\n",
    "        \"train_len\": train_len,\n",
    "        \"test_len\": test_len,\n",
    "        \"batch_size\": batch_size\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(data_name=\"mnist\", batch_size=64, image_scale=256, method=jax.image.ResizeMethod.LANCZOS3):\n",
    "    (train_ds, test_ds) = tfds.data_source(\"mnist\", split=[\"train\", \"test\"], try_gcs=False)\n",
    "    train_len = len(train_ds)\n",
    "    test_len = len(test_ds)\n",
    "\n",
    "    cpu_device = jax.devices(\"cpu\")[0]\n",
    "    try:\n",
    "        gpu_device = jax.devices(\"gpu\")[0]\n",
    "    except:\n",
    "        gpu_device = None\n",
    "        try:\n",
    "            tpu_devices = jax.devices(\"tpu\")\n",
    "        except:\n",
    "            tpu_devices = None\n",
    "    print(f\"Gpu Device: {gpu_device}, Cpu Device: {cpu_device}, TPU Devices: {tpu_devices}\")\n",
    "        \n",
    "    def preprocess(image):\n",
    "        image = (image - 127.5) / 127.5\n",
    "        # image = jax.image.resize(image, (image_scale, image_scale, 3), method=method)\n",
    "        image = jnp.clip(image, -1.0, 1.0)\n",
    "        # image = jax.device_put(image, device=jax.devices(\"gpu\")[0]) \n",
    "        return  image\n",
    "    \n",
    "    preprocess = jax.jit(preprocess, backend=\"cpu\")\n",
    "\n",
    "    class augmenter(pygrain.RandomMapTransform):\n",
    "        def random_map(self, element: Dict[str, Any], rng: np.random.Generator) ->  Dict[str, jnp.array]:\n",
    "            image = element['image']\n",
    "            image = preprocess(image)\n",
    "            # image = augments(rng.integers(0, 2**32, [2], dtype=np.uint32), image) \n",
    "            label = element['label']\n",
    "            return {'image':image, 'label':label}\n",
    "\n",
    "    transformations = [augmenter(), pygrain.Batch(batch_size, drop_remainder=True)]\n",
    "\n",
    "    train_sampler = pygrain.IndexSampler(\n",
    "        num_records=train_len,\n",
    "        shuffle=True,\n",
    "        seed=0,\n",
    "        num_epochs=None,\n",
    "        shard_options=pygrain.ShardByJaxProcess(),\n",
    "    )\n",
    "    \n",
    "    train_ds = pygrain.DataLoader(\n",
    "        data_source=train_ds,\n",
    "        sampler=train_sampler,\n",
    "        operations=transformations,\n",
    "        # worker_count=0,\n",
    "        # read_options=pygrain.ReadOptions(8, 500),\n",
    "        # worker_buffer_size=5\n",
    "        )\n",
    "    \n",
    "    test_sampler = pygrain.IndexSampler(\n",
    "        num_records=test_len,\n",
    "        shuffle=True,\n",
    "        seed=0,\n",
    "        num_epochs=None,\n",
    "        shard_options=pygrain.ShardByJaxProcess(),\n",
    "    )\n",
    "    \n",
    "    test_ds = pygrain.DataLoader(\n",
    "        data_source=test_ds,\n",
    "        sampler=test_sampler,\n",
    "        operations=transformations,\n",
    "        # worker_count=0,\n",
    "        # read_options=pygrain.ReadOptions(8, 500),\n",
    "        # worker_buffer_size=5\n",
    "        )\n",
    "    return {\n",
    "        \"train\": train_ds,\n",
    "        \"test\": test_ds,\n",
    "        \"train_len\": train_len,\n",
    "        \"test_len\": test_len,\n",
    "        \"batch_size\": batch_size\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = get_dataset(\"imagenette\", batch_size=128, splits=['train', 'validation'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-07-29 20:20:17.968071: W tensorflow/core/kernels/data/cache_dataset_ops.cc:913] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset  will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.\n"
     ]
    }
   ],
   "source": [
    "d = next(iter(data['train']()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.pmap\n",
    "def testfun(x):\n",
    "    x = x.reshape((-1))\n",
    "    print(x.shape)\n",
    "    return jnp.sin(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     TPU 0                                      </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                     TPU 2                                      </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                     TPU 1                                      </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                                                </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                                                </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                     TPU 3                                      </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                                                </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                      \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                                      \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;173;73;74m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;173;73;74m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                                      \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;173;73;74m                                                                                \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;181;207;107m                                                                                \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;181;207;107m                                     \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                                      \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;181;207;107m                                                                                \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "inp = d['image'].reshape((-1, 160*160))\n",
    "inp = jax.device_put(inp, named_sharding)\n",
    "jax.debug.visualize_array_sharding(inp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2457600,)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     TPU 0                                      </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                     TPU 1                                      </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                     TPU 2                                      </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                                                </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                                                </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                     TPU 3                                      </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                                                </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                      \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                                      \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;173;73;74m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;173;73;74m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                                      \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;173;73;74m                                                                                \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;181;207;107m                                                                                \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;181;207;107m                                     \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                                      \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;181;207;107m                                                                                \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)} None [[-0.24927016 -0.22547898 -0.23789191 ... -0.74764276 -0.7450645\n",
      "  -0.7432891 ]\n",
      " [ 0.1453397   0.40575063  0.44915524 ... -0.3922717  -0.4679144\n",
      "  -0.4169427 ]\n",
      " [ 0.30879593  0.33196172  0.3142953  ... -0.5827335  -0.5731142\n",
      "  -0.6314607 ]\n",
      " [ 0.39536574  0.48585674  0.5317951  ... -0.5890244  -0.567123\n",
      "  -0.68381965]]\n"
     ]
    }
   ],
   "source": [
    "k = testfun(inp.reshape((4, -1, *inp.shape[1:])))\n",
    "print(k.devices(), jax.debug.visualize_array_sharding(k), k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import flax.jax_utils\n",
    "\n",
    "\n",
    "@struct.dataclass\n",
    "class Metrics(metrics.Collection):\n",
    "  accuracy: metrics.Accuracy\n",
    "  loss: metrics.Average.from_output('loss')\n",
    "\n",
    "# Define the TrainState \n",
    "class SimpleTrainState(train_state.TrainState):\n",
    "    rngs: jax.random.PRNGKey\n",
    "    metrics: Metrics\n",
    "\n",
    "    def get_random_key(self):\n",
    "        rngs, subkey = jax.random.split(self.rngs)\n",
    "        return self.replace(rngs=rngs), subkey\n",
    "\n",
    "class SimpleTrainer:\n",
    "    state : SimpleTrainState\n",
    "    best_state : SimpleTrainState\n",
    "    best_loss : float\n",
    "    model : nn.Module\n",
    "    ema_decay:float = 0.999\n",
    "    \n",
    "    def __init__(self, \n",
    "                 model:nn.Module, \n",
    "                input_shapes:Dict[str, Tuple[int]],\n",
    "                 optimizer: optax.GradientTransformation,\n",
    "                 rngs:jax.random.PRNGKey,\n",
    "                 train_state:SimpleTrainState=None,\n",
    "                 name:str=\"Simple\",\n",
    "                 load_from_checkpoint:bool=False,\n",
    "                 checkpoint_suffix:str=\"\",\n",
    "                 loss_fn=optax.l2_loss,\n",
    "                 param_transforms:Callable=None,\n",
    "                 ):\n",
    "        self.model = model\n",
    "        self.name = name\n",
    "        self.loss_fn = loss_fn\n",
    "        self.input_shapes = input_shapes\n",
    "\n",
    "        checkpointer = orbax.checkpoint.PyTreeCheckpointer()\n",
    "        options = orbax.checkpoint.CheckpointManagerOptions(max_to_keep=4, create=True)\n",
    "        self.checkpointer = orbax.checkpoint.CheckpointManager(self.checkpoint_path() + checkpoint_suffix, checkpointer, options)\n",
    "\n",
    "        if load_from_checkpoint:\n",
    "            latest_step, old_state, old_best_state = self.load()\n",
    "        else:\n",
    "            latest_step, old_state, old_best_state = 0, None, None\n",
    "            \n",
    "        self.latest_step = latest_step\n",
    "\n",
    "        if train_state == None:\n",
    "            self.init_state(optimizer, rngs, existing_state=old_state, existing_best_state=old_best_state, model=model, param_transforms=param_transforms)\n",
    "        else:\n",
    "            self.state = train_state\n",
    "            self.best_state = train_state\n",
    "            self.best_loss = 1e9\n",
    "    \n",
    "    def get_input_ones(self):\n",
    "        return {k:jnp.ones((1, *v)) for k,v in self.input_shapes.items()}\n",
    "\n",
    "    def init_state(self,\n",
    "                   optimizer: optax.GradientTransformation, \n",
    "                   rngs:jax.random.PRNGKey,\n",
    "                   existing_state:dict=None,\n",
    "                   existing_best_state:dict=None,\n",
    "                   model:nn.Module=None,\n",
    "                   param_transforms:Callable=None\n",
    "                   ):\n",
    "        @partial(jax.pmap, axis_name=\"device\")\n",
    "        def init_fn(rngs):\n",
    "            rngs, subkey = jax.random.split(rngs)\n",
    "\n",
    "            if existing_state == None:\n",
    "                input_vars = self.get_input_ones()\n",
    "                params = model.init(subkey, **input_vars)\n",
    "\n",
    "            # if param_transforms is not None:\n",
    "            #     params = param_transforms(params)\n",
    "                \n",
    "            state = SimpleTrainState.create(\n",
    "                apply_fn=model.apply,\n",
    "                params=params,\n",
    "                tx=optimizer,\n",
    "                rngs=rngs,\n",
    "                metrics=Metrics.empty()\n",
    "            )\n",
    "            return state\n",
    "        self.state = init_fn(jax.device_put_replicated(rngs, jax.devices()))\n",
    "        self.best_loss = 1e9\n",
    "        if existing_best_state is not None:\n",
    "            self.best_state = self.state.replace(params=existing_best_state['params'], ema_params=existing_best_state['ema_params'])\n",
    "        else:\n",
    "            self.best_state = self.state\n",
    "\n",
    "    def checkpoint_path(self):\n",
    "        experiment_name = self.name\n",
    "        path = os.path.join(os.path.abspath('./checkpoints'), experiment_name)\n",
    "        if not os.path.exists(path):\n",
    "            os.makedirs(path)\n",
    "        return path\n",
    "    \n",
    "    def tensorboard_path(self):\n",
    "        experiment_name = self.name\n",
    "        path = os.path.join(os.path.abspath('./tensorboard'), experiment_name)\n",
    "        if not os.path.exists(path):\n",
    "            os.makedirs(path)\n",
    "        return path\n",
    "\n",
    "    def load(self):\n",
    "        step = self.checkpointer.latest_step()\n",
    "        print(\"Loading model from checkpoint\", step)\n",
    "        ckpt = self.checkpointer.restore(step)\n",
    "        state = ckpt['state']\n",
    "        best_state = ckpt['best_state']\n",
    "        # Convert the state to a TrainState\n",
    "        self.best_loss = ckpt['best_loss']\n",
    "        print(f\"Loaded model from checkpoint at step {step}\", ckpt['best_loss'])\n",
    "        return step, state, best_state\n",
    "\n",
    "    def save(self, epoch=0):\n",
    "        print(f\"Saving model at epoch {epoch}\")\n",
    "        ckpt = {\n",
    "            # 'model': self.model,\n",
    "            'state': self.state,\n",
    "            'best_state': self.best_state,\n",
    "            'best_loss': self.best_loss\n",
    "        }\n",
    "        try:\n",
    "            # save_args = orbax_utils.save_args_from_target(ckpt)\n",
    "            # self.checkpointer.save(epoch, ckpt, save_kwargs={'save_args': save_args}, force=True)\n",
    "            pass\n",
    "        except Exception as e:\n",
    "            print(\"Error saving checkpoint\", e)\n",
    "\n",
    "    def _define_train_step(self):\n",
    "        model = self.model\n",
    "        loss_fn = self.loss_fn\n",
    "        \n",
    "        @partial(jax.pmap, axis_name=\"device\")\n",
    "        def train_step(state:SimpleTrainState, batch):\n",
    "            \"\"\"Train for a single step.\"\"\"\n",
    "            images = batch['image']\n",
    "            labels= batch['label']\n",
    "            \n",
    "            def model_loss(params):\n",
    "                preds = model.apply(params, images)\n",
    "                expected_output = labels\n",
    "                nloss = loss_fn(preds, expected_output)\n",
    "                loss = jnp.mean(nloss)\n",
    "                return loss\n",
    "            loss, grads = jax.value_and_grad(model_loss)(state.params)\n",
    "            grads = jax.lax.pmean(grads, \"device\")\n",
    "            state = state.apply_gradients(grads=grads) \n",
    "            return state, loss\n",
    "        return train_step\n",
    "    \n",
    "    def _define_compute_metrics(self):\n",
    "        model = self.model\n",
    "        loss_fn = self.loss_fn\n",
    "        \n",
    "        @jax.jit\n",
    "        def compute_metrics(state:SimpleTrainState, batch):\n",
    "            preds = model.apply(state.params, batch['image'])\n",
    "            expected_output = batch['label']\n",
    "            loss = jnp.mean(loss_fn(preds, expected_output))\n",
    "            metric_updates = state.metrics.single_from_model_output(loss=loss, logits=preds, labels=expected_output)\n",
    "            metrics = state.metrics.merge(metric_updates)\n",
    "            state = state.replace(metrics=metrics)\n",
    "            return state\n",
    "        return compute_metrics\n",
    "\n",
    "    def summary(self):\n",
    "        input_vars = self.get_input_ones()\n",
    "        print(self.model.tabulate(jax.random.key(0), **input_vars, console_kwargs={\"width\": 200, \"force_jupyter\":True, }))\n",
    "    \n",
    "    def config(self):\n",
    "        return {\n",
    "            \"model\": self.model,\n",
    "            \"state\": self.state,\n",
    "            \"name\": self.name,\n",
    "            \"input_shapes\": self.input_shapes\n",
    "        }\n",
    "        \n",
    "    def init_tensorboard(self, batch_size, steps_per_epoch, epochs):\n",
    "        summary_writer = tensorboard.SummaryWriter(self.tensorboard_path())\n",
    "        summary_writer.hparams({\n",
    "            **self.config(),\n",
    "            \"steps_per_epoch\": steps_per_epoch,\n",
    "            \"epochs\": epochs,\n",
    "            \"batch_size\": batch_size\n",
    "        })\n",
    "        return summary_writer\n",
    "        \n",
    "    def fit(self, data, steps_per_epoch, epochs):\n",
    "        train_ds = iter(data['train']())\n",
    "        if 'test' in data:\n",
    "            test_ds = data['test']\n",
    "        else:\n",
    "            test_ds = None\n",
    "        train_step = self._define_train_step()\n",
    "        compute_metrics = self._define_compute_metrics()\n",
    "        state = self.state\n",
    "        device_count = jax.device_count()\n",
    "        # train_ds = flax.jax_utils.prefetch_to_device(train_ds, jax.devices())\n",
    "        \n",
    "        summary_writer = self.init_tensorboard(data['batch_size'], steps_per_epoch, epochs)\n",
    "        \n",
    "        for epoch in range(epochs):\n",
    "            current_epoch = self.latest_step + epoch + 1\n",
    "            print(f\"\\nEpoch {current_epoch}/{epochs}\")\n",
    "            start_time = time.time()\n",
    "            epoch_loss = 0\n",
    "            \n",
    "            with tqdm.tqdm(total=steps_per_epoch, desc=f'\\t\\tEpoch {current_epoch}', ncols=100, unit='step') as pbar:\n",
    "                for i in range(steps_per_epoch):\n",
    "                    batch = next(train_ds)\n",
    "                    batch = jax.tree.map(lambda x: x.reshape((device_count, -1, *x.shape[1:])), batch)\n",
    "                    # print(batch['image'].shape)\n",
    "                    state, loss = train_step(state, batch)\n",
    "                    loss = jnp.mean(loss)\n",
    "                    # print(\"==>\", loss)\n",
    "                    epoch_loss += loss\n",
    "                    if i % 100 == 0:\n",
    "                        pbar.set_postfix(loss=f'{loss:.4f}')\n",
    "                        pbar.update(100)\n",
    "                        current_step = current_epoch*steps_per_epoch + i\n",
    "                        summary_writer.scalar('Train Loss', loss, step=current_step)\n",
    "                        \n",
    "            end_time = time.time()\n",
    "            self.state = state\n",
    "            total_time = end_time - start_time\n",
    "            avg_time_per_step = total_time / steps_per_epoch\n",
    "            avg_loss = epoch_loss / steps_per_epoch\n",
    "            if avg_loss < self.best_loss:\n",
    "                self.best_loss = avg_loss\n",
    "                self.best_state = state\n",
    "                self.save(current_epoch)\n",
    "            \n",
    "            # Compute Metrics\n",
    "            metrics_str = ''\n",
    "            # if test_ds is not None:\n",
    "            #     for test_batch in iter(test_ds()):\n",
    "            #         state = compute_metrics(state, test_batch)\n",
    "            #     metrics = state.metrics.compute()\n",
    "            #     for metric,value in metrics.items():\n",
    "            #         summary_writer.scalar(f'Test {metric}', value, step=current_epoch)\n",
    "            #         metrics_str += f', Test {metric}: {value:.4f}'\n",
    "            #     state = state.replace(metrics=Metrics.empty())\n",
    "                    \n",
    "            print(f\"\\n\\tEpoch {current_epoch} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss} {metrics_str}\")\n",
    "            \n",
    "        self.save(epochs)\n",
    "        return self.state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "  \"\"\"A simple CNN model.\"\"\"\n",
    "\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n",
    "    x = nn.relu(x)\n",
    "    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n",
    "    x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n",
    "    x = nn.relu(x)\n",
    "    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n",
    "    x = x.reshape((x.shape[0], -1))  # flatten\n",
    "    x = nn.Dense(features=256)(x)\n",
    "    x = nn.relu(x)\n",
    "    x = nn.Dense(features=10)(x)\n",
    "    return x\n",
    "  \n",
    "  # Kernel initializer to use\n",
    "def kernel_init(scale):\n",
    "    scale = max(scale, 1e-10)\n",
    "    return nn.initializers.variance_scaling(scale=scale, mode=\"fan_avg\", distribution=\"truncated_normal\")\n",
    "\n",
    "class NormalAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Simple implementation of the normal attention.\n",
    "    \"\"\"\n",
    "    query_dim: int\n",
    "    heads: int = 4\n",
    "    dim_head: int = 64\n",
    "    dtype: Any = jnp.float32\n",
    "    precision: Any = jax.lax.Precision.HIGHEST\n",
    "    use_bias: bool = True\n",
    "    kernel_init: Callable = lambda : kernel_init(1.0)\n",
    "\n",
    "    def setup(self):\n",
    "        inner_dim = self.dim_head * self.heads\n",
    "        dense = functools.partial(\n",
    "            nn.DenseGeneral,\n",
    "            features=[self.heads, self.dim_head], \n",
    "            axis=-1, \n",
    "            precision=self.precision, \n",
    "            use_bias=self.use_bias, \n",
    "            kernel_init=self.kernel_init(), \n",
    "            dtype=self.dtype\n",
    "        )\n",
    "        self.query = dense(name=\"to_q\")\n",
    "        self.key = dense(name=\"to_k\")\n",
    "        self.value = dense(name=\"to_v\")\n",
    "\n",
    "        self.proj_attn = nn.DenseGeneral(\n",
    "            self.query_dim, \n",
    "            axis=(-2, -1), \n",
    "            precision=self.precision, \n",
    "            use_bias=self.use_bias, \n",
    "            dtype=self.dtype, \n",
    "            name=\"to_out_0\",\n",
    "            kernel_init=self.kernel_init()\n",
    "            # kernel_init=jax.nn.initializers.xavier_uniform()\n",
    "        )\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x, context=None):\n",
    "        # x has shape [B, H, W, C]\n",
    "        context = x if context is None else context\n",
    "        query = self.query(x)\n",
    "        key = self.key(context)\n",
    "        value = self.value(context)\n",
    "        \n",
    "        hidden_states = nn.dot_product_attention(\n",
    "            query, key, value, dtype=self.dtype, broadcast_dropout=False, dropout_rng=None, precision=self.precision\n",
    "        )\n",
    "        proj = self.proj_attn(hidden_states)\n",
    "        return proj\n",
    "    \n",
    "class AttentionBlock(nn.Module):\n",
    "    heads: int = 4\n",
    "    dim_head: int = 32\n",
    "    use_linear_attention: bool = True\n",
    "    dtype: Any = jnp.float32\n",
    "    precision: Any = jax.lax.Precision.HIGH\n",
    "    use_projection: bool = False\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        inner_dim = self.heads * self.dim_head\n",
    "        B, H, W, C = x.shape\n",
    "        normed_x = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)(x)\n",
    "        projected_x = NormalAttention(\n",
    "            query_dim=C,\n",
    "            heads=self.heads,\n",
    "            dim_head=self.dim_head,\n",
    "            name=f'Attention',\n",
    "            precision=self.precision,\n",
    "            use_bias=False,\n",
    "        )(normed_x)\n",
    "        out = x + projected_x\n",
    "        return out\n",
    "    \n",
    "class TimeEmbedding(nn.Module):\n",
    "    features:int\n",
    "    nax_positions:int=10000\n",
    "\n",
    "    def setup(self):\n",
    "        half_dim = self.features // 2\n",
    "        emb = jnp.log(self.nax_positions) / (half_dim - 1)\n",
    "        emb = jnp.exp(-emb * jnp.arange(half_dim, dtype=jnp.float32))\n",
    "        self.embeddings = emb\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = jax.lax.convert_element_type(x, jnp.float32)\n",
    "        emb = x[:, None] * self.embeddings[None, :]\n",
    "        emb = jnp.concatenate([jnp.sin(emb), jnp.cos(emb)], axis=-1)\n",
    "        return emb\n",
    "    \n",
    "class TimeProjection(nn.Module):\n",
    "    features:int\n",
    "    activation:Callable=jax.nn.gelu\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        x = nn.DenseGeneral(self.features, kernel_init=kernel_init(1.0))(x)\n",
    "        x = self.activation(x)\n",
    "        x = nn.DenseGeneral(self.features, kernel_init=kernel_init(1.0))(x)\n",
    "        x = self.activation(x)\n",
    "        return x\n",
    "\n",
    "class SeparableConv(nn.Module):\n",
    "    features:int\n",
    "    kernel_size:tuple=(3, 3)\n",
    "    strides:tuple=(1, 1)\n",
    "    use_bias:bool=False\n",
    "    kernel_init:Callable=kernel_init(1.0)\n",
    "    padding:str=\"SAME\"\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        in_features = x.shape[-1]\n",
    "        depthwise = nn.Conv(\n",
    "            features=in_features, kernel_size=self.kernel_size,\n",
    "            strides=self.strides, kernel_init=self.kernel_init,\n",
    "            feature_group_count=in_features, use_bias=self.use_bias,\n",
    "            padding=self.padding\n",
    "        )(x)\n",
    "        pointwise = nn.Conv(\n",
    "            features=self.features, kernel_size=(1, 1),\n",
    "            strides=(1, 1), kernel_init=self.kernel_init,\n",
    "            use_bias=self.use_bias\n",
    "        )(depthwise)\n",
    "        return pointwise\n",
    "\n",
    "class ConvLayer(nn.Module):\n",
    "    conv_type:str\n",
    "    features:int\n",
    "    kernel_size:tuple=(3, 3)\n",
    "    strides:tuple=(1, 1)\n",
    "    kernel_init:Callable=kernel_init(1.0)\n",
    "\n",
    "    def setup(self):\n",
    "        if self.conv_type == \"conv\":\n",
    "            self.conv = nn.Conv(\n",
    "                features=self.features,\n",
    "                kernel_size=self.kernel_size,\n",
    "                strides=self.strides,\n",
    "                kernel_init=self.kernel_init,\n",
    "            )\n",
    "        elif self.conv_type == \"separable\":\n",
    "            self.conv = SeparableConv(\n",
    "                features=self.features,\n",
    "                kernel_size=self.kernel_size,\n",
    "                strides=self.strides,\n",
    "                kernel_init=self.kernel_init,\n",
    "            )\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class Upsample(nn.Module):\n",
    "    features:int\n",
    "    scale:int\n",
    "    activation:Callable=jax.nn.swish\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x, residual=None):\n",
    "        out = x\n",
    "        # out = PixelShuffle(scale=self.scale)(out)\n",
    "        B, H, W, C = x.shape\n",
    "        out = jax.image.resize(x, (B, H * self.scale, W * self.scale, C), method=\"nearest\")\n",
    "        out = ConvLayer(\n",
    "            \"conv\",\n",
    "            features=self.features,\n",
    "            kernel_size=(3, 3),\n",
    "            strides=(1, 1),\n",
    "        )(out)\n",
    "        if residual is not None:\n",
    "            out = jnp.concatenate([out, residual], axis=-1)\n",
    "        return out\n",
    "\n",
    "class Downsample(nn.Module):\n",
    "    features:int\n",
    "    scale:int\n",
    "    activation:Callable=jax.nn.swish\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x, residual=None):\n",
    "        out = ConvLayer(\n",
    "            \"conv\",\n",
    "            features=self.features,\n",
    "            kernel_size=(3, 3),\n",
    "            strides=(2, 2)\n",
    "        )(x)\n",
    "        if residual is not None:\n",
    "            if residual.shape[1] > out.shape[1]:\n",
    "                residual = nn.avg_pool(residual, window_shape=(2, 2), strides=(2, 2), padding=\"SAME\")\n",
    "            out = jnp.concatenate([out, residual], axis=-1)\n",
    "        return out\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    conv_type:str\n",
    "    features:int\n",
    "    kernel_size:tuple=(3, 3)\n",
    "    strides:tuple=(1, 1)\n",
    "    padding:str=\"SAME\"\n",
    "    activation:Callable=jax.nn.swish\n",
    "    direction:str=None\n",
    "    res:int=2\n",
    "    norm_groups:int=8\n",
    "    kernel_init:Callable=kernel_init(1.0)\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x:jax.Array, extra_features:jax.Array=None):\n",
    "        residual = x\n",
    "        out = nn.GroupNorm(self.norm_groups)(x)\n",
    "        out = self.activation(out)\n",
    "\n",
    "        out = ConvLayer(\n",
    "            self.conv_type,\n",
    "            features=self.features,\n",
    "            kernel_size=self.kernel_size,\n",
    "            strides=self.strides,\n",
    "            kernel_init=self.kernel_init,\n",
    "            name=\"conv1\"\n",
    "        )(out)\n",
    "\n",
    "        out = nn.GroupNorm(self.norm_groups)(out)\n",
    "        out = self.activation(out)\n",
    "\n",
    "        out = ConvLayer(\n",
    "            self.conv_type,\n",
    "            features=self.features,\n",
    "            kernel_size=self.kernel_size,\n",
    "            strides=self.strides,\n",
    "            kernel_init=self.kernel_init,\n",
    "            name=\"conv2\"\n",
    "        )(out)\n",
    "\n",
    "        if residual.shape != out.shape:\n",
    "            residual = ConvLayer(\n",
    "                self.conv_type,\n",
    "                features=self.features,\n",
    "                kernel_size=(1, 1),\n",
    "                strides=1,\n",
    "                kernel_init=self.kernel_init,\n",
    "                name=\"residual_conv\"\n",
    "            )(residual)\n",
    "        out = out + residual\n",
    "\n",
    "        out = jnp.concatenate([out, extra_features], axis=-1) if extra_features is not None else out\n",
    "\n",
    "        return out\n",
    "    \n",
    "class Classifier(nn.Module):\n",
    "    feature_depths:list=[64, 128, 256, 512],\n",
    "    attention_configs:list=[{\"heads\":8}, {\"heads\":8}, {\"heads\":8}, {\"heads\":8}],\n",
    "    num_res_blocks:int=2,\n",
    "    num_middle_res_blocks:int=1,\n",
    "    activation:Callable = jax.nn.swish\n",
    "    norm_groups:int=8\n",
    "    major_conv_type:str=\"conv\"\n",
    "    mid_conv_type:str=\"conv\"\n",
    "    \n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        feature_depths = self.feature_depths\n",
    "        attention_configs = self.attention_configs\n",
    "\n",
    "        conv_type = \"conv\"\n",
    "        up_conv_type = down_conv_type = self.major_conv_type\n",
    "        middle_conv_type = self.mid_conv_type\n",
    "\n",
    "        x = ConvLayer(\n",
    "            conv_type,\n",
    "            features=self.feature_depths[0],\n",
    "            kernel_size=(3, 3),\n",
    "            strides=(1, 1),\n",
    "            kernel_init=kernel_init(1.0)\n",
    "        )(x)\n",
    "        downs = [x]\n",
    "\n",
    "        # Downscaling blocks\n",
    "        for i, (dim_out, attention_config) in enumerate(zip(feature_depths, attention_configs)):\n",
    "            dim_in = x.shape[-1]\n",
    "            # dim_in = dim_out\n",
    "            for j in range(self.num_res_blocks):\n",
    "                x = ResidualBlock(\n",
    "                    down_conv_type,\n",
    "                    name=f\"down_{i}_residual_{j}\",\n",
    "                    features=dim_in,\n",
    "                    kernel_init=kernel_init(1.0),\n",
    "                    kernel_size=(3, 3),\n",
    "                    strides=(1, 1),\n",
    "                    activation=self.activation,\n",
    "                    norm_groups=self.norm_groups\n",
    "                )(x)\n",
    "                if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block\n",
    "                    x = AttentionBlock(heads=attention_config['heads'], \n",
    "                                       dim_head=dim_in // attention_config['heads'],\n",
    "                                       name=f\"down_{i}_attention_{j}\")(x)\n",
    "                downs.append(x)\n",
    "            if i != len(feature_depths) - 1:\n",
    "                x = Downsample(\n",
    "                    features=dim_out,\n",
    "                    scale=2,\n",
    "                    activation=self.activation,\n",
    "                    name=f\"down_{i}_downsample\"\n",
    "                )(x)\n",
    "\n",
    "        # Middle Blocks\n",
    "        middle_dim_out = self.feature_depths[-1]\n",
    "        middle_attention = self.attention_configs[-1]\n",
    "        for j in range(self.num_middle_res_blocks):\n",
    "            x = ResidualBlock(\n",
    "                middle_conv_type,\n",
    "                name=f\"middle_res1_{j}\",\n",
    "                features=middle_dim_out,\n",
    "                kernel_init=kernel_init(1.0),\n",
    "                kernel_size=(3, 3),\n",
    "                strides=(1, 1),\n",
    "                activation=self.activation,\n",
    "                norm_groups=self.norm_groups\n",
    "            )(x)\n",
    "            if middle_attention is not None and j == self.num_middle_res_blocks - 1:   # Apply attention only on the last block\n",
    "                x = AttentionBlock(heads=attention_config['heads'], \n",
    "                                   dim_head=middle_dim_out // attention_config['heads'],\n",
    "                                   use_linear_attention=False, name=f\"middle_attention_{j}\")(x)\n",
    "            x = ResidualBlock(\n",
    "                middle_conv_type,\n",
    "                name=f\"middle_res2_{j}\",\n",
    "                features=middle_dim_out,\n",
    "                kernel_init=kernel_init(1.0),\n",
    "                kernel_size=(3, 3),\n",
    "                strides=(1, 1),\n",
    "                activation=self.activation,\n",
    "                norm_groups=self.norm_groups\n",
    "            )(x)\n",
    "            \n",
    "        x = Downsample(\n",
    "            features=middle_dim_out * 2,\n",
    "            scale=2,\n",
    "            activation=self.activation,\n",
    "            name=f\"middle_downsample\"\n",
    "        )(x)\n",
    "        \n",
    "        x = ResidualBlock(\n",
    "            conv_type,\n",
    "            name=\"final_residual\",\n",
    "            features=self.feature_depths[0],\n",
    "            kernel_init=kernel_init(1.0),\n",
    "            kernel_size=(3,3),\n",
    "            strides=(1, 1),\n",
    "            activation=self.activation,\n",
    "            norm_groups=self.norm_groups\n",
    "        )(x)\n",
    "\n",
    "        x = nn.GroupNorm(self.norm_groups)(x)\n",
    "        x = self.activation(x)\n",
    "        \n",
    "        # Average Pooling\n",
    "        x = nn.avg_pool(x, window_shape=(x.shape[1], x.shape[2]), strides=(1, 1))\n",
    "        \n",
    "        x = x.reshape((x.shape[0], -1))\n",
    "        \n",
    "        # Classifier head\n",
    "        x = nn.Dense(10)(x)\n",
    "        \n",
    "        return x\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                                                              Classifier Summary                                                              </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> path                                  </span>┃<span style=\"font-weight: bold\"> module          </span>┃<span style=\"font-weight: bold\"> inputs                 </span>┃<span style=\"font-weight: bold\"> outputs                </span>┃<span style=\"font-weight: bold\"> params                         </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│                                       │ Classifier      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,3]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10]          │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ ConvLayer_0                           │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,3]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ ConvLayer_0/conv                      │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,3]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,3,64]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,792 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(7.2 KB)</span>                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_downsample                     │ Downsample      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_downsample/ConvLayer_0         │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_downsample/ConvLayer_0/conv    │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,160,160,64]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1                    │ AttentionBlock  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/RMSNorm_0          │ RMSNorm         │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">64 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(256 B)</span>                     │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention          │ NormalAttention │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_q     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,8,8]   │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64,8,8]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">4,096 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(16.4 KB)</span>                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_k     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,8,8]   │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64,8,8]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">4,096 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(16.4 KB)</span>                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_v     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,8,8]   │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64,8,8]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">4,096 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(16.4 KB)</span>                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_out_0 │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,8,8]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[8,8,64]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">4,096 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(16.4 KB)</span>                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_downsample                     │ Downsample      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_downsample/ConvLayer_0         │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_downsample/ConvLayer_0/conv    │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,80,80,64]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">73,856 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(295.4 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">256 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(1.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">147,584 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(590.3 KB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">256 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(1.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">147,584 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(590.3 KB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">256 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(1.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">147,584 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(590.3 KB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">256 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(1.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">147,584 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(590.3 KB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1                    │ AttentionBlock  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/RMSNorm_0          │ RMSNorm         │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention          │ NormalAttention │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_q     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,8,16]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128,8,16]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">16,384 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(65.5 KB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_k     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,8,16]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128,8,16]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">16,384 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(65.5 KB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_v     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,8,16]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[128,8,16]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">16,384 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(65.5 KB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_out_0 │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,8,16]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[8,16,128]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">16,384 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(65.5 KB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_downsample                     │ Downsample      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_downsample/ConvLayer_0         │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_downsample/ConvLayer_0/conv    │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,40,40,128]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,128,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">295,168 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(1.2 MB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">512 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">590,080 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.4 MB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">512 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">590,080 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.4 MB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1                     │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/GroupNorm_0         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">512 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv1               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv1/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">590,080 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.4 MB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/GroupNorm_1         │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">512 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv2               │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv2/conv          │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">590,080 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.4 MB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1                    │ AttentionBlock  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/RMSNorm_0          │ RMSNorm         │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">256 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(1.0 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention          │ NormalAttention │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_q     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,8,32]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256,8,32]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">65,536 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(262.1 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_k     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,8,32]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256,8,32]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">65,536 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(262.1 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_v     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,8,32]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[256,8,32]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">65,536 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(262.1 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_out_0 │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,8,32]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[8,32,256]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">65,536 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(262.1 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_downsample                     │ Downsample      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_downsample/ConvLayer_0         │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_downsample/ConvLayer_0/conv    │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,20,20,256]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[512]             │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,256,512]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,180,160 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.7 MB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0                         │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/GroupNorm_0             │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[512]             │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[512]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,024 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.1 KB)</span>                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv1                   │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv1/conv              │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,512,1024]  │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">4,719,616 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(18.9 MB)</span>            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/GroupNorm_1             │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">2,048 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(8.2 KB)</span>                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv2                   │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv2/conv              │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,1024,1024] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">9,438,208 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(37.8 MB)</span>            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/residual_conv           │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/residual_conv/conv      │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,512]   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,1,512,1024]  │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">525,312 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.1 MB)</span>               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0                    │ AttentionBlock  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/RMSNorm_0          │ RMSNorm         │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,024 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.1 KB)</span>                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention          │ NormalAttention │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_q     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,8,128] │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024,8,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,048,576 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.2 MB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_k     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,8,128] │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024,8,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,048,576 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.2 MB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_v     │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,8,128] │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024,8,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,048,576 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.2 MB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_out_0 │ DenseGeneral    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,8,128] │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[8,128,1024]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,048,576 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.2 MB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0                         │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/GroupNorm_0             │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">2,048 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(8.2 KB)</span>                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv1                   │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv1/conv              │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,1024,1024] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">9,438,208 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(37.8 MB)</span>            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/GroupNorm_1             │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">2,048 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(8.2 KB)</span>                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv2                   │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv2/conv              │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,1024,1024] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">9,438,208 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(37.8 MB)</span>            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_downsample                     │ Downsample      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_downsample/ConvLayer_0         │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_downsample/ConvLayer_0/conv    │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10,10,1024]  │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[2048]            │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,1024,2048] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">18,876,416 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(75.5 MB)</span>           │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual                        │ ResidualBlock   │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/GroupNorm_0            │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[2048]            │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[2048]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">4,096 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(16.4 KB)</span>                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv1                  │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv1/conv             │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,2048,64]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">1,179,712 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(4.7 MB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/GroupNorm_1            │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv2                  │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv2/conv             │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">36,928 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(147.7 KB)</span>              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/residual_conv          │ ConvLayer       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/residual_conv/conv     │ Conv            │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,2048]    │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,1,2048,64]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">131,136 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(524.5 KB)</span>             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ GroupNorm_0                           │ GroupNorm       │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,5,5,64]      │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]              │\n",
       "│                                       │                 │                        │                        │ scale: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">128 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(512 B)</span>                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ Dense_0                               │ Dense           │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,64]          │ <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[1,10]          │ bias: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[10]              │\n",
       "│                                       │                 │                        │                        │ kernel: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span>[64,10]         │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ <span style=\"font-weight: bold\">650 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(2.6 KB)</span>                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│<span style=\"font-weight: bold\">                                       </span>│<span style=\"font-weight: bold\">                 </span>│<span style=\"font-weight: bold\">                        </span>│<span style=\"font-weight: bold\">                  Total </span>│<span style=\"font-weight: bold\"> 63,173,834 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(252.7 MB)</span><span style=\"font-weight: bold\">          </span>│\n",
       "└───────────────────────────────────────┴─────────────────┴────────────────────────┴────────────────────────┴────────────────────────────────┘\n",
       "<span style=\"font-weight: bold\">                                                                                                                                              </span>\n",
       "<span style=\"font-weight: bold\">                                                   Total Parameters: 63,173,834 </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">(252.7 MB)</span><span style=\"font-weight: bold\">                                                    </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                                                              Classifier Summary                                                              \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mpath                                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mmodule         \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1minputs                \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1moutputs               \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mparams                        \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│                                       │ Classifier      │ \u001b[2mfloat32\u001b[0m[1,160,160,3]   │ \u001b[2mfloat32\u001b[0m[1,10]          │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ ConvLayer_0                           │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,160,160,3]   │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ ConvLayer_0/conv                      │ Conv            │ \u001b[2mfloat32\u001b[0m[1,160,160,3]   │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,3,64]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,792 \u001b[0m\u001b[1;2m(7.2 KB)\u001b[0m                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_0/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_residual_1/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_downsample                     │ Downsample      │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_downsample/ConvLayer_0         │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_0_downsample/ConvLayer_0/conv    │ Conv            │ \u001b[2mfloat32\u001b[0m[1,160,160,64]  │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_0/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_residual_1/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1                    │ AttentionBlock  │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/RMSNorm_0          │ RMSNorm         │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m64 \u001b[0m\u001b[1;2m(256 B)\u001b[0m                     │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention          │ NormalAttention │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_q     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,8,8]   │ kernel: \u001b[2mfloat32\u001b[0m[64,8,8]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m4,096 \u001b[0m\u001b[1;2m(16.4 KB)\u001b[0m                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_k     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,8,8]   │ kernel: \u001b[2mfloat32\u001b[0m[64,8,8]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m4,096 \u001b[0m\u001b[1;2m(16.4 KB)\u001b[0m                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_v     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,80,80,8,8]   │ kernel: \u001b[2mfloat32\u001b[0m[64,8,8]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m4,096 \u001b[0m\u001b[1;2m(16.4 KB)\u001b[0m                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_attention_1/Attention/to_out_0 │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,80,80,8,8]   │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ kernel: \u001b[2mfloat32\u001b[0m[8,8,64]        │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m4,096 \u001b[0m\u001b[1;2m(16.4 KB)\u001b[0m                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_downsample                     │ Downsample      │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_downsample/ConvLayer_0         │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_1_downsample/ConvLayer_0/conv    │ Conv            │ \u001b[2mfloat32\u001b[0m[1,80,80,64]    │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m73,856 \u001b[0m\u001b[1;2m(295.4 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m256 \u001b[0m\u001b[1;2m(1.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m147,584 \u001b[0m\u001b[1;2m(590.3 KB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m256 \u001b[0m\u001b[1;2m(1.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_0/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m147,584 \u001b[0m\u001b[1;2m(590.3 KB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m256 \u001b[0m\u001b[1;2m(1.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m147,584 \u001b[0m\u001b[1;2m(590.3 KB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m256 \u001b[0m\u001b[1;2m(1.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_residual_1/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ bias: \u001b[2mfloat32\u001b[0m[128]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,128,128]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m147,584 \u001b[0m\u001b[1;2m(590.3 KB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1                    │ AttentionBlock  │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/RMSNorm_0          │ RMSNorm         │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ scale: \u001b[2mfloat32\u001b[0m[128]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention          │ NormalAttention │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_q     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,8,16]  │ kernel: \u001b[2mfloat32\u001b[0m[128,8,16]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m16,384 \u001b[0m\u001b[1;2m(65.5 KB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_k     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,8,16]  │ kernel: \u001b[2mfloat32\u001b[0m[128,8,16]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m16,384 \u001b[0m\u001b[1;2m(65.5 KB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_v     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,40,40,8,16]  │ kernel: \u001b[2mfloat32\u001b[0m[128,8,16]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m16,384 \u001b[0m\u001b[1;2m(65.5 KB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_attention_1/Attention/to_out_0 │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,40,40,8,16]  │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ kernel: \u001b[2mfloat32\u001b[0m[8,16,128]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m16,384 \u001b[0m\u001b[1;2m(65.5 KB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_downsample                     │ Downsample      │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_downsample/ConvLayer_0         │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_2_downsample/ConvLayer_0/conv    │ Conv            │ \u001b[2mfloat32\u001b[0m[1,40,40,128]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,128,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m295,168 \u001b[0m\u001b[1;2m(1.2 MB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m512 \u001b[0m\u001b[1;2m(2.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m590,080 \u001b[0m\u001b[1;2m(2.4 MB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m512 \u001b[0m\u001b[1;2m(2.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_0/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m590,080 \u001b[0m\u001b[1;2m(2.4 MB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1                     │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/GroupNorm_0         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m512 \u001b[0m\u001b[1;2m(2.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv1               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv1/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m590,080 \u001b[0m\u001b[1;2m(2.4 MB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/GroupNorm_1         │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m512 \u001b[0m\u001b[1;2m(2.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv2               │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_residual_1/conv2/conv          │ Conv            │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ bias: \u001b[2mfloat32\u001b[0m[256]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,256,256]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m590,080 \u001b[0m\u001b[1;2m(2.4 MB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1                    │ AttentionBlock  │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/RMSNorm_0          │ RMSNorm         │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ scale: \u001b[2mfloat32\u001b[0m[256]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m256 \u001b[0m\u001b[1;2m(1.0 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention          │ NormalAttention │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_q     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,8,32]  │ kernel: \u001b[2mfloat32\u001b[0m[256,8,32]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m65,536 \u001b[0m\u001b[1;2m(262.1 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_k     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,8,32]  │ kernel: \u001b[2mfloat32\u001b[0m[256,8,32]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m65,536 \u001b[0m\u001b[1;2m(262.1 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_v     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,20,20,8,32]  │ kernel: \u001b[2mfloat32\u001b[0m[256,8,32]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m65,536 \u001b[0m\u001b[1;2m(262.1 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_attention_1/Attention/to_out_0 │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,20,20,8,32]  │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ kernel: \u001b[2mfloat32\u001b[0m[8,32,256]      │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m65,536 \u001b[0m\u001b[1;2m(262.1 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_downsample                     │ Downsample      │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_downsample/ConvLayer_0         │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ down_3_downsample/ConvLayer_0/conv    │ Conv            │ \u001b[2mfloat32\u001b[0m[1,20,20,256]   │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ bias: \u001b[2mfloat32\u001b[0m[512]             │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,256,512]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,180,160 \u001b[0m\u001b[1;2m(4.7 MB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0                         │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/GroupNorm_0             │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ bias: \u001b[2mfloat32\u001b[0m[512]             │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[512]            │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,024 \u001b[0m\u001b[1;2m(4.1 KB)\u001b[0m                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv1                   │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv1/conv              │ Conv            │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,512,1024]  │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m4,719,616 \u001b[0m\u001b[1;2m(18.9 MB)\u001b[0m            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/GroupNorm_1             │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m2,048 \u001b[0m\u001b[1;2m(8.2 KB)\u001b[0m                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv2                   │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/conv2/conv              │ Conv            │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,1024,1024] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m9,438,208 \u001b[0m\u001b[1;2m(37.8 MB)\u001b[0m            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/residual_conv           │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res1_0/residual_conv/conv      │ Conv            │ \u001b[2mfloat32\u001b[0m[1,10,10,512]   │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[1,1,512,1024]  │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m525,312 \u001b[0m\u001b[1;2m(2.1 MB)\u001b[0m               │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0                    │ AttentionBlock  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/RMSNorm_0          │ RMSNorm         │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ scale: \u001b[2mfloat32\u001b[0m[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,024 \u001b[0m\u001b[1;2m(4.1 KB)\u001b[0m                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention          │ NormalAttention │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_q     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,8,128] │ kernel: \u001b[2mfloat32\u001b[0m[1024,8,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,048,576 \u001b[0m\u001b[1;2m(4.2 MB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_k     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,8,128] │ kernel: \u001b[2mfloat32\u001b[0m[1024,8,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,048,576 \u001b[0m\u001b[1;2m(4.2 MB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_v     │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,8,128] │ kernel: \u001b[2mfloat32\u001b[0m[1024,8,128]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,048,576 \u001b[0m\u001b[1;2m(4.2 MB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_attention_0/Attention/to_out_0 │ DenseGeneral    │ \u001b[2mfloat32\u001b[0m[1,10,10,8,128] │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ kernel: \u001b[2mfloat32\u001b[0m[8,128,1024]    │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,048,576 \u001b[0m\u001b[1;2m(4.2 MB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0                         │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/GroupNorm_0             │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m2,048 \u001b[0m\u001b[1;2m(8.2 KB)\u001b[0m                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv1                   │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv1/conv              │ Conv            │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,1024,1024] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m9,438,208 \u001b[0m\u001b[1;2m(37.8 MB)\u001b[0m            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/GroupNorm_1             │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[1024]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m2,048 \u001b[0m\u001b[1;2m(8.2 KB)\u001b[0m                 │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv2                   │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_res2_0/conv2/conv              │ Conv            │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ bias: \u001b[2mfloat32\u001b[0m[1024]            │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,1024,1024] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m9,438,208 \u001b[0m\u001b[1;2m(37.8 MB)\u001b[0m            │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_downsample                     │ Downsample      │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_downsample/ConvLayer_0         │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ middle_downsample/ConvLayer_0/conv    │ Conv            │ \u001b[2mfloat32\u001b[0m[1,10,10,1024]  │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ bias: \u001b[2mfloat32\u001b[0m[2048]            │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,1024,2048] │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m18,876,416 \u001b[0m\u001b[1;2m(75.5 MB)\u001b[0m           │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual                        │ ResidualBlock   │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/GroupNorm_0            │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ bias: \u001b[2mfloat32\u001b[0m[2048]            │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[2048]           │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m4,096 \u001b[0m\u001b[1;2m(16.4 KB)\u001b[0m                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv1                  │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv1/conv             │ Conv            │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,2048,64]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m1,179,712 \u001b[0m\u001b[1;2m(4.7 MB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/GroupNorm_1            │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv2                  │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/conv2/conv             │ Conv            │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[3,3,64,64]     │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m36,928 \u001b[0m\u001b[1;2m(147.7 KB)\u001b[0m              │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/residual_conv          │ ConvLayer       │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │                                │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ final_residual/residual_conv/conv     │ Conv            │ \u001b[2mfloat32\u001b[0m[1,5,5,2048]    │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[1,1,2048,64]   │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m131,136 \u001b[0m\u001b[1;2m(524.5 KB)\u001b[0m             │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ GroupNorm_0                           │ GroupNorm       │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ \u001b[2mfloat32\u001b[0m[1,5,5,64]      │ bias: \u001b[2mfloat32\u001b[0m[64]              │\n",
       "│                                       │                 │                        │                        │ scale: \u001b[2mfloat32\u001b[0m[64]             │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m128 \u001b[0m\u001b[1;2m(512 B)\u001b[0m                    │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│ Dense_0                               │ Dense           │ \u001b[2mfloat32\u001b[0m[1,64]          │ \u001b[2mfloat32\u001b[0m[1,10]          │ bias: \u001b[2mfloat32\u001b[0m[10]              │\n",
       "│                                       │                 │                        │                        │ kernel: \u001b[2mfloat32\u001b[0m[64,10]         │\n",
       "│                                       │                 │                        │                        │                                │\n",
       "│                                       │                 │                        │                        │ \u001b[1m650 \u001b[0m\u001b[1;2m(2.6 KB)\u001b[0m                   │\n",
       "├───────────────────────────────────────┼─────────────────┼────────────────────────┼────────────────────────┼────────────────────────────────┤\n",
       "│\u001b[1m \u001b[0m\u001b[1m                                     \u001b[0m\u001b[1m \u001b[0m│\u001b[1m \u001b[0m\u001b[1m               \u001b[0m\u001b[1m \u001b[0m│\u001b[1m \u001b[0m\u001b[1m                      \u001b[0m\u001b[1m \u001b[0m│\u001b[1m \u001b[0m\u001b[1m                 Total\u001b[0m\u001b[1m \u001b[0m│\u001b[1m \u001b[0m\u001b[1m63,173,834 \u001b[0m\u001b[1;2m(252.7 MB)\u001b[0m\u001b[1m         \u001b[0m\u001b[1m \u001b[0m│\n",
       "└───────────────────────────────────────┴─────────────────┴────────────────────────┴────────────────────────┴────────────────────────────────┘\n",
       "\u001b[1m                                                                                                                                              \u001b[0m\n",
       "\u001b[1m                                                   Total Parameters: 63,173,834 \u001b[0m\u001b[1;2m(252.7 MB)\u001b[0m\u001b[1m                                                    \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = Classifier(\n",
    "                  feature_depths=[64, 128, 256, 512, 1024],\n",
    "                  attention_configs=[None, {\"heads\": 8}, {\"heads\": 8}, {\"heads\": 8}],\n",
    "                  num_res_blocks=2,\n",
    "                  num_middle_res_blocks=1,\n",
    "                  major_conv_type=\"conv\"\n",
    "            )\n",
    "\n",
    "inp = jnp.ones((1, 160, 160, 3))\n",
    "print(model.tabulate(jax.random.key(0), inp,\n",
    "      console_kwargs={\"width\": 200, \"force_jupyter\": True, }))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate by August 1st, 2024.\n"
     ]
    }
   ],
   "source": [
    "# cnn = CNN()\n",
    "model = Classifier(\n",
    "                  feature_depths=[64, 128, 256, 512, 1024],\n",
    "                  attention_configs=[None, {\"heads\": 8}, {\"heads\": 8}, {\"heads\": 8}],\n",
    "                  num_res_blocks=2,\n",
    "                  num_middle_res_blocks=1,\n",
    "                  major_conv_type=\"conv\"\n",
    "            )\n",
    "solver = optax.adam(1e-3)\n",
    "\n",
    "def loss(preds, labels):\n",
    "      # lambda preds, targets: optax.softmax_cross_entropy_with_integer_labels(preds, targets)\n",
    "      # print(preds.shape, labels.shape)\n",
    "      result = optax.softmax_cross_entropy_with_integer_labels(preds, labels)\n",
    "      return result\n",
    "      \n",
    "trainer = SimpleTrainer(model, \n",
    "                        input_shapes={'x': (160, 160, 3)}, \n",
    "                        optimizer=solver, \n",
    "                        loss_fn=loss,\n",
    "                        rngs=jax.random.PRNGKey(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "jax.profiler.stop_server()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'jax' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax\u001b[49m\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mstart_server(\u001b[38;5;241m6009\u001b[39m)\n\u001b[1;32m      2\u001b[0m data \u001b[38;5;241m=\u001b[39m get_dataset(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mimagenette\u001b[39m\u001b[38;5;124m\"\u001b[39m, batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m256\u001b[39m, splits\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvalidation\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m      3\u001b[0m state \u001b[38;5;241m=\u001b[39m trainer\u001b[38;5;241m.\u001b[39mfit(data, steps_per_epoch\u001b[38;5;241m=\u001b[39mdata[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain_len\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m data[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mbatch_size\u001b[39m\u001b[38;5;124m'\u001b[39m], epochs\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'jax' is not defined"
     ]
    }
   ],
   "source": [
    "jax.profiler.start_server(6009)\n",
    "data = get_dataset(\"imagenette\", batch_size=256, splits=['train', 'validation'])\n",
    "state = trainer.fit(data, steps_per_epoch=data['train_len'] // data['batch_size'], epochs=10)\n",
    "jax.profiler.stop_server()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 1: 100step [00:56,  1.77step/s, loss=2.3902]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 1\n",
      "\n",
      "\tEpoch 1 completed. Avg Loss: 2.1881344318389893, Time: 56.66s, Best Loss: 2.1881344318389893 \n",
      "\n",
      "Epoch 2/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 2: 100step [00:05, 18.16step/s, loss=2.2113]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 2\n",
      "\n",
      "\tEpoch 2 completed. Avg Loss: 2.1151442527770996, Time: 5.51s, Best Loss: 2.1151442527770996 \n",
      "\n",
      "Epoch 3/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 3:   0%|                                              | 0/73 [00:00<?, ?step/s, loss=2.1399]2024-07-29 20:52:01.746798: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 2635800 nanoseconds and will start immediately.\n",
      "\t\tEpoch 3: 100step [00:05, 17.55step/s, loss=2.1399]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 3\n",
      "\n",
      "\tEpoch 3 completed. Avg Loss: 2.0795538425445557, Time: 5.70s, Best Loss: 2.0795538425445557 \n",
      "\n",
      "Epoch 4/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 4: 100step [00:05, 17.55step/s, loss=2.0495]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 4\n",
      "\n",
      "\tEpoch 4 completed. Avg Loss: 2.0166125297546387, Time: 5.70s, Best Loss: 2.0166125297546387 \n",
      "\n",
      "Epoch 5/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 5: 100step [00:05, 18.21step/s, loss=1.8212]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 5\n",
      "\n",
      "\tEpoch 5 completed. Avg Loss: 1.9245163202285767, Time: 5.49s, Best Loss: 1.9245163202285767 \n",
      "\n",
      "Epoch 6/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 6: 100step [00:05, 18.23step/s, loss=1.8679]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 6\n",
      "\n",
      "\tEpoch 6 completed. Avg Loss: 1.8269599676132202, Time: 5.49s, Best Loss: 1.8269599676132202 \n",
      "\n",
      "Epoch 7/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 7: 100step [00:05, 18.27step/s, loss=1.7161]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 7\n",
      "\n",
      "\tEpoch 7 completed. Avg Loss: 1.6975194215774536, Time: 5.48s, Best Loss: 1.6975194215774536 \n",
      "\n",
      "Epoch 8/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 8: 100step [00:05, 18.36step/s, loss=1.7334]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 8\n",
      "\n",
      "\tEpoch 8 completed. Avg Loss: 1.6169333457946777, Time: 5.45s, Best Loss: 1.6169333457946777 \n",
      "\n",
      "Epoch 9/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 9: 100step [00:05, 18.25step/s, loss=1.5885]                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 9\n",
      "\n",
      "\tEpoch 9 completed. Avg Loss: 1.5265167951583862, Time: 5.48s, Best Loss: 1.5265167951583862 \n",
      "\n",
      "Epoch 10/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 10: 100step [00:05, 18.30step/s, loss=1.4427]                                               "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 10\n",
      "\n",
      "\tEpoch 10 completed. Avg Loss: 1.4501495361328125, Time: 5.47s, Best Loss: 1.4501495361328125 \n",
      "Saving model at epoch 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "jax.profiler.start_server(6009)\n",
    "data = get_dataset(\"imagenette\", batch_size=128, splits=['train', 'validation'])\n",
    "state = trainer.fit(data, steps_per_epoch=data['train_len'] // data['batch_size'], epochs=10)\n",
    "jax.profiler.stop_server()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 1:   0%|                                                          | 0/468 [00:00<?, ?step/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 1: 500step [00:04, 114.14step/s, loss=0.0888]                                               \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 1\n",
      "\n",
      "\tEpoch 1 completed. Avg Loss: 0.18133877217769623, Time: 4.38s, Best Loss: 0.18133877217769623 \n",
      "\n",
      "Epoch 2/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 2: 500step [00:00, 1287.91step/s, loss=0.0585]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 2\n",
      "\n",
      "\tEpoch 2 completed. Avg Loss: 0.05323795974254608, Time: 0.39s, Best Loss: 0.05323795974254608 \n",
      "\n",
      "Epoch 3/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 3: 500step [00:00, 1323.98step/s, loss=0.0094]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 3\n",
      "\n",
      "\tEpoch 3 completed. Avg Loss: 0.03695608302950859, Time: 0.38s, Best Loss: 0.03695608302950859 \n",
      "\n",
      "Epoch 4/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 4: 500step [00:00, 1280.80step/s, loss=0.0233]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 4\n",
      "\n",
      "\tEpoch 4 completed. Avg Loss: 0.02820330113172531, Time: 0.39s, Best Loss: 0.02820330113172531 \n",
      "\n",
      "Epoch 5/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 5: 500step [00:00, 1298.43step/s, loss=0.0154]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 5\n",
      "\n",
      "\tEpoch 5 completed. Avg Loss: 0.02226649969816208, Time: 0.39s, Best Loss: 0.02226649969816208 \n",
      "\n",
      "Epoch 6/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 6: 500step [00:00, 1353.08step/s, loss=0.0297]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 6\n",
      "\n",
      "\tEpoch 6 completed. Avg Loss: 0.019302448257803917, Time: 0.37s, Best Loss: 0.019302448257803917 \n",
      "\n",
      "Epoch 7/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 7: 500step [00:00, 1381.89step/s, loss=0.0008]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 7\n",
      "\n",
      "\tEpoch 7 completed. Avg Loss: 0.014166121371090412, Time: 0.36s, Best Loss: 0.014166121371090412 \n",
      "\n",
      "Epoch 8/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 8: 500step [00:00, 1362.68step/s, loss=0.0132]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 8\n",
      "\n",
      "\tEpoch 8 completed. Avg Loss: 0.01146089006215334, Time: 0.37s, Best Loss: 0.01146089006215334 \n",
      "\n",
      "Epoch 9/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 9: 500step [00:00, 1345.12step/s, loss=0.0053]                                              \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 9\n",
      "\n",
      "\tEpoch 9 completed. Avg Loss: 0.011206524446606636, Time: 0.37s, Best Loss: 0.011206524446606636 \n",
      "\n",
      "Epoch 10/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 10: 500step [00:00, 1353.41step/s, loss=0.0065]                                             "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving model at epoch 10\n",
      "\n",
      "\tEpoch 10 completed. Avg Loss: 0.008855808526277542, Time: 0.37s, Best Loss: 0.008855808526277542 \n",
      "Saving model at epoch 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# jax.profiler.start_server(6009)\n",
    "data = get_dataset(\"mnist\", batch_size=128)\n",
    "state = trainer.fit(data, steps_per_epoch=data['train_len'] // data['batch_size'], epochs=10)\n",
    "# jax.profiler.stop_server()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}