{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Base Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-07-16 00:36:17.861915: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-07-16 00:36:25.276917: W external/xla/xla/service/gpu/nvptx_compiler.cc:718] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "935a82ff713149179c103451e4837c27",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Constructing model weights.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No registered CheckpointArgs found for handler type: <class 'paxml.checkpoints.FlaxCheckpointHandler'>\n",
      "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 May 1st, 2024.\n",
      "WARNING:absl:train_state_unpadded_shape_dtype_struct is not provided. We assume `train_state` is unpadded.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Constructed model weights in 2.84 seconds.\n",
      "Restoring checkpoint from /home/ubuntu/.cache/huggingface/hub/models--google--timesfm-1.0-200m/snapshots/8775f7531211ac864b739fe776b0b255c277e2be/checkpoints.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR:absl:For checkpoint version > 1.0, we require users to provide\n",
      "          `train_state_unpadded_shape_dtype_struct` during checkpoint\n",
      "          saving/restoring, to avoid potential silent bugs when loading\n",
      "          checkpoints to incompatible unpadded shapes of TrainState.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restored checkpoint in 0.86 seconds.\n",
      "Jitting decoding.\n",
      "Jitted decoding in 16.44 seconds.\n"
     ]
    }
   ],
   "source": [
    "from timesfm import TimesFm, freq_map, data_loader\n",
    "from adapter.utils import load_adapter_checkpoint\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "tfm = TimesFm(\n",
    "    context_len=512,\n",
    "    horizon_len=128,\n",
    "    input_patch_len=32,\n",
    "    output_patch_len=128,\n",
    "    num_layers=20,\n",
    "    model_dims=1280,\n",
    "    backend=\"cpu\",\n",
    ")\n",
    "tfm.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DICT = {\n",
    "    \"ettm2\": {\n",
    "        \"boundaries\": [34560, 46080, 57600],\n",
    "        \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
    "        \"freq\": \"15min\",\n",
    "    },\n",
    "    \"ettm1\": {\n",
    "        \"boundaries\": [34560, 46080, 57600],\n",
    "        \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
    "        \"freq\": \"15min\",\n",
    "    },\n",
    "    \"etth2\": {\n",
    "        \"boundaries\": [8640, 11520, 14400],\n",
    "        \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
    "        \"freq\": \"H\",\n",
    "    },\n",
    "    \"etth1\": {\n",
    "        \"boundaries\": [8640, 11520, 14400],\n",
    "        \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
    "        \"freq\": \"H\",\n",
    "    },\n",
    "    \"elec\": {\n",
    "        \"boundaries\": [18413, 21044, 26304],\n",
    "        \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
    "        \"freq\": \"H\",\n",
    "    },\n",
    "    \"traffic\": {\n",
    "        \"boundaries\": [12280, 14036, 17544],\n",
    "        \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
    "        \"freq\": \"H\",\n",
    "    },\n",
    "    \"weather\": {\n",
    "        \"boundaries\": [36887, 42157, 52696],\n",
    "        \"data_path\": \"../datasets/weather/weather.csv\",\n",
    "        \"freq\": \"10min\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tfm._train_state.mdl_vars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Adapter Checkpoint\n",
    "\n",
    "Specify the adapter checkpoint path, rank and the modules used to train the adapters and whether dora was employed or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restoring adapter checkpoint from /home/ubuntu/tanmay/timesfm/adapter_finetuning/checkpoints/run_ignore_20240715_234116.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No registered CheckpointArgs found for handler type: <class 'paxml.checkpoints.FlaxCheckpointHandler'>\n",
      "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 May 1st, 2024.\n",
      "WARNING:absl:train_state_unpadded_shape_dtype_struct is not provided. We assume `train_state` is unpadded.\n",
      "WARNING:absl:A possible mismatch (could be spurious) between the saved checkpoint structure and the current one has been detected (PyTreeDef(CustomNode(TrainState[()], [*, {'params': {'x_layers_0': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_1': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_10': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_11': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_12': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_13': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_14': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_15': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_16': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_17': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_18': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_19': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_2': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_3': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_4': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_5': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_6': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_7': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_8': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_9': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}}}, [(CustomNode(NestedMap[('count',)], [*]), CustomNode(NestedMap[('count',)], [*]), CustomNode(NestedMap[('count', 'm', 'v')], [*, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}}}}}, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}}}}}]), CustomNode(NestedMap[('count',)], [*]), CustomNode(NestedMap[('count', 'ema')], [*, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}}}}}]))], ()])) vs PyTreeDef(CustomNode(TrainState[()], [*, {'x_layers_0': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_1': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_10': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_11': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_12': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_13': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_14': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_15': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_16': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_17': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_18': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_19': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_2': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_3': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_4': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_5': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_6': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_7': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_8': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_9': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}}, [], ()]))).\n",
      "ERROR:absl:For checkpoint version > 1.0, we require users to provide\n",
      "          `train_state_unpadded_shape_dtype_struct` during checkpoint\n",
      "          saving/restoring, to avoid potential silent bugs when loading\n",
      "          checkpoints to incompatible unpadded shapes of TrainState.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restored adapter checkpoint in 8.77 seconds.\n",
      "Jitting decoding.\n",
      "Jitted decoding in 14.77 seconds.\n"
     ]
    }
   ],
   "source": [
    "load_adapter_checkpoint(\n",
    "    model=tfm,\n",
    "    adapter_checkpoint_path=\"/home/ubuntu/tanmay/timesfm/adapter_finetuning/checkpoints/run_ignore_20240715_234116\",\n",
    "    lora_rank=1,\n",
    "    lora_target_modules=\"all\",\n",
    "    use_dora=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = \"ettm1\"\n",
    "data_path = DATA_DICT[dataset][\"data_path\"]\n",
    "freq = DATA_DICT[dataset][\"freq\"]\n",
    "int_freq = freq_map(freq)\n",
    "boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
    "\n",
    "data_df = pd.read_csv(open(data_path, \"r\"))\n",
    "\n",
    "ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
    "num_cov_cols = None\n",
    "cat_cov_cols = None\n",
    "\n",
    "context_len = 512\n",
    "pred_len = 96\n",
    "\n",
    "num_ts = len(ts_cols)\n",
    "batch_size = 16\n",
    "\n",
    "dtl = data_loader.TimeSeriesdata(\n",
    "    data_path=data_path,\n",
    "    datetime_col=\"date\",\n",
    "    num_cov_cols=num_cov_cols,\n",
    "    cat_cov_cols=cat_cov_cols,\n",
    "    ts_cols=np.array(ts_cols),\n",
    "    train_range=[0, boundaries[0]],\n",
    "    val_range=[boundaries[0], boundaries[1]],\n",
    "    test_range=[boundaries[1], boundaries[2]],\n",
    "    hist_len=context_len,\n",
    "    pred_len=pred_len,\n",
    "    batch_size=num_ts,\n",
    "    freq=\"15min\",\n",
    "    normalize=True,\n",
    "    epoch_len=None,\n",
    "    holiday=False,\n",
    "    permute=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-07-16 00:37:10.346547: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-07-16 00:37:10.346626: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory\n",
      "2024-07-16 00:37:10.346677: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory\n",
      "2024-07-16 00:37:10.346724: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory\n",
      "2024-07-16 00:37:10.350324: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory\n",
      "2024-07-16 00:37:10.350358: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n"
     ]
    }
   ],
   "source": [
    "test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "120it [01:20,  1.50it/s]\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\">MAE: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.32633697986602783</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "MAE: \u001b[1;36m0.32633697986602783\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mae_losses = []\n",
    "for batch in tqdm(test_batches.as_numpy_iterator()):\n",
    "    past = batch[0]\n",
    "    actuals = batch[3]\n",
    "    _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
    "    forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
    "    mae_losses.append(np.abs(forecasts - actuals).mean())\n",
    "\n",
    "print(f\"MAE: {np.mean(mae_losses)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tanmay_tfm_env",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}