Inference after finetuning

[RamishSiddiqui]

Hi everyone,

I finetune the base model using the bash script provided here peft/finetune.sh using these parameters,

python3 finetune.py \
    --model-name="google/timesfm-1.0-200m" \
    --checkpoint-path="./model-checkpoints-paxml/checkpoints" \
    --checkpoint-dir="./finetune-checkpoints/" \
    --backend="gpu" \
    --horizon-len=32 \
    --context-len=128 \
    --batch-size=64 \
    --freq="D" \
    --data-path="../data/train.csv" \
    --num-epochs=100 \
    --learning-rate=1e-3 \
    --adam-epsilon=1e-7 \
    --adam-clip-threshold=1e2 \
    --early-stop-patience=10 \
    --datetime-col="ds" \
    --use-lora \
    --lora-rank=1 \
    --lora-target-modules="all" \
    --use-dora \
    --cos-initial-decay-value=1e-4 \
    --cos-decay-steps=40000 \
    --cos-final-decay-value=1e-5 \
    --ema-decay=0.9999 \
    --use-linear-probing

After finetuning i used this script to load the model,

context_len = 128
horizon_len = 32
INPUT_PATCH_LEN = 32
OUTPUT_PATCH_LEN = 128
NUM_LAYERS = 20
MODEL_DIMS = 1280
backend = "gpu"
batch_size = 64

# For JAX
tfm = timesfm.TimesFm(
    hparams=timesfm.TimesFmHparams(
        context_len=context_len,
        horizon_len=horizon_len,
        backend=backend,
        input_patch_len=INPUT_PATCH_LEN,
        output_patch_len=OUTPUT_PATCH_LEN,
        per_core_batch_size=batch_size,
        num_layers=NUM_LAYERS,
        model_dims=MODEL_DIMS,
    ),
    checkpoint=timesfm.TimesFmCheckpoint(
        # path=os.path.join(os.getcwd(), "peft/model-checkpoints-paxml/checkpoints/"),
        path=os.path.join(os.getcwd(), "peft/finetune-checkpoints/run_20241028_140830_tga4qg8v/"),
        # huggingface_repo_id="google/timesfm-1.0-200m-pytorch"
    ),
)

load_adapter_checkpoint(
    model=tfm,
    adapter_checkpoint_path="./checkpoints/run_20241028_140830_tga4qg8v",
    lora_rank=1,
    lora_target_modules="all",
    use_dora=True,
)

## Load Data here

Y_hat_df = tfm.forecast_on_df(
    inputs=Y_test_df,
    freq="D",  # Daily
    value_name="y",
    num_jobs=1,
)

and i got this following error.

/home/ec2-user/.cache/pypoetry/virtualenvs/timesfm-5aVyK813-py3.10/bin/python /home/ec2-user/timesfm-git/run_inference.py 
TimesFM v1.2.0. See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs.
Loaded Jax TimesFM.
2024-10-28 17:38:46.847744: 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
2024-10-28 17:38:57.828867: W external/xla/xla/service/gpu/nvptx_compiler.cc:718] The NVIDIA driver's CUDA version is 12.1 which is older than the ptxas CUDA version (12.6.77). 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.
Constructing model weights.
WARNING:absl:No registered CheckpointArgs found for handler type: <class 'paxml.checkpoints.FlaxCheckpointHandler'>
Constructed model weights in 2.76 seconds.
Restoring checkpoint from /home/ec2-user/timesfm-git/peft/finetune-checkpoints/run_20241028_140830_tga4qg8v/.
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.
WARNING:absl:train_state_unpadded_shape_dtype_struct is not provided. We assume `train_state` is unpadded.
Restored checkpoint in 0.39 seconds.
Jitting decoding.
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[()], [*, {'params': {'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': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}}}}, [], ()]))).
ERROR:absl:For checkpoint version > 1.0, we require users to provide
          `train_state_unpadded_shape_dtype_struct` during checkpoint
          saving/restoring, to avoid potential silent bugs when loading
          checkpoints to incompatible unpadded shapes of TrainState.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ec2-user/timesfm-git/run_inference.py", line 74, in <module>
    tfm = timesfm.TimesFm(
  File "/home/ec2-user/timesfm-git/src/timesfm/timesfm_base.py", line 180, in __init__
    self.load_from_checkpoint(checkpoint)
  File "/home/ec2-user/timesfm-git/src/timesfm/timesfm_jax.py", line 174, in load_from_checkpoint
    self.jit_decode()
  File "/home/ec2-user/timesfm-git/src/timesfm/timesfm_jax.py", line 207, in jit_decode
    _ = self._pmapped_decode(
  File "/home/ec2-user/timesfm-git/src/timesfm/timesfm_jax.py", line 183, in _decode
    return self._model.apply(
  File "/home/ec2-user/.cache/pypoetry/virtualenvs/timesfm-5aVyK813-py3.10/lib/python3.10/site-packages/praxis/base_layer.py", line 1968, in apply
    result = super().apply(
  File "/home/ec2-user/timesfm-git/src/timesfm/patched_decoder.py", line 450, in decode
    fprop_outputs = self(model_input)[_OUTPUT_TS]
  File "/home/ec2-user/timesfm-git/src/timesfm/patched_decoder.py", line 379, in __call__
    model_input, patched_padding, stats, _ = self._preprocess_input(
  File "/home/ec2-user/timesfm-git/src/timesfm/patched_decoder.py", line 332, in _preprocess_input
    model_input = self.input_ff_layer(concat_inputs)
  File "/home/ec2-user/timesfm-git/src/timesfm/patched_decoder.py", line 147, in __call__
    hidden = self.hidden_layer(inputs)
  File "/home/ec2-user/.cache/pypoetry/virtualenvs/timesfm-5aVyK813-py3.10/lib/python3.10/site-packages/praxis/layers/linears.py", line 222, in __call__
    projected_inputs = self.linear(inputs)
  File "/home/ec2-user/.cache/pypoetry/virtualenvs/timesfm-5aVyK813-py3.10/lib/python3.10/site-packages/praxis/base_layer.py", line 1901, in _try_setup
    super()._try_setup(shallow=shallow)
  File "/home/ec2-user/.cache/pypoetry/virtualenvs/timesfm-5aVyK813-py3.10/lib/python3.10/site-packages/praxis/layers/linears.py", line 88, in setup
    self.create_variable(
  File "/home/ec2-user/.cache/pypoetry/virtualenvs/timesfm-5aVyK813-py3.10/lib/python3.10/site-packages/praxis/base_layer.py", line 2453, in create_variable
    self.param(name, _initializer_fn)
flax.errors.ScopeParamShapeError: Initializer expected to generate shape (1280, 1) but got shape (64, 1280) instead for parameter "w" in "/input_ff_layer/hidden_layer/linear". (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeParamShapeError)

Process finished with exit code 1

Is there anything i am missing here ?

[rajatsen91]

@tanmayshishodia can you take a look if you have time?