[tab] /home/franz.srambical/.cache/huggingface/hub/datasets--p-doom--crowd-code-0.1/snapshots/7d826468379cf4bce9a2c819029f20e7cfe52d44/1de052c516cab686515c107385aaf7c3a7e3e5c23c9bc3c0be0cff3df28cd64d/crowd-code-2f4ff312-abac-4732-805a-ec6c245e532e1752745206273-2025_07_17-11.40.45.491/source.csv (plaintext)
Sequence,Time,File,RangeOffset,RangeLength,Text,Language,Type
98,1928682,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3352437.log ",,terminal_command
99,1928731,"TERMINAL",0,0,"",,terminal_output
100,1928792,"TERMINAL",0,0,"=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=2048 \
    --tokenizer_num_blocks=4 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2571288
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0519
SLURM_JOB_START_TIME=1752692957
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752779357
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3352437
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0519
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0519-0520]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1993.localdomain
SLURM_JOB_ID=3352437
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3352437.log",,terminal_output
101,1929130,"TERMINAL",0,0,"
",,terminal_output
102,1929258,"TERMINAL",0,0,"...skipping...
ry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_t ype='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_ type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Ar ray', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='ja x.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array',  skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array ', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_de serialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMeta dataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(va lue_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataE ntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_ty pe='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value _type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Ar ray', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax .Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip _deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_ty pe='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Ar ray', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.A rray', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_ deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', s kip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deseri alize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_des erialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=F alse, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax .Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', s kip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array',  skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deseria lize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_des erialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=F alse, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserializ e=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, w rite_shape=(64, 8, 64))}}})}}}}, None, None]}
pe=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 5 12))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 6 4)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))},  'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'k ernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STB lock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize= False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write _shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False , write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape =(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_sh ape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))},  'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)) , 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHead Attention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape =(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 6 4))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)),  'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'que ry': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kerne l': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {' bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': Val ueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttenti on_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64 )), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, ' out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel ': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {' bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': Val ueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetad ataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}}}, None, None]}
ESC[1;34mwandbESC[0m: 
ESC[1;34mwandbESC[0m: 🚀 View run ESC[33mdynamics-lr-1e-4-larger-ffnESC[0m at: ESC[34mhttps://wandb.ai/instan t-uv/jafar/runs/jktj3wowESC[0m
ESC[1;34mwandbESC[0m: Find logs at: ESC[1;35m../../../../hkfs/home/project/hk-project-p0023960/tum_dbd0378/ja far/wandb/run-20250716_211038-jktj3wow/logsESC[0m
srun: error: hkn0519: task 1: Exited with exit code 1
srun: error: hkn0519: task 2: Exited with exit code 1
srun: error: hkn0520: tasks 4-6: Exited with exit code 1
srun: error: hkn0520: task 7: Exited with exit code 1
srun: error: hkn0519: tasks 0,3: Exited with exit code 1

============================= JOB FEEDBACK =============================

Job ID: 3352437
Cluster: hk
User/Group: tum_dbd0378/hk-project-p0023960
Account: hk-project-p0023960
State: FAILED (exit code 1)
Partition: accelerated
Nodes: 2
Cores per node: 24
Nodelist: hkn[0519-0520]
CPU Utilized: 00:16:08
CPU Efficiency: 11.52% of 02:20:00 core-walltime
Job Wall-clock time: 00:02:55
Starttime: Wed Jul 16 21:09:17 2025
Endtime: Wed Jul 16 21:12:12 2025
Memory Utilized: 20.57 GB (estimated maximum)
Memory Efficiency: 0.00% of 0.00 MB (0.00 MB/node)
Energy Consumed: 192622 Joule / 53.5061111111111 Watthours
Average node power draw: 1100.69714285714 Watt
(END)",,terminal_output
103,1954250,"TERMINAL",0,0,"
MSLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
MSLURM_CONF=/etc/slurm/slurm.conf
MSLURM_NODEID=0
MSLURM_JOB_ID=3352437
MSLURM_SUBMIT_HOST=hkn1993.localdomain
MSLURM_NNODES=2
MSLURM_NPROCS=8
MSLURM_PRIO_PROCESS=0
MSLURM_JOB_ACCOUNT=hk-project-p0023960
MSLURM_NODELIST=hkn[0519-0520]
MSLURM_SCRIPT_CONTEXT=prolog_task
MSLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
MSLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0519
MSLURM_NTASKS=8
MSLURM_CPUS_PER_TASK=5
MSLURM_PROCID=0
MSLURM_JOB_QOS=normal
MSLURM_JOBID=3352437
MSLURM_JOB_NUM_NODES=2
MSLURM_OOM_KILL_STEP=0
MSLURM_TRES_PER_TASK=cpu=5
MSLURM_JOB_PARTITION=accelerated
MSLURM_GTIDS=0
MSLURM_GPUS_ON_NODE=4
MSLURM_JOB_CPUS_PER_NODE=24(x2)
MSLURM_CPUS_ON_NODE=24
MSLURM_JOB_END_TIME=1752779357
MSLURM_CLUSTER_NAME=hk
MSLURM_JOB_START_TIME=1752692957
MSLURMD_NODENAME=hkn0519
MSLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
MSLURM_LOCALID=0
MSLURM_JOB_GPUS=0,1,2,3
MSLURM_TASK_PID=2571288
MSLURM_JOB_UID=996262
MSLURM_TASKS_PER_NODE=4(x2)
MSLURM_JOB_USER=tum_dbd0378
M    --data_dir $array_records_dir \
M    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
M    --project jafar \
M    --entity instant-uv \
M    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
M    --name=dynamics-lr-1e-4-larger-ffn \
M    --log_checkpoint_interval=1000 \
M    --log \
M    --log_image_interval=1000 \
M    --max_lr=1.5e-4 \
M    --init_lr=0 \
M    --dyna_num_blocks=6 \
M    --dyna_ffn_dim=2048 \
M    --lam_num_blocks=4 \
M    --lam_ffn_dim=2048 \
M    --tokenizer_num_blocks=4 \
M    --tokenizer_ffn_dim=2048 \
M    --batch_size=96 \
M    --ckpt_dir $CHECKPOINT_DIR \
M    --save_ckpt \
Msrun python train_dynamics.py \
M
Menv | grep SLURM
M
Mscaling/train_tokenizer_lr_sweep_1e-4
Mtokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-M
Mmkdir -p $CHECKPOINT_DIR
Mng-modelsize-scaling/$job_name
MCHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotrainiM
Mslurm_job_id=$SLURM_JOB_ID
Mjob_name=$SLURM_JOB_NAME
M
Mords_chunked
Marray_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecM
Msource .venv/bin/activate
Mmodule unload devel/cuda/12.4
Mmodule unload mpi/openmpi/5.0
M
Mcat $0
M# Log the sbatch script
M
M#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn
M#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
M#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
M#SBATCH --gres=gpu:4
M#SBATCH --cpus-per-task=5
M#SBATCH --partition=accelerated
M#SBATCH --time=24:00:00
M#SBATCH --ntasks-per-node=4
M#SBATCH --nodes=2
M
M#!/usr/bin/env bash

:",,terminal_output
104,1954385,"TERMINAL",0,0,"

:",,terminal_output
105,1957531,"TERMINAL",0,0,"
SLURM_NTASKS_PER_NODE=4
SLURM_JOB_GID=502226
SLURM_JOB_NODELIST=hkn[0519-0520]
GpuFreq=control_disabled
GpuFreq=control_disabled
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val :",,terminal_output
106,1961728,"TERMINAL",0,0,"
ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
2025-07-16 21:10:00.136052: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we :",,terminal_output
107,1964480,"TERMINAL",0,0,"
re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.136050: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.147869: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.148033: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.269210: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.272749: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.287019: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.287044: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:26.516484: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:27.061049: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:27.166309: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:27.593704: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:29.445859: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:29.880469: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:30.345131: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:30.928066: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
:",,terminal_output
108,1965484,"TERMINAL",0,0,"
2025-07-16 21:10:31.827307: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:32.318324: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:32.468961: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:32.744868: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
2025-07-16 21:10:35.335235: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:35.404405: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:35.955390: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:36.822663: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
wandb: Currently logged in as: emergenz (p-doom) to https://api.wandb.ai. Use `wandb login --relogin` to forc e relogin
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
wandb: Tracking run with wandb version 0.19.11
wandb: Run data is saved locally in /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/wandb/run-202507 :",,terminal_output
109,1978330,"TERMINAL",0,0,"
16_211038-jktj3wow
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run dynamics-lr-1e-4-larger-ffn
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/jktj3wow
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca :",,terminal_output
110,1987507,"TERMINAL",0,0,"
ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
:",,terminal_output
111,1988012,"TERMINAL",0,0,"
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
:",,terminal_output
112,1988358,"TERMINAL",0,0,"
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca :",,terminal_output
113,1988826,"TERMINAL",0,0,"
ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
:",,terminal_output
114,1989109,"TERMINAL",0,0,"
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:",,terminal_output
115,1989360,"TERMINAL",0,0,"
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
:",,terminal_output
116,2014836,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py",0,0,"from dataclasses import dataclass, field
import os

import einops
from flax.training.train_state import TrainState
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh
import optax
import orbax.checkpoint as ocp
import numpy as np
import dm_pix as pix
import jax
import jax.numpy as jnp
import tyro
import wandb
import grain

from genie import Genie, restore_genie_components
from utils.dataloader import get_dataloader
from utils.lr_utils import get_lr_schedule
from utils.parameter_utils import count_parameters_by_component

@dataclass
class Args:
    # Experiment
    num_steps: int = 200_000
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = """"
    save_ckpt: bool = False
    restore_ckpt: bool = False
    # Optimization
    batch_size: int = 36
    init_lr: float = 0.0
    max_lr: float = 3e-5
    decay_end: float = 0.0
    wsd_decay_steps: int = 10000 # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    warmup_steps: int = 5000
    lr_schedule : str = ""wsd"" # supported options: wsd, cos
    # Tokenizer
    tokenizer_dim: int = 512
    tokenizer_ffn_dim: int = 2048
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 4
    tokenizer_num_heads: int = 8
    tokenizer_checkpoint: str = """"
    # LAM
    lam_dim: int = 512
    lam_ffn_dim: int = 2048
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 4
    lam_num_heads: int = 8
    lam_checkpoint: str = """"
    # Dynamics
    dyna_dim: int = 512
    dyna_ffn_dim: int = 2048
    dyna_num_blocks: int = 6
    dyna_num_heads: int = 8
    dropout: float = 0.0
    mask_limit: float = 0.5
    param_dtype: jnp.dtype = jnp.float32
    dtype: jnp.dtype = jnp.bfloat16
    use_flash_attention: bool = True
    # Logging
    log: bool = False
    entity: str = """"
    project: str = """"
    name: str = ""train_dynamics""
    tags: list[str] = field(default_factory=lambda: [""dynamics""])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = """"
    log_checkpoint_interval: int = 25000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False
    wandb_id: str = """"


args = tyro.cli(Args)


def dynamics_loss_fn(params, state, inputs):
    """"""Compute masked dynamics loss""""""
    inputs[""videos""] = inputs[""videos""].astype(args.dtype) / 255.0
    outputs = state.apply_fn(
        params,
        inputs,
        training=True,
        rngs={""params"": inputs[""rng""], ""dropout"": inputs[""dropout_rng""]},
    )
    mask = outputs[""mask""]
    outputs[""token_logits""] = outputs[""token_logits""].astype(jnp.float32)
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs[""token_logits""], outputs[""video_tokens""]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs[""token_logits""].argmax(-1) == outputs[""video_tokens""]
    acc = (mask * acc).sum() / mask.sum()
    select_probs = jax.nn.softmax(outputs[""token_logits""])
    gt = inputs[""videos""].clip(0, 1).reshape(-1, *inputs[""videos""].shape[2:])
    recon = outputs[""recon""].clip(0, 1).reshape(-1, *outputs[""recon""].shape[2:])
    psnr = pix.psnr(gt, recon).mean() # type: ignore
    ssim = pix.ssim(gt, recon).mean() # type: ignore
    _, index_counts_lam = jnp.unique_counts(
        jnp.ravel(outputs[""lam_indices""]), size=args.num_latent_actions, fill_value=0
    )
    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs[""video_tokens""]), size=args.num_patch_latents, fill_value=0
    )
    codebook_usage_lam = (index_counts_lam != 0).mean()
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    metrics = dict(
        cross_entropy_loss=ce_loss,
        masked_token_accuracy=acc,
        select_logit=outputs[""token_logits""].max(-1).mean(),
        select_p=select_probs.max(-1).mean(),
        entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
        psnr=psnr,
        ssim=ssim,
        codebook_usage_lam=codebook_usage_lam,
        codebook_usage_tokenizer=codebook_usage_tokenizer,
    )
    return ce_loss, (outputs[""recon""], metrics)


@jax.jit
def train_step(state, inputs):
    """"""Update state and compute metrics""""""
    grad_fn = jax.value_and_grad(dynamics_loss_fn, has_aux=True, allow_int=True)
    (loss, (recon, metrics)), grads = grad_fn(state.params, state, inputs)
    state = state.apply_gradients(grads=grads)
    if args.log_gradients:
        metrics[""gradients_std/""] = jax.tree.map(
            lambda x: x.std(), grads[""params""][""dynamics""]
        )
    return state, loss, recon, metrics


if __name__ == ""__main__"":
    jax.distributed.initialize()
    num_devices = jax.device_count()
    if num_devices == 0:
        raise ValueError(""No JAX devices found."")
    print(f""Running on {num_devices} devices."")

    if args.batch_size % num_devices != 0:
        raise ValueError(
            f""Global batch size {args.batch_size} must be divisible by ""
            f""number of devices {num_devices}.""
        )

    per_device_batch_size_for_init = args.batch_size // num_devices

    rng = jax.random.PRNGKey(args.seed)

    # --- Initialize model ---
    genie = Genie(
        # Tokenizer
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        tokenizer_ffn_dim=args.tokenizer_ffn_dim,
        latent_patch_dim=args.latent_patch_dim,
        num_patch_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        tokenizer_num_blocks=args.tokenizer_num_blocks,
        tokenizer_num_heads=args.tokenizer_num_heads,
        # LAM
        lam_dim=args.lam_dim,
        lam_ffn_dim=args.lam_ffn_dim,
        latent_action_dim=args.latent_action_dim,
        num_latent_actions=args.num_latent_actions,
        lam_patch_size=args.lam_patch_size,
        lam_num_blocks=args.lam_num_blocks,
        lam_num_heads=args.lam_num_heads,
        lam_co_train=not args.lam_checkpoint,
        # Dynamics
        dyna_dim=args.dyna_dim,
        dyna_ffn_dim=args.dyna_ffn_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        dropout=args.dropout,
        mask_limit=args.mask_limit,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
    )
    rng, _rng = jax.random.split(rng)
    image_shape = (args.image_height, args.image_width, args.image_channels)
    dummy_inputs = dict(
        videos=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, *image_shape),
            dtype=args.dtype,
        ),
        action=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len), dtype=args.dtype
        ),
        mask_rng=_rng,
    )
    rng, _rng = jax.random.split(rng)
    init_params = genie.init(_rng, dummy_inputs)

    param_counts = count_parameters_by_component(init_params)

    if args.log and jax.process_index() == 0:
        wandb_init_kwargs = {
            ""entity"": args.entity,
            ""project"": args.project,
            ""name"": args.name,
            ""tags"": args.tags,
            ""group"": ""debug"",
            ""config"": args,
        }

        if args.wandb_id:
            wandb_init_kwargs.update(
                {
                    ""id"": args.wandb_id,
                    ""resume"": ""allow"",
                }
            )
        wandb.init(**wandb_init_kwargs)

        wandb.config.update({""model_param_count"": param_counts})

    print(""Parameter counts:"")
    print(param_counts)

    # --- Initialize optimizer ---
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
    train_state = TrainState.create(apply_fn=genie.apply, params=init_params, tx=tx)

    device_mesh_arr = create_device_mesh((num_devices,))
    mesh = Mesh(devices=device_mesh_arr, axis_names=(""data"",))

    replicated_sharding = NamedSharding(mesh, PartitionSpec())
    videos_sharding = NamedSharding(mesh, PartitionSpec(""data"", None, None, None, None))
    train_state = jax.device_put(train_state, replicated_sharding)

    # --- Initialize checkpoint manager ---
    step = 0
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        ""model_state"", ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler
    )
    handler_registry.add(
        ""model_state"", ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler
    )
    handler_registry.add(""dataloader_state"", grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler)  # type: ignore
    handler_registry.add(""dataloader_state"", grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)  # type: ignore

    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        keep_period=args.log_checkpoint_keep_period,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )

    checkpoint_manager = ocp.CheckpointManager(
        args.ckpt_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )

    # --- Create DataLoaderIterator from dataloader ---
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith("".array_record"")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)

    # --- Restore checkpoint ---
    if args.restore_ckpt:
        # Restore full dynamics model
        abstract_train_state = jax.tree_util.tree_map(
            ocp.utils.to_shape_dtype_struct, train_state
        )
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_train_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
        train_state = restored[""model_state""]
        grain_iterator = restored[""dataloader_state""]
        step = checkpoint_manager.latest_step() or 0
        print(f""Restored dataloader and model state from step {step}"")
    else:
        # Restore from pre-trained tokenizer (and LAM)
        train_state = restore_genie_components(
            train_state, replicated_sharding, grain_iterator, dummy_inputs, rng, args
        )

    # --- TRAIN LOOP ---
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_iterator)  # type: ignore
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng, _rng_dropout, _rng_mask = jax.random.split(rng, 4)

            inputs = dict(
                videos=videos,
                rng=_rng,
                dropout_rng=_rng_dropout,
                mask_rng=_rng_mask,
            )
            train_state, loss, recon, metrics = train_step(train_state, inputs)
            metrics[""lr""] = lr_schedule(step)
            print(f""Step {step}, loss: {loss}"")
            step += 1

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    wandb.log(
                        {
                            ""loss"": loss,
                            ""step"": step,
                            **metrics,
                        }
                    )
                if step % args.log_image_interval == 0:
                    gt_seq = inputs[""videos""][0]
                    recon_seq = recon[0].clip(0, 1)
                    comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                    comparison_seq = einops.rearrange(
                        comparison_seq * 255, ""t h w c -> h (t w) c""
                    )
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[args.seq_len - 1])),
                            recon=wandb.Image(np.asarray(recon_seq[args.seq_len - 1])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
            # --- Checkpointing ---
            if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=grain.checkpoint.CheckpointSave(
                            grain_iterator
                        ),
                    ),
                )
                print(f""Saved checkpoint at step {step}"")
            if step >= args.num_steps:
                break

    checkpoint_manager.close()
",python,tab
117,2014841,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py",11200,0,"",python,selection_command
118,2018411,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",0,0,"from typing import Dict, Any

import optax
import jax
import jax.numpy as jnp
import flax.linen as nn
from flax.training.train_state import TrainState
import orbax.checkpoint as ocp

from models.dynamics import DynamicsMaskGIT
from models.lam import LatentActionModel
from models.tokenizer import TokenizerVQVAE

import grain


class Genie(nn.Module):
    """"""Genie model""""""

    # --- Tokenizer ---
    in_dim: int
    tokenizer_dim: int
    tokenizer_ffn_dim: int
    latent_patch_dim: int
    num_patch_latents: int
    patch_size: int
    tokenizer_num_blocks: int
    tokenizer_num_heads: int
    # --- LAM ---
    lam_dim: int
    lam_ffn_dim: int
    latent_action_dim: int
    num_latent_actions: int
    lam_patch_size: int
    lam_num_blocks: int
    lam_num_heads: int
    lam_co_train: bool
    # --- Dynamics ---
    dyna_dim: int
    dyna_ffn_dim: int
    dyna_num_blocks: int
    dyna_num_heads: int
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    dropout: float = 0.0
    mask_limit: float = 0.0

    def setup(self):
        self.tokenizer = TokenizerVQVAE(
            in_dim=self.in_dim,
            model_dim=self.tokenizer_dim,
            ffn_dim=self.tokenizer_ffn_dim,
            latent_dim=self.latent_patch_dim,
            num_latents=self.num_patch_latents,
            patch_size=self.patch_size,
            num_blocks=self.tokenizer_num_blocks,
            num_heads=self.tokenizer_num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        self.lam = LatentActionModel(
            in_dim=self.in_dim,
            model_dim=self.lam_dim,
            ffn_dim=self.lam_ffn_dim,
            latent_dim=self.latent_patch_dim,
            num_latents=self.num_latent_actions,
            patch_size=self.lam_patch_size,
            num_blocks=self.lam_num_blocks,
            num_heads=self.lam_num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        self.dynamics = DynamicsMaskGIT(
            model_dim=self.dyna_dim,
            ffn_dim=self.dyna_ffn_dim,
            num_latents=self.num_patch_latents,
            num_blocks=self.dyna_num_blocks,
            num_heads=self.dyna_num_heads,
            dropout=self.dropout,
            mask_limit=self.mask_limit,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        tokenizer_outputs = self.tokenizer.vq_encode(batch[""videos""], training=False)
        lam_outputs = self.lam.vq_encode(batch[""videos""], training=False)
        latent_actions = jax.lax.cond(
            self.lam_co_train,
            lambda: lam_outputs[""z_q""],
            lambda: jax.lax.stop_gradient(lam_outputs[""z_q""])
        )
        outputs = dict(
            video_tokens=jax.lax.stop_gradient(tokenizer_outputs[""indices""]),
            latent_actions=latent_actions,
        )
        outputs[""mask_rng""] = batch[""mask_rng""]
        dyna_outputs = self.dynamics(outputs, training)
        outputs.update(dyna_outputs)
        mle_indices = jnp.argmax(outputs[""token_logits""], axis=-1)
        outputs[""recon""] = self.tokenizer.decode(
            mle_indices, batch[""videos""].shape[2:4]
        )
        outputs[""lam_indices""] = lam_outputs[""indices""]
        return outputs

    @nn.compact
    def sample(
        self,
        batch: Dict[str, Any],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> Any:
        """"""
        Autoregressively samples up to `seq_len` future frames, following Figure 8 of the paper.

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.

        Note:
        - For interactive or step-wise sampling, detokenization should occur after each action.
        - To maintain consistent tensor shapes across timesteps, all current and future frames are decoded at every step.
        - Temporal causal structure is preserved by 
            a) reapplying the mask before each decoding step.
            b) a temporal causal mask is applied within each ST-transformer block.

        Dimension keys:
            B: batch size  
            T: number of input (conditioning) frames  
            N: patches per frame  
            S: sequence length  
            A: action space  
            D: model latent dimension
        """"""
        # --- Encode videos and actions ---
        tokenizer_out = self.tokenizer.vq_encode(batch[""videos""], training=False)
        token_idxs = tokenizer_out[""indices""] # (B, T, N)
        B, T, N = token_idxs.shape
        pad_shape = (B, seq_len - T, N)
        pad = jnp.zeros(pad_shape, dtype=token_idxs.dtype)
        token_idxs = jnp.concatenate([token_idxs, pad], axis=1) # (B, S, N)
        action_tokens = self.lam.vq.get_codes(batch[""latent_actions""])

        MaskGITLoop = nn.scan(
            MaskGITStep,
            variable_broadcast=""params"",
            split_rngs={""params"": False},
            in_axes=0,
            out_axes=0,
            length=steps,
        )
    
        loop_fn = MaskGITLoop(
            dynamics=self.dynamics,
            tokenizer=self.tokenizer,
            temperature=temperature,
            sample_argmax=sample_argmax,
            steps=steps,
        )

        def generation_step_fn(carry, step_t):
            rng, current_token_idxs = carry
            rng, step_rng = jax.random.split(rng)

            # Mask current and future frames (i.e., t >= step_t)
            mask = jnp.arange(seq_len) >= step_t # (S,)
            mask = jnp.broadcast_to(mask[None, :, None], (B, seq_len, N)) # (B, S, N)
            mask = mask.astype(bool)
            masked_token_idxs = current_token_idxs * ~mask

            # --- Initialize and run MaskGIT loop ---
            init_carry_maskgit = (
                step_rng,
                masked_token_idxs,
                mask,
                action_tokens,
            )
            final_carry_maskgit, _ = loop_fn(init_carry_maskgit, jnp.arange(steps))
            updated_token_idxs = final_carry_maskgit[1]
            new_carry = (rng, updated_token_idxs)
            return new_carry, None

        # --- Run the autoregressive generation using scan ---
        initial_carry = (batch[""rng""], token_idxs)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry, _ = jax.lax.scan(
            generation_step_fn,
            initial_carry,
            timesteps_to_scan
        )
        final_token_idxs = final_carry[1]

        # --- Decode all tokens at once at the end ---
        final_frames = self.tokenizer.decode(
            final_token_idxs,
            video_hw=batch[""videos""].shape[2:4],
        )
        return final_frames

    def vq_encode(self, batch, training) -> Dict[str, Any]:
        # --- Preprocess videos ---
        lam_output = self.lam.vq_encode(batch[""videos""], training=training)
        return lam_output[""indices""]


class MaskGITStep(nn.Module):
    dynamics: nn.Module
    tokenizer: nn.Module
    temperature: float
    sample_argmax: bool
    steps: int

    @nn.compact
    def __call__(self, carry, x):
        rng, token_idxs, mask, action_tokens = carry
        step = x
        N = token_idxs.shape[2]

        # --- Construct + encode video ---
        vid_embed = self.dynamics.patch_embed(token_idxs) # (B, S, N, D)
        mask_token = self.dynamics.mask_token  # (1, 1, 1, D,)
        mask_expanded = mask[..., None] # (B, S, N, 1) 
        vid_embed = jnp.where(mask_expanded, mask_token, vid_embed)

        # --- Predict transition ---
        act_embed = self.dynamics.action_up(action_tokens)
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (self.steps * 2))
        step_temp = self.temperature * (1.0 - unmasked_ratio)
        final_logits = self.dynamics.dynamics(vid_embed) / step_temp

        # --- Sample new tokens for final frame ---
        if self.sample_argmax:
            sampled_token_idxs = jnp.argmax(final_logits, axis=-1)
        else:
            rng, _rng = jax.random.split(rng)
            sampled_token_idxs = jax.random.categorical(_rng, final_logits)
        gather_fn = jax.vmap(jax.vmap(jax.vmap(lambda x, y: x[y])))
        final_token_probs = gather_fn(jax.nn.softmax(final_logits), sampled_token_idxs)
        final_token_probs += ~mask
        # Update masked tokens only
        token_idxs = jnp.where(mask, sampled_token_idxs, token_idxs)

        # --- Update mask ---
        num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
        idx_mask = jnp.arange(final_token_probs.shape[-1]) > num_unmasked_tokens
        sorted_idxs = jnp.argsort(final_token_probs, axis=-1, descending=True)
        mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask))
        new_mask = mask_update_fn(mask, sorted_idxs)

        new_carry = (rng, token_idxs, new_mask, action_tokens)
        return new_carry, None

def restore_genie_components(
    train_state: TrainState,
    sharding: jax.sharding.NamedSharding,
    grain_iterator: grain.DataLoaderIterator,
    inputs: Dict[str, jax.Array],
    rng: jax.Array,
    args,
):
    """"""Restore pre-trained Genie components""""""
    rng, _rng = jax.random.split(rng)

    # dummy values since we only use tx to initialize the dummy train states
    dummy_tx = optax.adamw(
        learning_rate=optax.constant_schedule(args.max_lr),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
    )
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
    

    checkpoint_options = ocp.CheckpointManagerOptions(
        step_format_fixed_length=6,
    )
    tokenizer_checkpoint_manager = ocp.CheckpointManager(
        directory=args.tokenizer_checkpoint,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
    dummy_tokenizer = TokenizerVQVAE(
        in_dim=args.image_channels,
        model_dim=args.tokenizer_dim,
        ffn_dim=args.tokenizer_ffn_dim,
        latent_dim=args.latent_patch_dim,
        num_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        num_blocks=args.tokenizer_num_blocks,
        num_heads=args.tokenizer_num_heads,
        dropout=args.dropout,
        codebook_dropout=args.dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
    )
    tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
    dummy_tokenizer_train_state = TrainState.create(
        apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
    )
    abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        step=tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.StandardRestore(abstract_sharded_tokenizer_state),
            dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
        ),
    )[""model_state""]
    restored_tokenizer_params = restored_tokenizer.params[""params""]
    train_state.params[""params""][""tokenizer""].update(restored_tokenizer_params)
    tokenizer_checkpoint_manager.close()

    if args.lam_checkpoint:
        lam_checkpoint_manager = ocp.CheckpointManager(
            directory=args.lam_checkpoint,
            options=checkpoint_options,
            handler_registry=handler_registry,
        )
        dummy_lam = LatentActionModel(
            in_dim=args.image_channels,
            model_dim=args.lam_dim,
            ffn_dim=args.lam_ffn_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_latent_actions,
            patch_size=args.lam_patch_size,
            num_blocks=args.lam_num_blocks,
            num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
        )
        lam_init_params = dummy_lam.init(_rng, inputs)
        dummy_lam_train_state = TrainState.create(
            apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
        )
        abstract_sharded_lam_state = _create_abstract_sharded_pytree(
            dummy_lam_train_state, sharding
        )
        restored_lam = lam_checkpoint_manager.restore(
            step=lam_checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_sharded_lam_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )[""model_state""]
        restored_lam_params = restored_lam.params[""params""]
        # Genie does not initialize all LAM modules, thus we omit those extra modules during restoration
        # (f.srambical) FIXME: Currently, this is a small HBM memory crunch since the LAM's decoder is loaded into HBM and immediately dicarded.
        # A workaround would be to restore to host memory first, and only move the weights to HBM after pruning the decoder
        restored_lam_params = {
            k: v
            for k, v in restored_lam_params.items()
            if k in train_state.params[""params""][""lam""]
        }
        train_state.params[""params""][""lam""].update(restored_lam_params)
        lam_checkpoint_manager.close()

    return train_state

def _create_abstract_sharded_pytree(pytree_template, sharding_spec):
    """"""Replaces arrays in a pytree with ShapeDtypeStructs having the given sharding.""""""

    def map_fn(leaf_template):
        if hasattr(leaf_template, ""shape"") and hasattr(leaf_template, ""dtype""):
            return jax.ShapeDtypeStruct(
                leaf_template.shape, leaf_template.dtype, sharding=sharding_spec
            )
        return leaf_template

    return jax.tree_util.tree_map(map_fn, pytree_template)",python,tab
119,2018413,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",11482,0,"",python,selection_command
120,2027221,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
121,2027423,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
:",,terminal_output
122,2027708,"TERMINAL",0,0,"
 ESCESCOOBB
    return self._impl.restore(
:",,terminal_output
123,2027787,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
124,2031240,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
:",,terminal_output
125,2031303,"TERMINAL",0,0,"
 ESCESCOOBB
Traceback (most recent call last):
:",,terminal_output
126,2031392,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
:",,terminal_output
127,2032670,"TERMINAL",0,0,"
 ESCESCOOBB
    train_state = restore_genie_components(
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon :
 ESCESCOOBB
ents
:",,terminal_output
128,2033235,"TERMINAL",0,0,"
 ESCESCOOBB
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/checkpoint_manager.py"", line 1608, in restore
:
 ESCESCOOBB
    return self._handler_impl.restore(directory, args=args)
:",,terminal_output
129,2033320,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
:",,terminal_output
130,2033380,"TERMINAL",0,0,"
 ESCESCOOBB
    return self._handler_impl.restore(directory, args=args)
:",,terminal_output
131,2033833,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
132,2033972,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:",,terminal_output
133,2034042,"TERMINAL",0,0,"
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:",,terminal_output
134,2034569,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:",,terminal_output
135,2034622,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
136,2035602,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:",,terminal_output
137,2035756,"TERMINAL",0,0,"
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:",,terminal_output
138,2035884,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:",,terminal_output
139,2037678,"TERMINAL",0,0,"
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:",,terminal_output
140,2037853,"TERMINAL",0,0,"
 ESCESCOOBB
    restored = self._restore(directory, args=ckpt_args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
141,2038380,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:
 ESCESCOOBB
    restored = self._restore(directory, args=ckpt_args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:
 ESCESCOOBB
    restored = self._restore(directory, args=ckpt_args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
142,2038458,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:",,terminal_output
143,2038764,"TERMINAL",0,0,"
 ESCESCOOBB
    restored = self._restore(directory, args=ckpt_args)
:",,terminal_output
144,2038832,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
145,2038946,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:
 ESCESCOOBB
    return self._handler.restore(directory, args=args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
:
 ESCESCOOBB
    return self._handler.restore(directory, args=args)
:",,terminal_output
146,2042185,"TERMINAL",0,0,"
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
Traceback (most recent call last):
    raise ValueError(
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=f loat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=de vice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.A rray', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_ deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array' , skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deseri alize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_des erialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=Fals e, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=F alse, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_s hape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deseri :",,terminal_output
147,2051695,"TERMINAL",0,0,"
 ESCESCOOAA
M    return self._handler_impl.restore(directory, args=args)

:
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:",,terminal_output
148,2051818,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:",,terminal_output
149,2052220,"TERMINAL",0,0,"
 ESCESCOOAA
M    train_state = restore_genie_components(

:",,terminal_output
150,2052416,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Mckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return self._impl.restore(

:",,terminal_output
151,2052474,"TERMINAL",0,0,"
 ESCESCOOAA
Mckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore

:",,terminal_output
152,2052881,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored[item_name] = handler.restore(

:
 ESCESCOOAA
Mckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return self._handler.restore(directory, args=args)

:",,terminal_output
153,2052967,"TERMINAL",0,0,"
 ESCESCOOAA
Mckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore

:",,terminal_output
154,2053388,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored = self._restore(directory, args=ckpt_args)

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:",,terminal_output
155,2053507,"TERMINAL",0,0,"
 ESCESCOOAA
M    return super().restore(directory, *args, **kwargs)

:",,terminal_output
156,2054871,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",11481,0,"",python,selection_mouse
157,2054873,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",11480,0,"",python,selection_command
158,2069374,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
159,2070055,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1649,0,"",shellscript,selection_command
160,2070301,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1600,0,"",shellscript,selection_command
161,2070412,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1578,0,"",shellscript,selection_command
162,2070415,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1552,0,"",shellscript,selection_command
163,2070418,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1479,0,"",shellscript,selection_command
164,2070433,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1438,0,"",shellscript,selection_command
165,2070591,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1401,0,"",shellscript,selection_command
166,2070594,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1389,0,"",shellscript,selection_command
167,2070596,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1357,0,"",shellscript,selection_command
168,2070598,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1335,0,"",shellscript,selection_command
169,2070602,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1317,0,"",shellscript,selection_command
170,2070638,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1291,0,"",shellscript,selection_command
171,2070691,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1265,0,"",shellscript,selection_command
172,2070777,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1240,0,"",shellscript,selection_command
173,2070778,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1215,0,"",shellscript,selection_command
174,2070779,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1184,0,"",shellscript,selection_command
175,2070804,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1153,0,"",shellscript,selection_command
176,2070838,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1131,0,"",shellscript,selection_command
177,2070937,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1098,0,"",shellscript,selection_command
178,2070938,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1080,0,"",shellscript,selection_command
179,2070961,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1048,0,"",shellscript,selection_command
180,2071068,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1047,0,"",shellscript,selection_command
181,2071071,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1030,0,"",shellscript,selection_command
182,2071170,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1029,0,"",shellscript,selection_command
183,2071247,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",882,0,"",shellscript,selection_command
184,2071465,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",881,0,"",shellscript,selection_command
185,2071511,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",882,0,"",shellscript,selection_command
186,2071654,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1029,0,"",shellscript,selection_command
187,2071838,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1030,0,"",shellscript,selection_command
188,2072026,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1034,0,"",shellscript,selection_command
189,2072208,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1029,0,"",shellscript,selection_command
190,2072382,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",886,0,"",shellscript,selection_command
191,2072516,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",900,0,"",shellscript,selection_command
192,2072689,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",902,0,"",shellscript,selection_command
193,2075429,"TERMINAL",0,0,"
>",,terminal_output
194,2077230,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",972,0,"",shellscript,selection_mouse
195,2079997,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",900,0,"",shellscript,selection_command
196,2082189,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,0,"",shellscript,selection_command
197,2082308,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,1,"/",shellscript,selection_command
198,2085092,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,127,"/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",shellscript,selection_command
199,2086099,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1027,0,"",shellscript,selection_command
200,2088409,"TERMINAL",0,0,"ls /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",,terminal_command
201,2088460,"TERMINAL",0,0,"",,terminal_output
202,2088528,"TERMINAL",0,0,"020000  040000  060000  080000  100000  120000  140000  144000  145000  146000
",,terminal_output
203,2349376,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3352437.log",,terminal_command
204,2349420,"TERMINAL",0,0,"",,terminal_output
205,2349494,"TERMINAL",0,0,"=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=2048 \
    --tokenizer_num_blocks=4 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2571288
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0519
SLURM_JOB_START_TIME=1752692957
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752779357
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3352437
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0519
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0519-0520]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1993.localdomain
SLURM_JOB_ID=3352437
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3352437.log",,terminal_output
206,2350819,"TERMINAL",0,0,"
SLURM_NTASKS_PER_NODE=4
SLURM_JOB_GID=502226
SLURM_JOB_NODELIST=hkn[0519-0520]
GpuFreq=control_disabled
GpuFreq=control_disabled
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
:",,terminal_output
207,2351808,"TERMINAL",0,0,"
Running on 8 devices.
2025-07-16 21:10:00.136052: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.136050: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.147869: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.148033: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.269210: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.272749: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.287019: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:00.287044: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:26.516484: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:27.061049: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:27.166309: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:27.593704: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:29.445859: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:29.880469: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:30.345131: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:30.928066: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:31.827307: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:32.318324: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:32.468961: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:32.744868: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
2025-07-16 21:10:35.335235: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:35.404405: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:35.955390: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-16 21:10:36.822663: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
wandb: Currently logged in as: emergenz (p-doom) to https://api.wandb.ai. Use `wandb login --relogin` to forc e relogin
Counting all components: ['tokenizer', 'lam', 'dynamics']
:",,terminal_output
208,2352299,"TERMINAL",0,0,"
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
wandb: Tracking run with wandb version 0.19.11
wandb: Run data is saved locally in /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/wandb/run-202507 16_211038-jktj3wow
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run dynamics-lr-1e-4-larger-ffn
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/jktj3wow
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
:",,terminal_output
209,2352994,"TERMINAL",0,0,"
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
:",,terminal_output
210,2353414,"TERMINAL",0,0,"
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon :",,terminal_output
211,2353996,"TERMINAL",0,0,"
ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
:",,terminal_output
212,2355096,"TERMINAL",0,0,"
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_componM    train_state = restore_genie_components(
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
MTraceback (most recent call last):
Ments
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_componM    train_state = restore_genie_components(
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
MTraceback (most recent call last):
Ments
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_componM    train_state = restore_genie_components(
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
MTraceback (most recent call last):
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 60000
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 20000
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 100000
Mling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 144000
Mling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 146000
Mling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 120000
Mling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 40000
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 140000
Mling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 80000
Mling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 145000
Mling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 146000

:",,terminal_output
213,2355907,"TERMINAL",0,0,"
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 100000
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 140000
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 60000
Mling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 145000
Mling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 80000
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 60000
Mling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 144000
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 20000
MWARNING:absl:Missing metrics for step 20000
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 20000
Mling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 80000
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 100000
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 100000
Mling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 80000
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 140000
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:",,terminal_output
214,2356530,"TERMINAL",0,0,"
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
:",,terminal_output
215,2373222,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
216,2373399,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:",,terminal_output
217,2373581,"TERMINAL",0,0,"
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:",,terminal_output
218,2373762,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
219,2374064,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:
 ESCESCOOBB
    restored = self._restore(directory, args=ckpt_args)
:",,terminal_output
220,2374151,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
221,2375016,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:",,terminal_output
222,2375082,"TERMINAL",0,0,"
 ESCESCOOBB
    return self._handler.restore(directory, args=args)
:",,terminal_output
223,2375191,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
224,2375336,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
:",,terminal_output
225,2375406,"TERMINAL",0,0,"
 ESCESCOOBB
    restored[item_name] = handler.restore(
:",,terminal_output
226,2375518,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
:",,terminal_output
227,2375598,"TERMINAL",0,0,"
 ESCESCOOBB
    return self._impl.restore(
:",,terminal_output
228,2376100,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
:
 ESCESCOOBB
Traceback (most recent call last):
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
:
 ESCESCOOBB
    train_state = restore_genie_components(
:",,terminal_output
229,2376201,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon :",,terminal_output
230,2376649,"TERMINAL",0,0,"
 ESCESCOOBB
ents
:
 ESCESCOOBB
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
231,2376701,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/checkpoint_manager.py"", line 1608, in restore
:",,terminal_output
232,2389444,"TERMINAL",0,0,"
 ESCESCOOBB
    return self._handler_impl.restore(directory, args=args)
:",,terminal_output
233,2389503,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
234,2389686,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
:",,terminal_output
235,2390539,"TERMINAL",0,0,"
 ESCESCOOBB
    return self._handler_impl.restore(directory, args=args)
:",,terminal_output
236,2390661,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:",,terminal_output
237,2390790,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
238,2391152,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:",,terminal_output
239,2391312,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
240,2391949,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:
 ESCESCOOBB
    restored = self._checkpointer.restore(restore_directory, args=args)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
:
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
241,2392033,"TERMINAL",0,0,"
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:",,terminal_output
242,2392673,"TERMINAL",0,0,"
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
:
 ESCESCOOBB
    return super().restore(directory, *args, **kwargs)
:",,terminal_output
243,2412982,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:",,terminal_output
244,2413049,"TERMINAL",0,0,"
 ESCESCOOAA
M    restored = self._checkpointer.restore(restore_directory, args=args)

:",,terminal_output
245,2413115,"TERMINAL",0,0,"
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:",,terminal_output
246,2413267,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Ments

:",,terminal_output
247,2413798,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:",,terminal_output
248,2413867,"TERMINAL",0,0,"
 ESCESCOOAA
M    train_state = restore_genie_components(

:",,terminal_output
249,2467211,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:",,terminal_output
250,2467670,"TERMINAL",0,0,"
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 60000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 20000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 100000

:",,terminal_output
251,2468112,"TERMINAL",0,0,"
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 144000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 146000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:",,terminal_output
252,2468176,"TERMINAL",0,0,"
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 120000

:",,terminal_output
253,2472885,"TERMINAL",0,0,"
 ESCESCOOBB
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
 ESCESCOOBB
ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
:
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:",,terminal_output
254,2473162,"TERMINAL",0,0,"
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
Traceback (most recent call last):
    raise ValueError(
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=f loat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=de vice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.A rray', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_ :",,terminal_output
255,2474676,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Mckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return self._impl.restore(

:
 ESCESCOOAA
Mckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored[item_name] = handler.restore(

:
 ESCESCOOAA
Mckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore

:",,terminal_output
256,2475018,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return self._handler.restore(directory, args=args)

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored = self._restore(directory, args=ckpt_args)

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore

:",,terminal_output
257,2475085,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:",,terminal_output
258,2475471,"TERMINAL",0,0,"
 ESCESCOOAA
M    return super().restore(directory, *args, **kwargs)

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored = self._checkpointer.restore(restore_directory, args=args)

:
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:",,terminal_output
259,2475821,"TERMINAL",0,0,"
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Mckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return self._impl.restore(

:
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:",,terminal_output
260,2476193,"TERMINAL",0,0,"
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Mckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:",,terminal_output
261,2476550,"TERMINAL",0,0,"
 ESCESCOOAA
M    restored[item_name] = handler.restore(

:
 ESCESCOOAA
Mckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return self._handler.restore(directory, args=args)

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored = self._restore(directory, args=ckpt_args)

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:",,terminal_output
262,2476901,"TERMINAL",0,0,"
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore

:",,terminal_output
263,2476983,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    return super().restore(directory, *args, **kwargs)

:
 ESCESCOOAA
Mckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored = self._checkpointer.restore(restore_directory, args=args)

:",,terminal_output
264,2477365,"TERMINAL",0,0,"
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:",,terminal_output
265,2477465,"TERMINAL",0,0,"
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Ments

:",,terminal_output
266,2477647,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:",,terminal_output
267,2477820,"TERMINAL",0,0,"
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:
 ESCESCOOAA
M    train_state = restore_genie_components(

:",,terminal_output
268,2478791,"TERMINAL",0,0,"
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
MTraceback (most recent call last):
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 60000
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 20000
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 100000
Mling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 144000
Mling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 146000
Mling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 120000
Mling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 40000
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 140000
Mling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 80000
Mling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 145000
Mling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 146000
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 100000
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 140000
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 60000
Mling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaMWARNING:absl:Missing metrics for step 145000

:",,terminal_output
269,2480349,"TERMINAL",0,0,"
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
:",,terminal_output
270,2481677,"TERMINAL",0,0,"
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:",,terminal_output
271,2482194,"TERMINAL",0,0,"
MTraceback (most recent call last):

:
Mling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.

:
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
MWARNING:absl:Missing metrics for step 60000

:
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.

:
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
MWARNING:absl:Missing metrics for step 20000

:
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.

:
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:",,terminal_output
272,2495566,"TERMINAL",0,0,"
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :",,terminal_output
273,2501288,"TERMINAL",0,0,"
ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
Traceback (most recent call last):
    raise ValueError(
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=f loat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=de vice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.A rray', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_ deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array' , skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deseri alize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_des erialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=Fals e, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=F alse, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_s hape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deseri alize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=Fals e, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=Fal se, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_sh ape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, writ e_shape=(1, 64)), 'k  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line  318, in <module>
    train_state = restore_genie_components(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=f loat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=de vice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 51 2), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), mem ory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dty pe=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kin d=device)), 'kernel': ShapeDtypeStruct(shape=(2048, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': Diff(lhs=N one, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape= (64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))},  'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,) ), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNor m_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale' : ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bi as': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMet adataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key ': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel' : ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias ': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMeta dataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': Value MetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'k    return self._impl.r estore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
ernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value' : {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': :",,terminal_output
274,2511477,"TERMINAL",0,0,"
 ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAtt ention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1 , 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64)) }, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'ke rnel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query' : {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bia s': ValueMetadataEntry(value_type='jax.Arra    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
ernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value' : {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAtt ention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1 , 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64)) }, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'ke rnel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query' : {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bia s': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueM etadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff( lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_s hape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 51 2))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape= (64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'Lay erNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 's cale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2':  {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': Val ueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0':  {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'ke rnel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': { 'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': Valu eMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetad ataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMe tadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntr y(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key':  {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias':  ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetada taEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMe tadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntr y(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataE ntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value _type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'De nse_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kern el': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0 :",,terminal_output
275,2610920,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1173,0,"",shellscript,selection_mouse
276,2611469,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1176,0,"",shellscript,selection_command
277,2611594,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
278,2612835,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,4,"",shellscript,content
279,2613599,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"5",shellscript,content
280,2613600,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1178,0,"",shellscript,selection_keyboard
281,2613711,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1178,0,"1",shellscript,content
282,2613712,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1179,0,"",shellscript,selection_keyboard
283,2613796,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1179,0,"2",shellscript,content
284,2613797,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,0,"",shellscript,selection_keyboard
285,2613980,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1179,0,"",shellscript,selection_command
286,2614972,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1209,0,"",shellscript,selection_command
287,2615047,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,0,"",shellscript,selection_command
288,2615853,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"8",shellscript,content
289,2630426,"train_tokenizer.py",0,0,"from dataclasses import dataclass, field
import os

import einops
from flax.training import orbax_utils
from flax.training.train_state import TrainState
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh
import optax
import orbax.checkpoint as ocp
import numpy as np
import dm_pix as pix
import jax
import jax.numpy as jnp
import tyro
import wandb
import grain

from models.tokenizer import TokenizerVQVAE
from utils.dataloader import get_dataloader
from utils.lr_utils import get_lr_schedule
from utils.parameter_utils import count_parameters_by_component


@dataclass
class Args:
    # Experiment
    num_steps: int = 300_000
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = """"
    save_ckpt: bool = False
    restore_ckpt: bool = False
    # Optimization
    vq_beta: float = 0.25
    batch_size: int = 48
    init_lr: float = 0.0
    max_lr: float = 3e-4
    decay_end: float = 0.0
    wsd_decay_steps: int = 20000 # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    lr_schedule: str = ""wsd"" # supported options: wsd, cos 
    warmup_steps: int = 10000
    # Tokenizer
    model_dim: int = 512
    ffn_dim: int = 2048
    latent_dim: int = 32
    num_latents: int = 1024
    patch_size: int = 4
    num_blocks: int = 4
    num_heads: int = 8
    dropout: float = 0.0
    codebook_dropout: float = 0.01
    param_dtype: jnp.dtype = jnp.float32
    dtype: jnp.dtype = jnp.bfloat16
    # Logging
    log: bool = False
    entity: str = """"
    project: str = """"
    name: str = ""train_tokenizer""
    tags: list[str] = field(default_factory=lambda: [""tokenizer""])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = """"
    log_checkpoint_interval: int = 10000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False
    wandb_id: str = """"
    use_flash_attention: bool = True


args = tyro.cli(Args)


def tokenizer_loss_fn(params, state, inputs):
    # --- Compute loss ---
    # FIXME (f.srambical): Can we even do native int8 training without casting the video at all?
    # FIXME (f.srambical): If the tokenizer is the reason for the dynamics model being memory-bound,
    # should we at least train the tokenizer natively in int8?
    inputs[""videos""] = inputs[""videos""].astype(args.dtype) / 255.0
    outputs = state.apply_fn(
        params,
        inputs,
        training=True,
        rngs={""params"": inputs[""rng""], ""dropout"": inputs[""dropout_rng""]},
    )
    mse = jnp.square(inputs[""videos""] - outputs[""recon""]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs[""emb""]) - outputs[""z""]).mean()
    commitment_loss = jnp.square(
        outputs[""emb""] - jax.lax.stop_gradient(outputs[""z""])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss

    # --- Compute validation metrics ---
    gt = inputs[""videos""].clip(0, 1).reshape(-1, *inputs[""videos""].shape[2:])
    recon = outputs[""recon""].clip(0, 1).reshape(-1, *outputs[""recon""].shape[2:])
    psnr = pix.psnr(gt, recon).mean()  # type: ignore
    ssim = pix.ssim(gt, recon).mean()  # type: ignore
    _, index_counts = jnp.unique_counts(
        jnp.ravel(outputs[""indices""]), size=args.num_latents, fill_value=0
    )
    codebook_usage = (index_counts != 0).mean()
    metrics = dict(
        loss=loss,
        mse=mse,
        q_loss=q_loss,
        commitment_loss=commitment_loss,
        psnr=psnr,
        ssim=ssim,
        codebook_usage=codebook_usage,
    )
    return loss, (outputs[""recon""], metrics)


@jax.jit
def train_step(state, inputs):
    grad_fn = jax.value_and_grad(tokenizer_loss_fn, has_aux=True, allow_int=True)
    (loss, (recon, metrics)), grads = grad_fn(state.params, state, inputs)
    state = state.apply_gradients(grads=grads)
    if args.log_gradients:
        metrics[""encoder_gradients_std/""] = jax.tree.map(
            lambda x: x.std(), grads[""params""][""encoder""]
        )
        metrics[""vq_gradients_std/""] = jax.tree.map(
            lambda x: x.std(), grads[""params""][""vq""]
        )
        metrics[""decoder_gradients_std/""] = jax.tree.map(
            lambda x: x.std(), grads[""params""][""decoder""]
        )
    return state, loss, recon, metrics


if __name__ == ""__main__"":
    jax.distributed.initialize()
    num_devices = jax.device_count()
    if num_devices == 0:
        raise ValueError(""No JAX devices found."")
    print(f""Running on {num_devices} devices."")

    if args.batch_size % num_devices != 0:
        raise ValueError(
            f""Global batch size {args.batch_size} must be divisible by ""
            f""number of devices {num_devices}.""
        )

    per_device_batch_size_for_init = args.batch_size // num_devices

    rng = jax.random.PRNGKey(args.seed)

    # --- Initialize model ---
    tokenizer = TokenizerVQVAE(
        in_dim=args.image_channels,
        model_dim=args.model_dim,
        ffn_dim=args.ffn_dim,
        latent_dim=args.latent_dim,
        num_latents=args.num_latents,
        patch_size=args.patch_size,
        num_blocks=args.num_blocks,
        num_heads=args.num_heads,
        dropout=args.dropout,
        codebook_dropout=args.codebook_dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
    )
    rng, _rng = jax.random.split(rng)
    image_shape = (args.image_height, args.image_width, args.image_channels)
    inputs = dict(
        videos=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, *image_shape),
            dtype=args.dtype,
        ),
    )
    init_params = tokenizer.init(_rng, inputs)

    param_counts = count_parameters_by_component(init_params)

    if args.log and jax.process_index() == 0:
        wandb_init_kwargs = {
            ""entity"": args.entity,
            ""project"": args.project,
            ""name"": args.name,
            ""tags"": args.tags,
            ""group"": ""debug"",
            ""config"": args,
        }

        if args.wandb_id:
            wandb_init_kwargs.update(
                {
                    ""id"": args.wandb_id,
                    ""resume"": ""allow"",
                }
            )
        wandb.init(**wandb_init_kwargs)

        wandb.config.update({""model_param_count"": param_counts})

    print(""Parameter counts:"")
    print(param_counts)

    # --- Initialize optimizer ---
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
    train_state = TrainState.create(apply_fn=tokenizer.apply, params=init_params, tx=tx)

    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
    device_mesh_arr = create_device_mesh((num_devices,))
    mesh = Mesh(devices=device_mesh_arr, axis_names=(""data"",))

    replicated_sharding = NamedSharding(mesh, PartitionSpec())
    videos_sharding = NamedSharding(mesh, PartitionSpec(""data"", None, None, None, None))
    train_state = jax.device_put(train_state, replicated_sharding)

    # --- Initialize checkpoint manager ---
    step = 0
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        ""model_state"", ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler
    )
    handler_registry.add(
        ""model_state"", ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler
    )
    handler_registry.add(""dataloader_state"", grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler)  # type: ignore
    handler_registry.add(""dataloader_state"", grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)  # type: ignore

    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        keep_period=args.log_checkpoint_keep_period,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )

    checkpoint_manager = ocp.CheckpointManager(
        args.ckpt_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )

    # --- Create DataLoaderIterator from dataloader ---
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith("".array_record"")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)

    # --- Restore checkpoint ---
    if args.restore_ckpt:
        abstract_train_state = jax.tree_util.tree_map(
            ocp.utils.to_shape_dtype_struct, train_state
        )
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_train_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
        train_state = restored[""model_state""]
        grain_iterator = restored[""dataloader_state""]
        step = checkpoint_manager.latest_step() or 0
        print(f""Restored dataloader and model state from step {step}"")

    # --- TRAIN LOOP ---
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_iterator)  # type: ignore
    print(f""Starting training from step {step}..."")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng, _rng_dropout = jax.random.split(rng, 3)

            inputs = dict(videos=videos, rng=_rng, dropout_rng=_rng_dropout)
            train_state, loss, recon, metrics = train_step(train_state, inputs)
            metrics[""lr""] = lr_schedule(step)
            print(f""Step {step}, loss: {loss}"")
            step += 1

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    wandb.log(
                        {
                            ""loss"": loss,
                            ""step"": step,
                            **metrics,
                        }
                    )
                if step % args.log_image_interval == 0:
                    gt_seq = inputs[""videos""][0]
                    recon_seq = recon[0].clip(0, 1)
                    comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                    comparison_seq = einops.rearrange(
                        comparison_seq * 255, ""t h w c -> h (t w) c""
                    )
                    # NOTE: Process-dependent control flow deliberately happens
                    # after indexing operation since it must not contain code
                    # sections that lead to cross-accelerator communication.
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[0])),
                            recon=wandb.Image(np.asarray(recon_seq[0])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
            # --- Checkpointing ---
            if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=grain.checkpoint.CheckpointSave(
                            grain_iterator
                        ),
                    ),
                )
                print(f""Saved checkpoint at step {step}"")
            if step >= args.num_steps:
                break

    checkpoint_manager.close()
",python,tab
290,2640052,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
291,3175801,"TERMINAL",0,0,"
>",,terminal_output
292,3177210,"TERMINAL",0,0,"idle",,terminal_command
293,3177242,"TERMINAL",0,0,"Partition dev_cpuonly             :      8 nodes idle
Partition cpuonly                 :      1 nodes idle
Partition dev_accelerated         :      2 nodes idle
Partition accelerated             :     61 nodes idle
Partition dev_accelerated-h100    :      0 nodes idle
Partition accelerated-h100        :      0 nodes idle
Partition large                   :      8 nodes idle
",,terminal_output
294,3186753,"TERMINAL",0,0,"sbatch experiments/train_dynamics_lr_1e-4_larger_ffn.sh",,terminal_command
295,3186794,"TERMINAL",0,0,"Submitted batch job 3353271
",,terminal_output
296,3188920,"TERMINAL",0,0,"queue",,terminal_command
297,3189026,"TERMINAL",0,0,"",,terminal_output
298,3189096,"TERMINAL",0,0,"(BEvery 1.0s: squeue --mehkn1990.localdomain: Thu Jul 17 12:33:54 2025JOBID PARTITION     NAME     USER ST	TIME  NODES NODELIST(REASON)3353271 accelerat train_dy tum_dbd0 PD	0:00	  2 (None)",,terminal_output
299,3190084,"TERMINAL",0,0,"5",,terminal_output
300,3191131,"TERMINAL",0,0,"6",,terminal_output
301,3192173,"TERMINAL",0,0,"7",,terminal_output
302,3193220,"TERMINAL",0,0,"8",,terminal_output
303,3194282,"TERMINAL",0,0,"9",,terminal_output
304,3195315,"TERMINAL",0,0,"4:00",,terminal_output
305,3196380,"TERMINAL",0,0,"1",,terminal_output
306,3197410,"TERMINAL",0,0,"2",,terminal_output
307,3198447,"TERMINAL",0,0,"3",,terminal_output
308,3199497,"TERMINAL",0,0,"4",,terminal_output
309,3200544,"TERMINAL",0,0,"5",,terminal_output
310,3201595,"TERMINAL",0,0,"6",,terminal_output
311,3202650,"TERMINAL",0,0,"8",,terminal_output
312,3203703,"TERMINAL",0,0,"9",,terminal_output
313,3204763,"TERMINAL",0,0,"10",,terminal_output
314,3205812,"TERMINAL",0,0,"1",,terminal_output
315,3206862,"TERMINAL",0,0,"2",,terminal_output
316,3207908,"TERMINAL",0,0,"3",,terminal_output
317,3208948,"TERMINAL",0,0,"4",,terminal_output
318,3210000,"TERMINAL",0,0,"5",,terminal_output
319,3211082,"TERMINAL",0,0,"6",,terminal_output
320,3212113,"TERMINAL",0,0,"7",,terminal_output
321,3213151,"TERMINAL",0,0,"8",,terminal_output
322,3214207,"TERMINAL",0,0,"9",,terminal_output
323,3215248,"TERMINAL",0,0,"20",,terminal_output
324,3216308,"TERMINAL",0,0,"1",,terminal_output
325,3217394,"TERMINAL",0,0,"2",,terminal_output
326,3217942,"TERMINAL",0,0,"
>",,terminal_output
327,3222390,"TERMINAL",0,0,"salloc --nodes=2 --ntasks-per-node=4 --gres=gpu:4 --cpus-per-task=8 --time=9:00:00 --partition=accelerated --account=hk-project-p0023960",,terminal_command
328,3222466,"TERMINAL",0,0,"salloc: Pending job allocation 3353272
salloc: job 3353272 queued and waiting for resources
",,terminal_output
329,3233618,"TERMINAL",0,0,"salloc: job 3353272 has been allocated resources
salloc: Granted job allocation 3353272
",,terminal_output
330,3233726,"TERMINAL",0,0,"salloc: Waiting for resource configuration
",,terminal_output
331,3260790,"TERMINAL",0,0,"salloc: Nodes hkn[0819,0821] are ready for job
",,terminal_output
332,3261574,"TERMINAL",0,0,"[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
333,3395382,"TERMINAL",0,0,"q",,terminal_output
334,3395489,"TERMINAL",0,0,"u",,terminal_output
335,3395592,"TERMINAL",0,0,"e",,terminal_output
336,3395659,"TERMINAL",0,0,"u",,terminal_output
337,3395746,"TERMINAL",0,0,"e",,terminal_output
338,3395889,"TERMINAL",0,0,"

(BEvery 1.0s: squeue --mehkn0819.localdomain: Thu Jul 17 12:37:21 2025JOBID PARTITION     NAME     USER ST	TIME  NODES NODELIST(REASON)3353271 accelerat train_dy tum_dbd0  R	2:42	  2 hkn[0516,0518]3353272 accelerat interact tum_dbd0  R	2:42	  2 hkn[0819,0821]",,terminal_output
339,3396902,"TERMINAL",0,0,"233",,terminal_output
340,3397921,"TERMINAL",0,0,"344",,terminal_output
341,3398948,"TERMINAL",0,0,"455",,terminal_output
342,3399969,"TERMINAL",0,0,"566",,terminal_output
343,3400969,"TERMINAL",0,0,"677",,terminal_output
344,3401985,"TERMINAL",0,0,"788",,terminal_output
345,3403001,"TERMINAL",0,0,"899",,terminal_output
346,3404020,"TERMINAL",0,0,"95050",,terminal_output
347,3405044,"TERMINAL",0,0,"3011",,terminal_output
348,3406067,"TERMINAL",0,0,"122",,terminal_output
349,3407071,"TERMINAL",0,0,"233",,terminal_output
350,3408084,"TERMINAL",0,0,"344",,terminal_output
351,3409106,"TERMINAL",0,0,"455",,terminal_output
352,3410122,"TERMINAL",0,0,"5CG6",,terminal_output
353,3411137,"TERMINAL",0,0,"67",,terminal_output
354,3412157,"TERMINAL",0,0,"78",,terminal_output
355,3413175,"TERMINAL",0,0,"89",,terminal_output
356,3414186,"TERMINAL",0,0,"93:00",,terminal_output
357,3415209,"TERMINAL",0,0,"401",,terminal_output
358,3416221,"TERMINAL",0,0,"12",,terminal_output
359,3417238,"TERMINAL",0,0,"23",,terminal_output
360,3418167,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
361,3419047,"TERMINAL",0,0,"
(reverse-i-search)`': ",,terminal_output
362,3419292,"TERMINAL",0,0,"l': idlee': idle",,terminal_output
363,3419378,"TERMINAL",0,0,"s': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4_3345470.log 
",,terminal_output
364,3419492,"TERMINAL",0,0,"s': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tokenize",,terminal_output
365,3419865,"TERMINAL",0,0,"
[tum_dbd0378@hkn0819 jafar]$ less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tokeniz

",,terminal_output
366,3420525,"TERMINAL",0,0,"",,terminal_output
367,3420973,"TERMINAL",0,0,"",,terminal_output
368,3421133,"TERMINAL",0,0,"",,terminal_output
369,3421296,"TERMINAL",0,0,"",,terminal_output
370,3421815,"TERMINAL",0,0,"",,terminal_output
371,3421996,"TERMINAL",0,0,"",,terminal_output
372,3422199,"TERMINAL",0,0,"3345470.log ",,terminal_output
373,3422332,"TERMINAL",0,0,"",,terminal_output
374,3422952,"TERMINAL",0,0,"",,terminal_output
375,3423370,"TERMINAL",0,0,"",,terminal_output
376,3423996,"TERMINAL",0,0,"",,terminal_output
377,3429511,"TERMINAL",0,0,"",,terminal_output
378,3430796,"TERMINAL",0,0,"",,terminal_output
379,3430907,"TERMINAL",0,0,"",,terminal_output
380,3431001,"TERMINAL",0,0,"",,terminal_output
381,3431507,"TERMINAL",0,0,"",,terminal_output
382,3431940,"TERMINAL",0,0,"",,terminal_output
383,3432017,"TERMINAL",0,0,"",,terminal_output
384,3432264,"TERMINAL",0,0,"",,terminal_output
385,3432541,"TERMINAL",0,0,"


",,terminal_output
386,3432787,"TERMINAL",0,0,"",,terminal_output
387,3433062,"TERMINAL",0,0,"",,terminal_output
388,3433312,"TERMINAL",0,0,"",,terminal_output
389,3433608,"TERMINAL",0,0,"
big-runs/
dynamics-cotraining-modelsize-scaling/
tokenizer-lr-scaling/
train_dynamics_lr_1e-4_larger_ffn_3352437.log
train_dynamics_lr_1e-4_larger_ffn_3353271.log
",,terminal_output
390,3433697,"TERMINAL",0,0,"train_tokenizer_batch_size_scaling_linear_lr_scaling_16_nodes_3322407.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_16_nodes_3325607.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_2_nodes_3322404.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_2_nodes_3325603.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_4_nodes_3322405.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_4_nodes_3325605.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_8_nodes_3322406.log
train_tokenizer_batch_size_scaling_linear_lr_scaling_8_nodes_3325606.log
train_tokenizer_no_warmup_ablation_3322446.log
train_tokenizer_no_warmup_ablation_3325608.log
[tum_dbd0378@hkn0819 jafar]$ less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/",,terminal_output
391,3435165,"TERMINAL",0,0,"t",,terminal_output
392,3435422,"TERMINAL",0,0,"r",,terminal_output
393,3435500,"TERMINAL",0,0,"a 
in",,terminal_output
394,3435693,"TERMINAL",0,0,"_",,terminal_output
395,3435796,"TERMINAL",0,0,"d",,terminal_output
396,3435936,"TERMINAL",0,0,"y",,terminal_output
397,3436023,"TERMINAL",0,0,"namics_lr_1e-4_larger_ffn_335",,terminal_output
398,3443118,"TERMINAL",0,0,"3",,terminal_output
399,3443363,"TERMINAL",0,0,"271.log ",,terminal_output
400,3443999,"TERMINAL",0,0,"

",,terminal_output
401,3444062,"TERMINAL",0,0,"=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2771207
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0516
SLURM_JOB_START_TIME=1752748479
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752834879
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3353271
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0516,0518]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1990.localdomain
SLURM_JOB_ID=3353271
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
402,3445049,"TERMINAL",0,0,"
SLURM_NTASKS_PER_NODE=4
SLURM_JOB_GID=502226
SLURM_JOB_NODELIST=hkn[0516,0518]
GpuFreq=control_disabled
GpuFreq=control_disabled
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c :",,terminal_output
403,3445377,"TERMINAL",0,0,"
lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
2025-07-17 12:35:34.171919: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we :
re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.177867: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.183723: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.185789: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.410170: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.414120: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.414129: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.424831: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:59.978488: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:00.116831: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:00.861762: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:01.235616: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.001680: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.557867: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.629343: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.831381: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
:
2025-07-17 12:36:06.893702: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:06.981293: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:07.887393: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:07.975065: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
2025-07-17 12:36:10.468323: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
2025-07-17 12:36:11.034549: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:11.398745: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:11.644909: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
wandb: Currently logged in as: emergenz (p-doom) to https://api.wandb.ai. Use `wandb login --relogin` to forc e relogin
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
wandb: Tracking run with wandb version 0.19.11
wandb: Run data is saved locally in /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/wandb/run-202507 17_123614-a9l2va8l
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run dynamics-lr-1e-4-larger-ffn
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/a9l2va8l
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca :",,terminal_output
404,3445576,"TERMINAL",0,0,"
ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
:
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
WARNING:absl:Missing metrics for step 100000
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
:
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca :
ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
:
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b :",,terminal_output
405,3445901,"TERMINAL",0,0,"
ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={' bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8Traceback (most recent call last):
, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'opt_state': [{'mu': {'par ams': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s :
harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axi    raise ValueError(
s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff( lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memor y_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data' : 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': D iff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), d type=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_k ind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8 , axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff (lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtyp e=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind =device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device :
)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, a xis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lh s={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=f loat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=de vice))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, aVal ueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {'par ams': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, xis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs= None)}}}}, 'nu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(5 12,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), m emory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh= Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {' Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('d ata': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(5 12, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec( ), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512, ), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None) }, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_5': {'Despec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dty pe=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kin d=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=floa t32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devic e)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(l hs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ty pes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'binse_1': Diff(lhs={'bias': ShapeDtypeStruct(shape= (512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), :
 memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6':  {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape= (512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpe c(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(51 2,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(as': S hapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs =None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Na medSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': S hapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': Shap eDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), sp ec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8,   File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_d ynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Den se_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data ': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512,  512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh(' data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_ 1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(sha pe=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec (), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1 ': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Me sh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(sha pe=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partition Spec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape= (512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,))), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t :
ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStru, spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memor y_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data' : 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': D iff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), d type=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_k ind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'ct(shape=(512,), dtype=float32, sharding=Named Sharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Shap eDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDt ypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None) }, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,))bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8, spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}}, None,  None]}
, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'opt_state': [{'mu': {'par ams': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axiTraceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
:
    restored = self._checkpointer.restore(restore_directory, args=args)
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff( lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device))}, rhs=None)}, 'STBlock_5': {'Dense    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mes    restored = self._checkpointer.restore(restore_directory, args= args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': { 'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtype Struct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStr uct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'ST Block_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), s    return super().restore(directory, *args, **kwargs)
pec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=N one)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sha peDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeD typeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding= :
NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None )}, 'STBlock_5': {'Dense_1': Diff(lhs={'bia  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
    restored = self._restore(directory, args=ckpt_args)
s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8, a  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
xis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': {'decoder':  {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'S TBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)),     return self._handler.restore(directory, args=args)
spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shard ing=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs= None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sh apeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': Shape DtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding =NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=Non e)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bi  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py"", line 837,  in restore
as': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Au to,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) }, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shard ing=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kern el': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias' :
: ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8,     restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
    return self._impl.restore(
axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Den se_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data ': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512,  512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh(' data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_ 1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in  restore
), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(m esh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3' : {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mes h('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shap e=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionS pec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=( 512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh =Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': { 'Dense_1': Diff(lhs={'bias': ShapeDtypeStru    return self._handler_impl.restore(directory, args=args)
ct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partiti onSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STB lock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(m esh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStru ct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct( shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionS pec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,))  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737 , in restore
, spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}}, None, None]}
    raise ValueError(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
:
    return self._handler.restore(directory, args=args)
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8,   File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py"", line 837,  in restore
    restored[item_name] = handler.restore(
axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Di ff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_5': {'De  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, i n restore
    return self._impl.restore(
nse_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('dat a': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512 , 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense _1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in  restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder' : {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s :
pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,))ValueError: User-provided restore item and on-disk value metadata  tree structures do not match: {'params': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': S hapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs =None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Na medSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': S hapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': Shap eDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, , spec=PartitionSpec(),  memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3':  {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape= (512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpe c(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(51 2,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=M esh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'D ense_1': Diff(lhs={'axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStru ct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct( shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionS pec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBloc k_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh =Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct( shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Debias': ShapeDtypeStruct(shape=(512,), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8 , axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff (lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtyp e=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind =device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device )), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8nse _1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dt :",,terminal_output
406,3446098,"TERMINAL",0,0,"
ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, r hs=None)}}}}, 'opt_state': [{'mu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDt ypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None) }, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, aximesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memor y_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data' : 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': D iff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,))s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(sha pe=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partition Spec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape= (512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4':  {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape= (512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpe c(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense, spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'_1': D iff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), d type=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_k ind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=fl :
oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8 , axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff (lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtyp e=float32, sharding=NamedSharding(mesbias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'S TBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeSt ruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8h=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff (lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtyp e=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind =device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device )), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, a xis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lh s={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), s, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'opt_state':  [{'mu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memor y_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axipec=PartitionSpec(), memory_kind=device)), 'kernel':  ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Au to,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': Sh apeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shard ing=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs= None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sh apeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke :
rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_5': {'Denses': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(me sh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruc t(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Part itionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(s hape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSp ec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShardin g(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock _7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh= Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(s hape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, a_1': Diff(lhs={'bias': ShapeDtyp eStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesxis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': { 'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding =NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel' : ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': S hapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs =None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Na medSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSp ec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'S TBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeSt ruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=P artitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruc t(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), sspec=Partiti onSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STB lock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(m esh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStru :
ct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct( shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionS pec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBloc k_5': {'Dense_1': Diff(lhs={'bipec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(5 12, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec( ), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512, ), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memo ry_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mes h('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Den se_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data ': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512,  512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'biaas': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense _1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ty pes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDt",,terminal_output
407,3446428,"TERMINAL",0,0,"ypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, aaxis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memor y_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(xis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': {'decoder': {'STBlock_0': { 'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh(' data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=( 512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec :
(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512 ,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), mem ory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Me sh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'De nse_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('dat a': 8, axis_types=(Auto,)), ), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruspec=PartitionSpec(), memory_kind=devi ce)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff( lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs ={'bict(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,))as': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel':  ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': Sha peDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=N one)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sha peDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, , spec=PartitionSp ec(), memory_kind=device))}, rhs=None)}}}}}, None, None]}
axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Den se_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data ': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, :
 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh(' data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_ 1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 51 2), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), mem ory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dty pe=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kin d=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('dat a': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1':  Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Di ff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}}, None, None]}
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che :
ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8,     return super().restore(directory, *args, **kwargs)
axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Di ff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_5': {'De  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
nse_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('dat a': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512 , 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh( :
'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense _1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(    restored = self._restore(directory, args=ckpt_args)
mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder' : {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,))  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
, spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel':  ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': Sha peDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=N one)}, 'STBlock_5': {'Dense_1': Diff(lhs={'    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}}}}, 'opt_state': [{'mu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': Shap eDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), sp ec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=No ne)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Named Sharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Shap eDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDt ypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=N :
amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None) }, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtype Struct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStr uct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'ST Block_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStr uct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mes    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': { 'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtype Struct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStr uct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'ST Block_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), s    return self._impl.restore(
pec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=N one)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sha peDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeD typeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding= NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None )}, 'STBlock_5': {'Dense_1': Diff(lhs={'bia  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in  restore
s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne :
l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8, a    return self._handler_impl.restore(directory, args=args)
xis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': {'decoder':  {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'S TBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)),   File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/. venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737 , in restore
spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shard ing=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs= None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sh apeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': Shape DtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding =NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=Non e)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bi    raise ValueError(
as': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Au to,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) }, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shard ing=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kern el': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias' : ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, ValueError: User-provided restore item and on-disk value metadata  tree structures do not match: {'params': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': S hapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs =None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Na medSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': S hapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': Shap eDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), sp ec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias' :
: ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding =NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel' : ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': S hapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStru ct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct( shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionS pec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBloc k_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh =Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct( shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'De), memory_kind=device)), 'kernel': ShapeDtypeS truct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStru ct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partiti onSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STB lock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(m esh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStru ct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStrunse _1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('da ta': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1' : Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(ct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('dat a': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512 , 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,),  dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory _kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense _1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data':  8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 5 12), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,))mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1' :
: Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512) , dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memor y_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data' : 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': D iff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}}, None, None]}
, spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel':  ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': Sha peDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=N one)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': Sha peDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeD typeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding= NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None )}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeD typeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'opt_state': [{'mu': {'params': {'decoder': {'STB lock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(m esh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStru ct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct( shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionS pec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBloc k_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh =Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct( shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(sha pe=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec (), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4 ': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Me sh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(sha pe=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partition Spec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape= :
(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(),  memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6':  {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh( 'data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape= (512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpe c(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(51 2,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), me mory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=M esh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'ST Block_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStr uct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct (shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partition Spec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlo ck_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct (shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(sh ape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpe c(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_ 4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=M esh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(sh ape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partitio nSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape =(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec() , memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(me sh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6':  {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh ('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape =(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSp ec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(5 12,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), m emory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh= Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'para ms': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'k ernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ty pes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bi as': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Au to,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) }, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shard ing=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kern :
el': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias' : ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto, )), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding =NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel' : ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': S hapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs =None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=Na medSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': S hapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': Shap eDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), sp ec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=No ne)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel':  ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': Sha peDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=N :
one)}}}}}, None, None]}
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar :
ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={' bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, r hs=None)}}}}, 'opt_state': [{'mu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDt ypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None) }, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt :",,terminal_output
408,3446622,"TERMINAL",0,0,"
ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtype Struct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStr uct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'ST Block_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStr uct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct (shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partition Spec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlo ck_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct (shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeD typeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding= NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None )}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeD typeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtyp eStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'S TBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeSt ruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=P artitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruc t(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partitio nSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShar ding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBl :
ock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(me sh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruc t(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Part itionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Dif f(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis _types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dty pe=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kin d=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=floa t32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devic e)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(l hs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ty pes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Di ff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff( lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type :
s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={' bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}}, None, None]}
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi :
s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={' bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))} , rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardi :
ng=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kerne l': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types= (Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,) ), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sh arding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, r hs=None)}}}}, 'opt_state': [{'mu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDt ypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec= PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=N amedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None) }, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSha rding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDt ypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), s pec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtype Struct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Par titionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Name dSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShardi ng(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtype Struct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStr uct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partit ionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'ST Block_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding( mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStr uct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct (shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partition Spec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlo ck_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mes h=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct (shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeD typeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec =PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding= NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None )}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSh arding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeD typeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)),  spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtyp eStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Pa rtitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=Nam edSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)},  'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedShard ing(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtyp :
eStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spe c=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeSt ruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Parti tionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedS harding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'S TBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding (mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeSt ruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=P artitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruc t(shape=(512,), dtype=fl",,terminal_output
409,3446739,"TERMINAL",0,0,"oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Partitio nSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedShar ding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBl ock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(me sh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruc t(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Part itionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Dif f(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis _types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dty pe=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kin d=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=floa t32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devic e)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(l hs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ty pes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Di ff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki :
MSLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
MSLURM_CONF=/etc/slurm/slurm.conf
MSLURM_NODEID=0
MSLURM_JOB_ID=3353271
MSLURM_SUBMIT_HOST=hkn1990.localdomain
MSLURM_NNODES=2
MSLURM_NPROCS=8
MSLURM_PRIO_PROCESS=0
MSLURM_JOB_ACCOUNT=hk-project-p0023960
MSLURM_NODELIST=hkn[0516,0518]
MSLURM_SCRIPT_CONTEXT=prolog_task
MSLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
MSLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
MSLURM_NTASKS=8
MSLURM_CPUS_PER_TASK=5
MSLURM_PROCID=0
MSLURM_JOB_QOS=normal
MSLURM_JOBID=3353271
MSLURM_JOB_NUM_NODES=2
MSLURM_OOM_KILL_STEP=0
MSLURM_TRES_PER_TASK=cpu=5
MSLURM_JOB_PARTITION=accelerated
MSLURM_GTIDS=0
MSLURM_GPUS_ON_NODE=4
MSLURM_JOB_CPUS_PER_NODE=24(x2)
MSLURM_CPUS_ON_NODE=24
MSLURM_JOB_END_TIME=1752834879
MSLURM_CLUSTER_NAME=hk
MSLURM_JOB_START_TIME=1752748479
MSLURMD_NODENAME=hkn0516
MSLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
MSLURM_LOCALID=0
MSLURM_JOB_GPUS=0,1,2,3
MSLURM_TASK_PID=2771207
MSLURM_JOB_UID=996262
MSLURM_TASKS_PER_NODE=4(x2)
MSLURM_JOB_USER=tum_dbd0378
M    --data_dir $array_records_dir \
M    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
M    --project jafar \
M    --entity instant-uv \
M    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
M    --name=dynamics-lr-1e-4-larger-ffn \
M    --log_checkpoint_interval=1000 \
M    --log \
M    --log_image_interval=1000 \
M    --max_lr=1.5e-4 \
M    --init_lr=0 \
M    --dyna_num_blocks=6 \
M    --dyna_ffn_dim=2048 \
M    --lam_num_blocks=4 \
M    --lam_ffn_dim=2048 \
M    --tokenizer_num_blocks=8 \
M    --tokenizer_ffn_dim=512 \
M    --batch_size=96 \
M    --ckpt_dir $CHECKPOINT_DIR \
M    --save_ckpt \
Msrun python train_dynamics.py \
M
Menv | grep SLURM
M
Mscaling/train_tokenizer_lr_sweep_1e-4
Mtokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-M
Mmkdir -p $CHECKPOINT_DIR
Mng-modelsize-scaling/$job_name
MCHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotrainiM
Mslurm_job_id=$SLURM_JOB_ID
Mjob_name=$SLURM_JOB_NAME
M
Mords_chunked
Marray_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecM
Msource .venv/bin/activate
Mmodule unload devel/cuda/12.4
Mmodule unload mpi/openmpi/5.0
M
Mcat $0
M# Log the sbatch script
M
M#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn
M#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
M#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
M#SBATCH --gres=gpu:4
M#SBATCH --cpus-per-task=5
M#SBATCH --partition=accelerated
M#SBATCH --time=24:00:00
M#SBATCH --ntasks-per-node=4
M#SBATCH --nodes=2
M
M#!/usr/bin/env bash

:",,terminal_output
410,3446950,"TERMINAL",0,0,"

:",,terminal_output
411,3447997,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
412,3448522,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log ",,terminal_output
413,3449107,"TERMINAL",0,0,"",,terminal_output
414,3449570,"TERMINAL",0,0,"",,terminal_output
415,3452784,"TERMINAL",0,0,"

",,terminal_output
416,3453631,"TERMINAL",0,0,"q",,terminal_output
417,3453764,"TERMINAL",0,0,"u",,terminal_output
418,3453962,"TERMINAL",0,0,"eu",,terminal_output
419,3454094,"TERMINAL",0,0,"e

(BEvery 1.0s: squeue --mehkn0819.localdomain: Thu Jul 17 12:38:19 2025JOBID PARTITION     NAME     USER ST	TIME  NODES NODELIST(REASON)3353272 accelerat interact tum_dbd0  R	3:40	  2 hkn[0819,0821]",,terminal_output
420,3455112,"TERMINAL",0,0,"201",,terminal_output
421,3456128,"TERMINAL",0,0,"12",,terminal_output
422,3457169,"TERMINAL",0,0,"23",,terminal_output
423,3458218,"TERMINAL",0,0,"34",,terminal_output
424,3459263,"TERMINAL",0,0,"45",,terminal_output
425,3460252,"TERMINAL",0,0,"56",,terminal_output
426,3461024,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
427,3463788,"TERMINAL",0,0,"queue",,terminal_output
428,3464598,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
429,3464751,"TERMINAL",0,0,"",,terminal_output
430,3466613,"TERMINAL",0,0,"",,terminal_output
431,3466752,"TERMINAL",0,0,"",,terminal_output
432,3466904,"TERMINAL",0,0,"",,terminal_output
433,3467065,"TERMINAL",0,0,"",,terminal_output
434,3467231,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dyamics_lr_1e-4_larger_ffn_335",,terminal_output
435,3467708,"TERMINAL",0,0,"l /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
436,3467766,"TERMINAL",0,0,"s /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
437,3467846,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
438,3467915,"TERMINAL",0,0,"- /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
439,3468083,"TERMINAL",0,0,"l /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
440,3468174,"TERMINAL",0,0,"h /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
441,3468288,"TERMINAL",0,0,"

",,terminal_output
442,3468970,"TERMINAL",0,0,"",,terminal_output
443,3470298,"TERMINAL",0,0,"",,terminal_output
444,3470501,"TERMINAL",0,0,"",,terminal_output
445,3470653,"TERMINAL",0,0,"
",,terminal_output
446,3470838,"TERMINAL",0,0,"

",,terminal_output
447,3471381,"TERMINAL",0,0,"

total 1.7M
drwxr-x---+ 2 tum_dbd0378 hk-project-p0023960 4.0K Jul 16 21:06 big-runs
drwxr-x---+ 2 tum_dbd0378 hk-project-p0023960 4.0K Jul 15 12:12 dynamics-cotraining-modelsize-scaling
drwxr-x---+ 2 tum_dbd0378 hk-project-p0023960 4.0K Jul 15 11:46 tokenizer-lr-scaling
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 616K Jul 16 21:12 train_dynamics_lr_1e-4_larger_ffn_3352437.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 200K Jul 17 12:37 train_dynamics_lr_1e-4_larger_ffn_3353271.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 167K Jul  6 22:11 train_tokenizer_batch_size_scaling_linear_lr_scaling_16_nodes_3322407.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 243K Jul  7 10:23 train_tokenizer_batch_size_scaling_linear_lr_scaling_16_nodes_3325607.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  24K Jul  6 21:33 train_tokenizer_batch_size_scaling_linear_lr_scaling_2_nodes_3322404.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  34K Jul  7 10:22 train_tokenizer_batch_size_scaling_linear_lr_scaling_2_nodes_3325603.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  44K Jul  6 21:33 train_tokenizer_batch_size_scaling_linear_lr_scaling_4_nodes_3322405.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  64K Jul  7 10:22 train_tokenizer_batch_size_scaling_linear_lr_scaling_4_nodes_3325605.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  85K Jul  6 21:36 train_tokenizer_batch_size_scaling_linear_lr_scaling_8_nodes_3322406.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 124K Jul  7 10:22 train_tokenizer_batch_size_scaling_linear_lr_scaling_8_nodes_3325606.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  14K Jul  6 21:33 train_tokenizer_no_warmup_ablation_3322446.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  19K Jul  7 10:23 train_tokenizer_no_warmup_ablation_3325608.log
[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
448,3476947,"TERMINAL",0,0,"ls -lh /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/",,terminal_output
449,3477091,"TERMINAL",0,0,"
queue",,terminal_output
450,3477147,"TERMINAL",0,0,"",,terminal_output
451,3477255,"TERMINAL",0,0,"",,terminal_output
452,3477815,"TERMINAL",0,0,"ls -lh /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/",,terminal_output
453,3478131,"TERMINAL",0,0,"
",,terminal_output
454,3478261,"TERMINAL",0,0,"",,terminal_output
455,3478707,"TERMINAL",0,0,"",,terminal_output
456,3478984,"TERMINAL",0,0,"",,terminal_output
457,3479147,"TERMINAL",0,0,"",,terminal_output
458,3479343,"TERMINAL",0,0,"",,terminal_output
459,3479831,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
460,3480219,"TERMINAL",0,0,"l /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
461,3480279,"TERMINAL",0,0,"e /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
462,3480346,"TERMINAL",0,0,"s /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
463,3480500,"TERMINAL",0,0,"s /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
464,3480691,"TERMINAL",0,0,"",,terminal_output
465,3481389,"TERMINAL",0,0,"train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
466,3481618,"TERMINAL",0,0,"train_dynamics_lr_1e-4_larger_ffn_3353271.log ",,terminal_output
467,3481827,"TERMINAL",0,0,"

=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2771207
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0516
SLURM_JOB_START_TIME=1752748479
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752834879
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3353271
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0516,0518]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1990.localdomain
SLURM_JOB_ID=3353271
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
468,3482964,"TERMINAL",0,0,"
...skipping...
ding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBl ock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(me sh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruc t(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=Part itionSpec(), memory_kind=device))}, rhs=None)}}}}, 'nu': {'params': {'decoder': {'STBlock_0': {'Dense_1': Dif f(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis _types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dty pe=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kin d=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=floa t32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devic e)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(l hs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ty pes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype= float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=d evice))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Di ff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dt ype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_ki nd=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8,  axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff( lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype =float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind= device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={' bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}}}}}, None, None]}
ESC[1;34mwandbESC[0m: 
ESC[1;34mwandbESC[0m: 🚀 View run ESC[33mdynamics-lr-1e-4-larger-ffnESC[0m at: ESC[34mhttps://wandb.ai/instan t-uv/jafar/runs/a9l2va8lESC[0m
ESC[1;34mwandbESC[0m: Find logs at: ESC[1;35m../../../../hkfs/home/project/hk-project-p0023960/tum_dbd0378/ja far/wandb/run-20250717_123614-a9l2va8l/logsESC[0m
srun: error: hkn0516: task 1: Exited with exit code 1
srun: error: hkn0518: task 4: Exited with exit code 1
srun: error: hkn0518: tasks 5-6: Exited with exit code 1
srun: error: hkn0516: tasks 0,2-3: Exited with exit code 1
srun: error: hkn0518: task 7: Exited with exit code 1

============================= JOB FEEDBACK =============================

Job ID: 3353271
Cluster: hk
User/Group: tum_dbd0378/hk-project-p0023960
Account: hk-project-p0023960
State: FAILED (exit code 1)
Partition: accelerated
Nodes: 2
Cores per node: 24
Nodelist: hkn[0516,0518]
CPU Utilized: 00:16:26
CPU Efficiency: 11.74% of 02:20:00 core-walltime
Job Wall-clock time: 00:02:55
Starttime: Thu Jul 17 12:34:39 2025
Endtime: Thu Jul 17 12:37:34 2025
Memory Utilized: 21.24 GB (estimated maximum)
Memory Efficiency: 0.00% of 0.00 MB (0.00 MB/node)
Energy Consumed: 189456 Joule / 52.6266666666667 Watthours
Average node power draw: 1082.60571428571 Watt
(END)",,terminal_output
469,3484253,"TERMINAL",0,0,"
MSLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
MSLURM_CONF=/etc/slurm/slurm.conf
MSLURM_NODEID=0
MSLURM_JOB_ID=3353271
MSLURM_SUBMIT_HOST=hkn1990.localdomain
MSLURM_NNODES=2
MSLURM_NPROCS=8
MSLURM_PRIO_PROCESS=0
MSLURM_JOB_ACCOUNT=hk-project-p0023960
MSLURM_NODELIST=hkn[0516,0518]
MSLURM_SCRIPT_CONTEXT=prolog_task
MSLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
MSLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
MSLURM_NTASKS=8
MSLURM_CPUS_PER_TASK=5
MSLURM_PROCID=0
MSLURM_JOB_QOS=normal
MSLURM_JOBID=3353271
MSLURM_JOB_NUM_NODES=2
MSLURM_OOM_KILL_STEP=0
MSLURM_TRES_PER_TASK=cpu=5
MSLURM_JOB_PARTITION=accelerated
MSLURM_GTIDS=0
MSLURM_GPUS_ON_NODE=4
MSLURM_JOB_CPUS_PER_NODE=24(x2)
MSLURM_CPUS_ON_NODE=24
MSLURM_JOB_END_TIME=1752834879
MSLURM_CLUSTER_NAME=hk
MSLURM_JOB_START_TIME=1752748479
MSLURMD_NODENAME=hkn0516
MSLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
MSLURM_LOCALID=0
MSLURM_JOB_GPUS=0,1,2,3
MSLURM_TASK_PID=2771207
MSLURM_JOB_UID=996262
MSLURM_TASKS_PER_NODE=4(x2)
MSLURM_JOB_USER=tum_dbd0378
M    --data_dir $array_records_dir \
M    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
M    --project jafar \
M    --entity instant-uv \
M    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
M    --name=dynamics-lr-1e-4-larger-ffn \
M    --log_checkpoint_interval=1000 \
M    --log \
M    --log_image_interval=1000 \
M    --max_lr=1.5e-4 \
M    --init_lr=0 \
M    --dyna_num_blocks=6 \
M    --dyna_ffn_dim=2048 \
M    --lam_num_blocks=4 \
M    --lam_ffn_dim=2048 \
M    --tokenizer_num_blocks=8 \
M    --tokenizer_ffn_dim=512 \
M    --batch_size=96 \
M    --ckpt_dir $CHECKPOINT_DIR \
M    --save_ckpt \
Msrun python train_dynamics.py \
M
Menv | grep SLURM
M
Mscaling/train_tokenizer_lr_sweep_1e-4
Mtokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-M
Mmkdir -p $CHECKPOINT_DIR
Mng-modelsize-scaling/$job_name
MCHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotrainiM
Mslurm_job_id=$SLURM_JOB_ID
Mjob_name=$SLURM_JOB_NAME
M
Mords_chunked
Marray_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecM
Msource .venv/bin/activate
Mmodule unload devel/cuda/12.4
Mmodule unload mpi/openmpi/5.0
M
Mcat $0
M# Log the sbatch script
M
M#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn
M#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
M#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
M#SBATCH --gres=gpu:4
M#SBATCH --cpus-per-task=5
M#SBATCH --partition=accelerated
M#SBATCH --time=24:00:00
M#SBATCH --ntasks-per-node=4
M#SBATCH --nodes=2
M
M#!/usr/bin/env bash

:",,terminal_output
470,3484368,"TERMINAL",0,0,"

:",,terminal_output
471,3485842,"TERMINAL",0,0,"
SLURM_NTASKS_PER_NODE=4
SLURM_JOB_GID=502226
SLURM_JOB_NODELIST=hkn[0516,0518]
GpuFreq=control_disabled
GpuFreq=control_disabled
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
:",,terminal_output
472,3496079,"TERMINAL",0,0,"
Running on 8 devices.
2025-07-17 12:35:34.171919: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.177867: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.183723: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.185789: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.410170: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.414120: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.414129: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:34.424831: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:35:59.978488: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:00.116831: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:00.861762: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:01.235616: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.001680: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.557867: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.629343: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:03.831381: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we :",,terminal_output
473,3497158,"TERMINAL",0,0,"
re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:06.893702: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:06.981293: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:07.887393: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:07.975065: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
2025-07-17 12:36:10.468323: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
2025-07-17 12:36:11.034549: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:11.398745: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 12:36:11.644909: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
wandb: Currently logged in as: emergenz (p-doom) to https://api.wandb.ai. Use `wandb login --relogin` to forc e relogin
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
wandb: Tracking run with wandb version 0.19.11
wandb: Run data is saved locally in /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/wandb/run-202507 17_123614-a9l2va8l
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run dynamics-lr-1e-4-larger-ffn
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/a9l2va8l
Parameter counts:
{'tokenizer': 42192112, 'lam': 17229792, 'dynamics': 26555904, 'total': 85977808}
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Dropping 2 examples of 89394 examples (shard 8).
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
:",,terminal_output
474,3506528,"TERMINAL",0,0,"
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
WARNING:absl:Missing metrics for step 100000
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca :",,terminal_output
475,3507053,"TERMINAL",0,0,"
ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca :",,terminal_output
476,3507473,"TERMINAL",0,0,"
ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 144000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/144000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 60000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/060000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 146000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/146000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 120000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 :",,terminal_output
477,3520524,"TERMINAL",0,0,"
, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl oat32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=dev ice))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)),  'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_ types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={' bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float 32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device ))}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ke rnel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_typ es=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bia s': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Aut o,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, :",,terminal_output
478,3606171,"TERMINAL",0,0,"
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:",,terminal_output
479,3606228,"TERMINAL",0,0,"
 ESCESCOOAA
MTraceback (most recent call last):

:",,terminal_output
480,3606306,"TERMINAL",0,0,"
 ESCESCOOAA
Mckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore

:",,terminal_output
481,3606695,"TERMINAL",0,0,"
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored = self._checkpointer.restore(restore_directory, args=args)

:
 ESCESCOOAA
Mckpoint/checkpoint_manager.py"", line 1608, in restore

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che
:
 ESCESCOOAA
M    restored_tokenizer = tokenizer_checkpoint_manager.restore(

:
 ESCESCOOAA
Ments

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon
:",,terminal_output
482,3607106,"TERMINAL",0,0,"
 ESCESCOOAA
M    train_state = restore_genie_components(

:
 ESCESCOOAA
M  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>

:
 ESCESCOOAA
MTraceback (most recent call last):

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 140000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:",,terminal_output
483,3607439,"TERMINAL",0,0,"
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 20000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 100000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:",,terminal_output
484,3607499,"TERMINAL",0,0,"
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 145000

:",,terminal_output
485,3607832,"TERMINAL",0,0,"
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 40000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.

:
 ESCESCOOAA
MERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca
:
 ESCESCOOAA
MWARNING:absl:Missing metrics for step 80000

:
 ESCESCOOAA
Mling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.

:",,terminal_output
486,3609278,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",0,0,"",python,tab
487,3609279,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",11482,0,"",python,selection_command
488,3675729,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
489,3676557,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"4",shellscript,content
490,3676852,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,0,"048",shellscript,content
491,3676854,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,2,"",shellscript,content
492,3676856,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
493,3679072,"TERMINAL",0,0,"ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 :",,terminal_output
494,3680051,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1151,0,"",shellscript,selection_command
495,3680301,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1122,0,"",shellscript,selection_command
496,3680331,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1096,0,"",shellscript,selection_command
497,3680364,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1072,0,"",shellscript,selection_command
498,3680397,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1047,0,"",shellscript,selection_command
499,3680429,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1045,0,"",shellscript,selection_command
500,3680548,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1029,0,"",shellscript,selection_command
501,3680741,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",906,0,"",shellscript,selection_command
502,3681324,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",902,0,"",shellscript,selection_command
503,3681446,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",900,0,"",shellscript,selection_command
504,3681787,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,0,"",shellscript,selection_command
505,3682127,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,1,"/",shellscript,selection_command
506,3682441,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,128,"/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4
",shellscript,selection_command
507,3682634,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",901,127,"/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",shellscript,selection_command
508,3682855,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1027,0,"",shellscript,selection_command
509,3683935,"TERMINAL",0,0,"
...skipping...
ling/train_tokenizer_lr_sweep_1e-4/120000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 80000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/080000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 40000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/040000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 145000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/145000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 100000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/100000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 20000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/020000/metrics/metrics not found.
WARNING:absl:Missing metrics for step 140000
ERROR:absl:File /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-sca ling/train_tokenizer_lr_sweep_1e-4/140000/metrics/metrics not found.
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 :",,terminal_output
510,3684893,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
511,3685349,"TERMINAL",0,0,"[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
512,3685722,"TERMINAL",0,0,"/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",,terminal_output
513,3686154,"TERMINAL",0,0,"/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",,terminal_output
514,3686599,"TERMINAL",0,0,"l/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4s/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",,terminal_output
515,3686688,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4",,terminal_output
516,3686862,"TERMINAL",0,0,"

",,terminal_output
517,3688806,"TERMINAL",0,0,"",,terminal_output
518,3689750,"TERMINAL",0,0,"",,terminal_output
519,3690070,"TERMINAL",0,0,"",,terminal_output
520,3690586,"TERMINAL",0,0,"",,terminal_output
521,3693864,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",0,0,"#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/%x_%j.log
#SBATCH --job-name=train_tokenizer_lr_sweep_1e-4_larger_ffn_5_blocks

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$PWD/checkpoints/$job_name
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --ffn_dim=2048 \
    --num_blocks=5 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=mixed-precision-tokenizer-lr-1e-4-larger-ffn_5_blocks \
    --tags tokenizer mixed-precision 1e-4 mixed-precision flash-attention larger-ffn 5-blocks \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir",shellscript,tab
522,3696514,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",902,0,"",shellscript,selection_command
523,3696777,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",869,0,"",shellscript,selection_command
524,3696795,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",836,0,"",shellscript,selection_command
525,3696815,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",835,0,"",shellscript,selection_command
526,3696850,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",818,0,"",shellscript,selection_command
527,3696883,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",817,0,"",shellscript,selection_command
528,3697033,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",792,0,"",shellscript,selection_command
529,3697201,"experiments/train_tokenizer_lr_1e-4_larger_ffn.sh",750,0,"",shellscript,selection_command
530,3731395,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
531,3731748,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"51",shellscript,content
532,3731750,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,3,"",shellscript,content
533,3731752,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
534,3731910,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"8",shellscript,content
535,3731912,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,0,"",shellscript,selection_command
536,3733283,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"4",shellscript,content
537,3733630,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,0,"048",shellscript,content
538,3733633,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,2,"",shellscript,content
539,3733635,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
540,3734529,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"51",shellscript,content
541,3734534,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,3,"",shellscript,content
542,3734535,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
543,3734617,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"8",shellscript,content
544,3734621,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,0,"",shellscript,selection_command
545,3735905,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"4",shellscript,content
546,3735941,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,0,"048",shellscript,content
547,3735944,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,2,"",shellscript,content
548,3735946,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
549,3739050,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"51",shellscript,content
550,3739059,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,3,"",shellscript,content
551,3739061,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1177,0,"",shellscript,selection_command
552,3739246,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,1,"8",shellscript,content
553,3739248,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1210,0,"",shellscript,selection_command
554,3744659,"TERMINAL",0,0,"",,terminal_output
555,3744817,"TERMINAL",0,0,"",,terminal_output
556,3745092,"TERMINAL",0,0,"ess /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log ",,terminal_output
557,3745664,"TERMINAL",0,0,"

=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2771207
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0516
SLURM_JOB_START_TIME=1752748479
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752834879
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3353271
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0516,0518]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1990.localdomain
SLURM_JOB_ID=3353271
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
558,3753467,"TERMINAL",0,0,"
/",,terminal_output
559,3753581,"TERMINAL",0,0,"rr",,terminal_output
560,3753675,"TERMINAL",0,0,"eess",,terminal_output
561,3753776,"TERMINAL",0,0,"tt",,terminal_output
562,3753885,"TERMINAL",0,0,"oo",,terminal_output
563,3753957,"TERMINAL",0,0,"rr",,terminal_output
564,3754010,"TERMINAL",0,0,"ee",,terminal_output
565,3754133,"TERMINAL",0,0,"
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2771207
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0516
SLURM_JOB_START_TIME=1752748479
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752834879
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3353271
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0516,0518]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1990.localdomain
SLURM_JOB_ID=3353271
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2771207
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0516
SLURM_JOB_START_TIME=1752748479
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752834879
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3353271
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0516,0518]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1990.localdomain
SLURM_JOB_ID=3353271
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
...skipping...
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl :",,terminal_output
566,3756003,"TERMINAL",0,0,"
 ESCESC",,terminal_output
567,3756909,"TERMINAL",0,0,"ESCESC
/",,terminal_output
568,3759717,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",0,0,"",python,tab
569,3759718,"/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py",11482,0,"",python,selection_command
570,3868485,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
571,3895185,"TERMINAL",0,0,"ESCESCESCESCESCESCqq",,terminal_output
572,3896595,"TERMINAL",0,0,"    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
Traceback (most recent call last):
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_dynamics.py"", line 318, in <module>
    train_state = restore_genie_components(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/genie.py"", line 318, in restore_genie_compon ents
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/checkpoint_manager.py"", line 1608, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/async_checkpointer.py"", line 558, in restore
    return super().restore(directory, *args, **kwargs)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 300, in restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    restored = self._restore(directory, args=ckpt_args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/checkpointers/checkpointer.py"", line 319, in _restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    return self._handler.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/composite_checkpoint_handler.py"", line 837, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    restored[item_name] = handler.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/standard_checkpoint_handler.py"", line 246, in restore
    return self._impl.restore(
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/pytree_checkpoint_handler.py"", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
  File ""/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/orbax/che ckpoint/_src/handlers/base_pytree_checkpoint_handler.py"", line 737, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'params': {' params': {'decoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32 , sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)) , 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axi s_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs= {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types =(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=flo at32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=devi ce))}, rhs=None)}, 'STBlock_2': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, s harding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), ' kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_t ypes=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_3': {'Dense_1': Diff(lhs={'b ias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(A uto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) )}, rhs=None)}, 'STBlock_4': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shar ding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'ker nel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_5': {'Dense_1': Diff(lhs={'bias ': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto ,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32,  sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))},  rhs=None)}, 'STBlock_6': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, shardin g=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel ': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=( Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_7': {'Dense_1': Diff(lhs={'bias':  ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)) , spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sha rding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rh s=None)}}, 'encoder': {'STBlock_0': {'Dense_1': Diff(lhs={'bias': ShapeDtypeStruct(shape=(512,), dtype=float3 2, sharding=NamedSharding(mesh=Mesh('data': 8, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) ), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, ax is_types=(Auto,)), spec=PartitionSpec(), memory_kind=device))}, rhs=None)}, 'STBlock_1': {'Dense_1': Diff(lhs ={'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=NamedSharding(mesh=Mesh('data': 8, axis_type s=(Auto,)), spec=PartitionSpec(), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=fl :",,terminal_output
573,3897164,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
574,3898145,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log ",,terminal_output
575,3898941,"TERMINAL",0,0,"",,terminal_output
576,3899376,"TERMINAL",0,0,"",,terminal_output
577,3899536,"TERMINAL",0,0,"",,terminal_output
578,3899709,"TERMINAL",0,0,"",,terminal_output
579,3899854,"TERMINAL",0,0,"",,terminal_output
580,3900084,"TERMINAL",0,0,"",,terminal_output
581,3900188,"TERMINAL",0,0,"335",,terminal_output
582,3900393,"TERMINAL",0,0,"",,terminal_output
583,3902587,"TERMINAL",0,0,"
train_dynamics_lr_1e-4_larger_ffn_3352437.log  train_dynamics_lr_1e-4_larger_ffn_3353271.log
[tum_dbd0378@hkn0819 jafar]$ less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_335",,terminal_output
584,3907413,"TERMINAL",0,0,"",,terminal_output
585,3907784,"TERMINAL",0,0,"",,terminal_output
586,3908117,"TERMINAL",0,0,"",,terminal_output
587,3908309,"TERMINAL",0,0,"",,terminal_output
588,3908468,"TERMINAL",0,0,"",,terminal_output
589,3908531,"TERMINAL",0,0,"",,terminal_output
590,3908722,"TERMINAL",0,0,"",,terminal_output
591,3908835,"TERMINAL",0,0,"",,terminal_output
592,3908936,"TERMINAL",0,0,"",,terminal_output
593,3909162,"TERMINAL",0,0,"",,terminal_output
594,3909254,"TERMINAL",0,0,"",,terminal_output
595,3909475,"TERMINAL",0,0,"
",,terminal_output
596,3909685,"TERMINAL",0,0,"

",,terminal_output
597,3909881,"TERMINAL",0,0,"",,terminal_output
598,3910210,"TERMINAL",0,0,"",,terminal_output
599,3910408,"TERMINAL",0,0,"
",,terminal_output
600,3910487,"TERMINAL",0,0,"",,terminal_output
601,3910611,"TERMINAL",0,0,"",,terminal_output
602,3910817,"TERMINAL",0,0,"",,terminal_output
603,3911018,"TERMINAL",0,0,"",,terminal_output
604,3911536,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
605,3911866,"TERMINAL",0,0,"l /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
s /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
606,3911981,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
607,3912064,"TERMINAL",0,0,"- /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
608,3912211,"TERMINAL",0,0,"l /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
609,3912375,"TERMINAL",0,0,"h /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/
",,terminal_output
610,3912503,"TERMINAL",0,0,"",,terminal_output
611,3912611,"TERMINAL",0,0,"

total 1.7M
drwxr-x---+ 2 tum_dbd0378 hk-project-p0023960 4.0K Jul 16 21:06 big-runs
drwxr-x---+ 2 tum_dbd0378 hk-project-p0023960 4.0K Jul 15 12:12 dynamics-cotraining-modelsize-scaling
drwxr-x---+ 2 tum_dbd0378 hk-project-p0023960 4.0K Jul 15 11:46 tokenizer-lr-scaling
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 616K Jul 16 21:12 train_dynamics_lr_1e-4_larger_ffn_3352437.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 200K Jul 17 12:37 train_dynamics_lr_1e-4_larger_ffn_3353271.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 167K Jul  6 22:11 train_tokenizer_batch_size_scaling_linear_lr_scaling_16_nodes_3322407.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 243K Jul  7 10:23 train_tokenizer_batch_size_scaling_linear_lr_scaling_16_nodes_3325607.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  24K Jul  6 21:33 train_tokenizer_batch_size_scaling_linear_lr_scaling_2_nodes_3322404.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  34K Jul  7 10:22 train_tokenizer_batch_size_scaling_linear_lr_scaling_2_nodes_3325603.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  44K Jul  6 21:33 train_tokenizer_batch_size_scaling_linear_lr_scaling_4_nodes_3322405.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  64K Jul  7 10:22 train_tokenizer_batch_size_scaling_linear_lr_scaling_4_nodes_3325605.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  85K Jul  6 21:36 train_tokenizer_batch_size_scaling_linear_lr_scaling_8_nodes_3322406.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960 124K Jul  7 10:22 train_tokenizer_batch_size_scaling_linear_lr_scaling_8_nodes_3325606.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  14K Jul  6 21:33 train_tokenizer_no_warmup_ablation_3322446.log
-rw-rw----+ 1 tum_dbd0378 hk-project-p0023960  19K Jul  7 10:23 train_tokenizer_no_warmup_ablation_3325608.log
[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
612,3948927,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1213,0,"",shellscript,selection_mouse
613,3948930,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1212,0,"",shellscript,selection_command
614,3951626,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",403,0,"",shellscript,selection_mouse
615,3951747,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",386,20,"train_dynamics_lr_1e",shellscript,selection_mouse
616,3997738,"TERMINAL",0,0,"les",,terminal_output
617,3997818,"TERMINAL",0,0,"s",,terminal_output
618,3997896,"TERMINAL",0,0," ",,terminal_output
619,3998786,"TERMINAL",0,0,"s -lh /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/",,terminal_output
620,3999311,"TERMINAL",0,0,"
ess /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log ",,terminal_output
621,3999723,"TERMINAL",0,0,"",,terminal_output
622,4001210,"TERMINAL",0,0,"

",,terminal_output
623,4001297,"TERMINAL",0,0,"",,terminal_output
624,4001736,"TERMINAL",0,0,"",,terminal_output
625,4001932,"TERMINAL",0,0,"",,terminal_output
626,4002262,"TERMINAL",0,0,"",,terminal_output
627,4002457,"TERMINAL",0,0,"",,terminal_output
628,4002686,"TERMINAL",0,0,"",,terminal_output
629,4002866,"TERMINAL",0,0,"",,terminal_output
630,4003063,"TERMINAL",0,0,"",,terminal_output
631,4003175,"TERMINAL",0,0,"",,terminal_output
632,4003391,"TERMINAL",0,0,"",,terminal_output
633,4003743,"TERMINAL",0,0,"


",,terminal_output
634,4003932,"TERMINAL",0,0,"",,terminal_output
635,4004450,"TERMINAL",0,0,"/",,terminal_output
636,4004761,"TERMINAL",0,0,"train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
637,4005221,"TERMINAL",0,0,"train_dynamics_lr_1e-4_larger_ffn_3353271.log

=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2771207
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0516
SLURM_JOB_START_TIME=1752748479
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752834879
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_JOBID=3353271
SLURM_JOB_QOS=normal
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e10.hkn0516
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SCRIPT_CONTEXT=prolog_task
SLURM_NODELIST=hkn[0516,0518]
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn1990.localdomain
SLURM_JOB_ID=3353271
SLURM_NODEID=0
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
638,4006127,"TERMINAL",0,0,"
SLURM_NTASKS_PER_NODE=4
SLURM_JOB_GID=502226
SLURM_JOB_NODELIST=hkn[0516,0518]
GpuFreq=control_disabled
GpuFreq=control_disabled
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c :",,terminal_output908,7635533,"genie.py",0,0,"",python,tab
909,7639083,"utils/nn.py",0,0,"import math
from typing import Tuple

from flax import linen as nn
import jax
import jax.numpy as jnp
import einops


class PositionalEncoding(nn.Module):
    """"""https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.html""""""

    d_model: int  # Hidden dimensionality of the input.
    max_len: int = 5000  # Maximum length of a sequence to expect.

    def setup(self):
        # Create matrix of [SeqLen, HiddenDim] representing the positional encoding for max_len inputs
        self.pe = jnp.zeros((self.max_len, self.d_model))
        position = jnp.arange(0, self.max_len, dtype=jnp.float32)[:, None]
        div_term = jnp.exp(
            jnp.arange(0, self.d_model, 2) * (-math.log(10000.0) / self.d_model)
        )
        self.pe = self.pe.at[:, 0::2].set(jnp.sin(position * div_term))
        self.pe = self.pe.at[:, 1::2].set(jnp.cos(position * div_term))

    def __call__(self, x):
        x = x + self.pe[: x.shape[2]]
        return x


class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x


class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)


def normalize(x):
    return x / (jnp.linalg.norm(x, ord=2, axis=-1, keepdims=True) + 1e-8)


class VectorQuantizer(nn.Module):
    latent_dim: int
    num_latents: int
    dropout: float

    def setup(self):
        self.codebook = normalize(
            self.param(
                ""codebook"",
                nn.initializers.lecun_uniform(),
                (self.num_latents, self.latent_dim),
            )
        )
        self.drop = nn.Dropout(self.dropout, deterministic=False)

    def __call__(
        self, x: jax.Array, training: bool
    ) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
        # --- Compute distances ---
        x = normalize(x)
        codebook = normalize(self.codebook)
        distance = -jnp.matmul(x, codebook.T)
        if training:
            dropout_key = self.make_rng(""dropout"")
            distance = self.drop(distance, rng=dropout_key)

        # --- Get indices and embeddings ---
        indices = jnp.argmin(distance, axis=-1)
        z = self.codebook[indices]

        # --- Straight through estimator ---
        z_q = x + jax.lax.stop_gradient(z - x)
        return z_q, z, x, indices

    def get_codes(self, indices: jax.Array):
        return self.codebook[indices]


def _create_flash_attention_fn(use_flash_attention: bool, is_causal: bool):
    """"""
    Create an attention function that uses flash attention if enabled.

    Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
    jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim).

    We need to reshape to ensure compatibility. cuDNN's flash attention additionally
    requires a sequence length that is a multiple of 4. We pad the sequence length to the nearest
    multiple of 4 and mask accordingly.
    """"""
        
    def attention_fn(query, key, value, bias=None, mask=None, **kwargs):
        implementation = 'cudnn' if use_flash_attention else None

        def _rearrange(x):
            return einops.rearrange(x, '... l h d -> (...) l h d')
        def _pad(x):
            return jnp.pad(x, ((0, 0), (0, pad_size), (0, 0), (0, 0)))
        def _fuse_masks(mask: jax.Array, attention_mask: jax.Array) -> jax.Array:
            mask_bool = mask.astype(jnp.bool_)
            expanded_mask = jnp.pad(mask_bool, ((0, pad_size), (0, pad_size)), constant_values=False)
            return jnp.logical_and(attention_mask, expanded_mask)
        
        original_shape = query.shape
        original_seq_len = query.shape[-3]
        
        # Pad to nearest multiple of 4
        target_seq_len = ((original_seq_len + 3) // 4) * 4
        pad_size = target_seq_len - original_seq_len

        query_4d = _pad(_rearrange(query))
        key_4d = _pad(_rearrange(key))
        value_4d = _pad(_rearrange(value))
        
        attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        mask_4d = _fuse_masks(mask, attention_mask) if mask is not None else attention_mask
        mask_4d = mask_4d[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
        
        bias_4d = _pad(_rearrange(bias)) if bias is not None else None
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
            implementation=implementation,
            is_causal=is_causal,
            **kwargs
        )
        return output_4d[..., :original_seq_len, :, :].reshape(original_shape)
    
    return attention_fn

",python,tab
910,7640689,"utils/nn.py",3171,0,"",python,selection_command
911,7642306,"utils/nn.py",3203,0,"",python,selection_command
912,7642534,"utils/nn.py",3222,0,"",python,selection_command
913,7642564,"utils/nn.py",3239,0,"",python,selection_command
914,7642711,"utils/nn.py",3222,0,"",python,selection_command
915,7642883,"utils/nn.py",3203,0,"",python,selection_command
916,7643046,"utils/nn.py",3171,0,"",python,selection_command
917,7662116,"utils/nn.py",3165,31,"class STTransformer(nn.Module):",python,selection_command
918,7662229,"utils/nn.py",3165,221,"class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
",python,selection_command
919,7662424,"utils/nn.py",3165,1313,"class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)
",python,selection_command
920,7662585,"utils/nn.py",3165,1407,"class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)


def normalize(x):
    return x / (jnp.linalg.norm(x, ord=2, axis=-1, keepdims=True) + 1e-8)
",python,selection_command
921,7663747,"utils/nn.py",4572,0,"",python,selection_command
922,7663862,"utils/nn.py",4498,0,"",python,selection_command
923,7664151,"utils/nn.py",4480,0,"",python,selection_command
924,7664155,"utils/nn.py",4479,0,"",python,selection_command
925,7664160,"utils/nn.py",4478,0,"",python,selection_command
926,7664186,"utils/nn.py",4448,0,"",python,selection_command
927,7664222,"utils/nn.py",4435,0,"",python,selection_command
928,7664256,"utils/nn.py",4405,0,"",python,selection_command
929,7664289,"utils/nn.py",4363,0,"",python,selection_command
930,7664321,"utils/nn.py",4337,0,"",python,selection_command
931,7664363,"utils/nn.py",4315,0,"",python,selection_command
932,7664472,"utils/nn.py",4298,0,"",python,selection_command
933,7664475,"utils/nn.py",4236,0,"",python,selection_command
934,7664477,"utils/nn.py",4202,0,"",python,selection_command
935,7664488,"utils/nn.py",4156,0,"",python,selection_command
936,7664521,"utils/nn.py",4118,0,"",python,selection_command
937,7664556,"utils/nn.py",4076,0,"",python,selection_command
938,7664667,"utils/nn.py",4038,0,"",python,selection_command
939,7664732,"utils/nn.py",4002,0,"",python,selection_command
940,7664899,"utils/nn.py",3977,0,"",python,selection_command
941,7665049,"utils/nn.py",3989,0,"",python,selection_command
942,7665203,"utils/nn.py",3991,0,"",python,selection_command
943,7665371,"utils/nn.py",3993,0,"",python,selection_command
944,7666316,"utils/nn.py",1018,0,"",python,selection_command
945,7667134,"utils/nn.py",1012,25,"class STBlock(nn.Module):",python,selection_command
946,7667294,"utils/nn.py",1012,172,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
",python,selection_command
947,7667401,"utils/nn.py",1012,791,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z
",python,selection_command
948,7667617,"utils/nn.py",1012,1548,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)
",python,selection_command
949,7667691,"utils/nn.py",1012,2133,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z
",python,selection_command
950,7667841,"utils/nn.py",1012,2151,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x
",python,selection_command
951,7668704,"utils/nn.py",1012,2150,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
952,7668951,"utils/nn.py",3146,0,"",python,selection_command
953,7686454,"utils/nn.py",3145,0,"",python,selection_command
954,7686709,"utils/nn.py",3127,0,"",python,selection_command
955,7686733,"utils/nn.py",3114,0,"",python,selection_command
956,7686765,"utils/nn.py",3084,0,"",python,selection_command
957,7686815,"utils/nn.py",3042,0,"",python,selection_command
958,7686849,"utils/nn.py",3020,0,"",python,selection_command
959,7686866,"utils/nn.py",2998,0,"",python,selection_command
960,7686900,"utils/nn.py",2975,0,"",python,selection_command
961,7686933,"utils/nn.py",2962,0,"",python,selection_command
962,7686966,"utils/nn.py",2932,0,"",python,selection_command
963,7687001,"utils/nn.py",2890,0,"",python,selection_command
964,7687034,"utils/nn.py",2864,0,"",python,selection_command
965,7687068,"utils/nn.py",2842,0,"",python,selection_command
966,7687103,"utils/nn.py",2702,0,"",python,selection_command
967,7688648,"utils/nn.py",2702,140,"",python,content
968,7688661,"utils/nn.py",2710,0,"",python,selection_command
969,7696009,"utils/nn.py",0,0,"",python,tab
970,7696084,"utils/nn.py",2702,0,"",python,selection_command
971,7697497,"utils/nn.py",0,0,"import math
from typing import Tuple

from flax import linen as nn
import jax
import jax.numpy as jnp
import einops


class PositionalEncoding(nn.Module):
    """"""https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.html""""""

    d_model: int  # Hidden dimensionality of the input.
    max_len: int = 5000  # Maximum length of a sequence to expect.

    def setup(self):
        # Create matrix of [SeqLen, HiddenDim] representing the positional encoding for max_len inputs
        self.pe = jnp.zeros((self.max_len, self.d_model))
        position = jnp.arange(0, self.max_len, dtype=jnp.float32)[:, None]
        div_term = jnp.exp(
            jnp.arange(0, self.d_model, 2) * (-math.log(10000.0) / self.d_model)
        )
        self.pe = self.pe.at[:, 0::2].set(jnp.sin(position * div_term))
        self.pe = self.pe.at[:, 1::2].set(jnp.cos(position * div_term))

    def __call__(self, x):
        x = x + self.pe[: x.shape[2]]
        return x


class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x


class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)


def normalize(x):
    return x / (jnp.linalg.norm(x, ord=2, axis=-1, keepdims=True) + 1e-8)


class VectorQuantizer(nn.Module):
    latent_dim: int
    num_latents: int
    dropout: float

    def setup(self):
        self.codebook = normalize(
            self.param(
                ""codebook"",
                nn.initializers.lecun_uniform(),
                (self.num_latents, self.latent_dim),
            )
        )
        self.drop = nn.Dropout(self.dropout, deterministic=False)

    def __call__(
        self, x: jax.Array, training: bool
    ) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
        # --- Compute distances ---
        x = normalize(x)
        codebook = normalize(self.codebook)
        distance = -jnp.matmul(x, codebook.T)
        if training:
            dropout_key = self.make_rng(""dropout"")
            distance = self.drop(distance, rng=dropout_key)

        # --- Get indices and embeddings ---
        indices = jnp.argmin(distance, axis=-1)
        z = self.codebook[indices]

        # --- Straight through estimator ---
        z_q = x + jax.lax.stop_gradient(z - x)
        return z_q, z, x, indices

    def get_codes(self, indices: jax.Array):
        return self.codebook[indices]


def _create_flash_attention_fn(use_flash_attention: bool, is_causal: bool):
    """"""
    Create an attention function that uses flash attention if enabled.

    Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
    jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim).

    We need to reshape to ensure compatibility. cuDNN's flash attention additionally
    requires a sequence length that is a multiple of 4. We pad the sequence length to the nearest
    multiple of 4 and mask accordingly.
    """"""
        
    def attention_fn(query, key, value, bias=None, mask=None, **kwargs):
        implementation = 'cudnn' if use_flash_attention else None

        def _rearrange(x):
            return einops.rearrange(x, '... l h d -> (...) l h d')
        def _pad(x):
            return jnp.pad(x, ((0, 0), (0, pad_size), (0, 0), (0, 0)))
        def _fuse_masks(mask: jax.Array, attention_mask: jax.Array) -> jax.Array:
            mask_bool = mask.astype(jnp.bool_)
            expanded_mask = jnp.pad(mask_bool, ((0, pad_size), (0, pad_size)), constant_values=False)
            return jnp.logical_and(attention_mask, expanded_mask)
        
        original_shape = query.shape
        original_seq_len = query.shape[-3]
        
        # Pad to nearest multiple of 4
        target_seq_len = ((original_seq_len + 3) // 4) * 4
        pad_size = target_seq_len - original_seq_len

        query_4d = _pad(_rearrange(query))
        key_4d = _pad(_rearrange(key))
        value_4d = _pad(_rearrange(value))
        
        attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        mask_4d = _fuse_masks(mask, attention_mask) if mask is not None else attention_mask
        mask_4d = mask_4d[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
        
        bias_4d = _pad(_rearrange(bias)) if bias is not None else None
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
            implementation=implementation,
            is_causal=is_causal,
            **kwargs
        )
        return output_4d[..., :original_seq_len, :, :].reshape(original_shape)
    
    return attention_fn

",python,tab
972,7697498,"utils/nn.py",3144,0,"",python,selection_mouse
973,7697502,"utils/nn.py",3143,0,"",python,selection_command
974,7697904,"utils/nn.py",0,0,"",python,tab
975,7697911,"utils/nn.py",0,0,"",python,tab
976,7737916,"utils/nn.py",2732,0,"",python,selection_command
977,7738168,"utils/nn.py",2758,0,"",python,selection_command
978,7738190,"utils/nn.py",2800,0,"",python,selection_command
979,7738216,"utils/nn.py",2830,0,"",python,selection_command
980,7738268,"utils/nn.py",2843,0,"",python,selection_command
981,7738291,"utils/nn.py",2866,0,"",python,selection_command
982,7738395,"utils/nn.py",2888,0,"",python,selection_command
983,7738398,"utils/nn.py",2910,0,"",python,selection_command
984,7738402,"utils/nn.py",2952,0,"",python,selection_command
985,7738424,"utils/nn.py",2982,0,"",python,selection_command
986,7738591,"utils/nn.py",2995,0,"",python,selection_command
987,7738695,"utils/nn.py",3005,0,"",python,selection_command
988,7738920,"utils/nn.py",3014,0,"",python,selection_command
989,7739115,"utils/nn.py",3006,16,"        return x",python,selection_command
990,7739191,"utils/nn.py",3005,17,"
        return x",python,selection_command
991,7739642,"utils/nn.py",2560,462,"
        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
992,7739716,"utils/nn.py",1803,1219,"
        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
993,7739846,"utils/nn.py",1184,1838,"
    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
994,7740011,"utils/nn.py",1011,2011,"
class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
995,7740304,"utils/nn.py",1011,2011,"class STBlock(nn.Module): dim: int ffn_dim: int num_heads: int dropout: float param_dtype: jnp.dtype dtype: jnp.dtype use_flash_attention: bool @nn.remat @nn.compact def __call__(self, x: jax.Array) -> jax.Array: # --- Spatial attention --- z = PositionalEncoding(self.dim)(x) z = nn.LayerNorm( param_dtype=self.param_dtype, dtype=self.dtype,)(z) z = nn.MultiHeadAttention( num_heads=self.num_heads, qkv_features=self.dim, dropout_rate=self.dropout, param_dtype=self.param_dtype, dtype=self.dtype, attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),)(z) x = x + z # --- Temporal attention --- x = x.swapaxes(1, 2) z = PositionalEncoding(self.dim)(x) z = nn.LayerNorm( param_dtype=self.param_dtype, dtype=self.dtype,)(z) causal_mask = jnp.tri(z.shape[-2]) z = nn.MultiHeadAttention( num_heads=self.num_heads, qkv_features=self.dim, dropout_rate=self.dropout, param_dtype=self.param_dtype, dtype=self.dtype, attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True), # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True)(z, mask=causal_mask) x = x + z x = x.swapaxes(1, 2) # --- Feedforward --- z = nn.LayerNorm( param_dtype=self.param_dtype, dtype=self.dtype,)(x) z = nn.Dense( self.ffn_dim, param_dtype=self.param_dtype, dtype=self.dtype,)(z) z = nn.gelu(z) z = nn.Dense( self.dim, param_dtype=self.param_dtype, dtype=self.dtype,)(z) x = x + z return x",python,content
996,7740330,"utils/nn.py",2449,0,"",python,selection_command
997,7740510,"utils/nn.py",2459,0,"",python,selection_command
998,7742787,"utils/nn.py",2449,0,"

       ",python,content
999,7742792,"utils/nn.py",2439,0,"
       ",python,content
1000,7742797,"utils/nn.py",2435,0,"
        ",python,content
1001,7742800,"utils/nn.py",2417,0,"
           ",python,content
1002,7742803,"utils/nn.py",2387,0,"
           ",python,content
1003,7742806,"utils/nn.py",2377,0,"
           ",python,content
1004,7742809,"utils/nn.py",2363,0,"
       ",python,content
1005,7742813,"utils/nn.py",2348,0,"
       ",python,content
1006,7742816,"utils/nn.py",2344,0,"
        ",python,content
1007,7742820,"utils/nn.py",2326,0,"
           ",python,content
1008,7742823,"utils/nn.py",2296,0,"
           ",python,content
1009,7742828,"utils/nn.py",2282,0,"
           ",python,content
1010,7742831,"utils/nn.py",2268,0,"
       ",python,content
1011,7742833,"utils/nn.py",2264,0,"
        ",python,content
1012,7742835,"utils/nn.py",2246,0,"
           ",python,content
1013,7742838,"utils/nn.py",2216,0,"
           ",python,content
1014,7742841,"utils/nn.py",2198,0,"
       ",python,content
1015,7742844,"utils/nn.py",2176,0,"

       ",python,content
1016,7742846,"utils/nn.py",2155,0,"
       ",python,content
1017,7742849,"utils/nn.py",2145,0,"
       ",python,content
1018,7742853,"utils/nn.py",2123,0,"
        ",python,content
1019,7742856,"utils/nn.py",2031,0,"
       ",python,content
1020,7742859,"utils/nn.py",1948,0,"
           ",python,content
1021,7742862,"utils/nn.py",1930,0,"
           ",python,content
1022,7742865,"utils/nn.py",1900,0,"
           ",python,content
1023,7742868,"utils/nn.py",1873,0,"
           ",python,content
1024,7742871,"utils/nn.py",1850,0,"
           ",python,content
1025,7742874,"utils/nn.py",1824,0,"
           ",python,content
1026,7742877,"utils/nn.py",1797,0,"
       ",python,content
1027,7742879,"utils/nn.py",1762,0,"
       ",python,content
1028,7742882,"utils/nn.py",1758,0,"
        ",python,content
1029,7742885,"utils/nn.py",1740,0,"
           ",python,content
1030,7742888,"utils/nn.py",1710,0,"
           ",python,content
1031,7742891,"utils/nn.py",1692,0,"
       ",python,content
1032,7742893,"utils/nn.py",1656,0,"
       ",python,content
1033,7742896,"utils/nn.py",1635,0,"
       ",python,content
1034,7742899,"utils/nn.py",1606,0,"

       ",python,content
1035,7742902,"utils/nn.py",1596,0,"
       ",python,content
1036,7742904,"utils/nn.py",1592,0,"
        ",python,content
1037,7742907,"utils/nn.py",1508,0,"
           ",python,content
1038,7742910,"utils/nn.py",1490,0,"
           ",python,content
1039,7742912,"utils/nn.py",1460,0,"
           ",python,content
1040,7742916,"utils/nn.py",1433,0,"
           ",python,content
1041,7742918,"utils/nn.py",1410,0,"
           ",python,content
1042,7742921,"utils/nn.py",1384,0,"
           ",python,content
1043,7742923,"utils/nn.py",1357,0,"
       ",python,content
1044,7742926,"utils/nn.py",1353,0,"
        ",python,content
1045,7742929,"utils/nn.py",1335,0,"
           ",python,content
1046,7742932,"utils/nn.py",1305,0,"
           ",python,content
1047,7742934,"utils/nn.py",1287,0,"
       ",python,content
1048,7742938,"utils/nn.py",1251,0,"
       ",python,content
1049,7742940,"utils/nn.py",1223,0,"
       ",python,content
1050,7742943,"utils/nn.py",1176,0,"
   ",python,content
1051,7742946,"utils/nn.py",1164,0,"
   ",python,content
1052,7742950,"utils/nn.py",1154,0,"

   ",python,content
1053,7742952,"utils/nn.py",1128,0,"
   ",python,content
1054,7742955,"utils/nn.py",1111,0,"
   ",python,content
1055,7742958,"utils/nn.py",1088,0,"
   ",python,content
1056,7742961,"utils/nn.py",1073,0,"
   ",python,content
1057,7742964,"utils/nn.py",1058,0,"
   ",python,content
1058,7742967,"utils/nn.py",1045,0,"
   ",python,content
1059,7742970,"utils/nn.py",1036,0,"
   ",python,content
1060,7742974,"utils/nn.py",1011,0,"
",python,content
1061,7742978,"utils/nn.py",1011,0,"",python,selection_command
1062,7744047,"utils/nn.py",1012,0,"",python,selection_command
1063,7744246,"utils/nn.py",1012,25,"class STBlock(nn.Module):",python,selection_command
1064,7744359,"utils/nn.py",1012,172,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
",python,selection_command
1065,7744782,"utils/nn.py",1012,791,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z
",python,selection_command
1066,7744934,"utils/nn.py",1012,1548,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)
",python,selection_command
1067,7745209,"utils/nn.py",1012,1993,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z
",python,selection_command
1068,7745409,"utils/nn.py",1012,2011,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x
",python,selection_command
1069,7745738,"utils/nn.py",1012,1993,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z
",python,selection_command
1070,7746259,"utils/nn.py",1012,2010,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
1071,7746455,"utils/nn.py",1012,2011,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x
",python,selection_command
1072,7746630,"utils/nn.py",1012,2010,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
1073,7746980,"utils/nn.py",3006,0,"",python,selection_command
1074,7763922,"utils/nn.py",3005,0,"",python,selection_command
1075,7764096,"utils/nn.py",2987,0,"",python,selection_command
1076,7764292,"utils/nn.py",2974,0,"",python,selection_command
1077,7764359,"utils/nn.py",2944,0,"",python,selection_command
1078,7764613,"utils/nn.py",2902,0,"",python,selection_command
1079,7764644,"utils/nn.py",2880,0,"",python,selection_command
1080,7764675,"utils/nn.py",2858,0,"",python,selection_command
1081,7781383,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
1082,7782337,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1180,0,"",shellscript,selection_command
1083,7782480,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1151,0,"",shellscript,selection_command
1084,7782918,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1125,0,"",shellscript,selection_command
1085,7783046,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1096,0,"",shellscript,selection_command
1086,7783167,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1075,0,"",shellscript,selection_command
1087,7783363,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1047,0,"",shellscript,selection_command
1088,7809069,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1030,0,"",shellscript,selection_command
1089,7809191,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1029,0,"",shellscript,selection_command
1090,7809351,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1030,0,"",shellscript,selection_command
1091,7809539,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1047,0,"",shellscript,selection_command
1092,7814517,"train_dynamics.py",0,0,"",python,tab
1093,7815717,"train_dynamics.py",5372,0,"",python,selection_command
1094,7816081,"train_dynamics.py",5353,0,"",python,selection_command
1095,7816346,"train_dynamics.py",5355,0,"",python,selection_command
1096,7816791,"genie.py",0,0,"",python,tab
1097,7816792,"genie.py",334,0,"",python,selection_command
1098,7818031,"genie.py",385,0,"",python,selection_command
1099,7819233,"genie.py",419,0,"",python,selection_command
1100,7819393,"genie.py",442,0,"",python,selection_command
1101,7819546,"genie.py",542,0,"",python,selection_command
1102,7819680,"genie.py",572,0,"",python,selection_command
1103,7819933,"genie.py",1080,0,"",python,selection_command
1104,7819991,"genie.py",1092,0,"",python,selection_command
1105,7820015,"genie.py",1167,0,"",python,selection_command
1106,7820041,"genie.py",1207,0,"",python,selection_command
1107,7820499,"genie.py",1167,0,"",python,selection_command
1108,7820931,"genie.py",1092,0,"",python,selection_command
1109,7821493,"models/tokenizer.py",0,0,"",python,tab
1110,7821494,"models/tokenizer.py",196,0,"",python,selection_command
1111,7822534,"models/tokenizer.py",229,0,"",python,selection_command
1112,7822782,"models/tokenizer.py",249,0,"",python,selection_command
1113,7822811,"models/tokenizer.py",256,0,"",python,selection_command
1114,7822843,"models/tokenizer.py",272,0,"",python,selection_command
1115,7822867,"models/tokenizer.py",291,0,"",python,selection_command
1116,7822901,"models/tokenizer.py",308,0,"",python,selection_command
1117,7822935,"models/tokenizer.py",328,0,"",python,selection_command
1118,7822969,"models/tokenizer.py",349,0,"",python,selection_command
1119,7823003,"models/tokenizer.py",369,0,"",python,selection_command
1120,7823036,"models/tokenizer.py",389,0,"",python,selection_command
1121,7823070,"models/tokenizer.py",408,0,"",python,selection_command
1122,7823104,"models/tokenizer.py",427,0,"",python,selection_command
1123,7823140,"models/tokenizer.py",455,0,"",python,selection_command
1124,7823174,"models/tokenizer.py",482,0,"",python,selection_command
1125,7823207,"models/tokenizer.py",503,0,"",python,selection_command
1126,7823351,"models/tokenizer.py",527,0,"",python,selection_command
1127,7823516,"models/tokenizer.py",534,0,"",python,selection_command
1128,7823629,"models/tokenizer.py",555,0,"",python,selection_command
1129,7823880,"models/tokenizer.py",560,0,"",python,selection_command
1130,7824854,"models/tokenizer.py",561,0,"",python,selection_command
1131,7824994,"models/tokenizer.py",562,0,"",python,selection_command
1132,7825336,"models/tokenizer.py",570,0,"",python,selection_command
1133,7825538,"models/tokenizer.py",572,0,"",python,selection_command
1134,7826339,"models/tokenizer.py",584,0,"",python,selection_command
1135,7827865,"models/tokenizer.py",585,0,"",python,selection_command
1136,7828493,"models/tokenizer.py",585,0,"2",python,content
1137,7828494,"models/tokenizer.py",586,0,"",python,selection_keyboard
1138,7828640,"models/tokenizer.py",585,0,"",python,selection_command
1139,7829492,"models/tokenizer.py",614,0,"",python,selection_command
1140,7829496,"models/tokenizer.py",640,0,"",python,selection_command
1141,7829500,"models/tokenizer.py",669,0,"",python,selection_command
1142,7829505,"models/tokenizer.py",698,0,"",python,selection_command
1143,7829510,"models/tokenizer.py",726,0,"",python,selection_command
1144,7829512,"models/tokenizer.py",752,0,"",python,selection_command
1145,7829515,"models/tokenizer.py",782,0,"",python,selection_command
1146,7829518,"models/tokenizer.py",806,0,"",python,selection_command
1147,7829520,"models/tokenizer.py",844,0,"",python,selection_command
1148,7829524,"models/tokenizer.py",874,0,"",python,selection_command
1149,7829526,"models/tokenizer.py",909,0,"",python,selection_command
1150,7829527,"models/tokenizer.py",938,0,"",python,selection_command
1151,7829528,"models/tokenizer.py",968,0,"",python,selection_command
1152,7829533,"models/tokenizer.py",1003,0,"",python,selection_command
1153,7829818,"models/tokenizer.py",1013,0,"",python,selection_command
1154,7829924,"models/tokenizer.py",1051,0,"",python,selection_command
1155,7830145,"models/tokenizer.py",1107,0,"",python,selection_command
1156,7830960,"models/tokenizer.py",1107,0,"2",python,content
1157,7830962,"models/tokenizer.py",1108,0,"",python,selection_keyboard
1158,7831193,"models/tokenizer.py",1107,0,"",python,selection_command
1159,7831256,"models/tokenizer.py",1051,0,"",python,selection_command
1160,7831518,"models/tokenizer.py",1013,0,"",python,selection_command
1161,7831538,"models/tokenizer.py",1003,0,"",python,selection_command
1162,7831574,"models/tokenizer.py",968,0,"",python,selection_command
1163,7831608,"models/tokenizer.py",938,0,"",python,selection_command
1164,7831719,"models/tokenizer.py",909,0,"",python,selection_command
1165,7831722,"models/tokenizer.py",874,0,"",python,selection_command
1166,7831725,"models/tokenizer.py",844,0,"",python,selection_command
1167,7831741,"models/tokenizer.py",806,0,"",python,selection_command
1168,7831776,"models/tokenizer.py",782,0,"",python,selection_command
1169,7831914,"models/tokenizer.py",752,0,"",python,selection_command
1170,7831917,"models/tokenizer.py",726,0,"",python,selection_command
1171,7831919,"models/tokenizer.py",698,0,"",python,selection_command
1172,7831921,"models/tokenizer.py",669,0,"",python,selection_command
1173,7832024,"models/tokenizer.py",640,0,"",python,selection_command
1174,7832248,"models/tokenizer.py",614,0,"",python,selection_command
1175,7832333,"models/tokenizer.py",585,0,"",python,selection_command
1176,7834660,"utils/nn.py",0,0,"",python,tab
1177,7836695,"utils/nn.py",3031,0,"",python,selection_command
1178,7838629,"utils/nn.py",3025,31,"class STTransformer(nn.Module):",python,selection_command
1179,7838746,"utils/nn.py",3025,221,"class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
",python,selection_command
1180,7839108,"utils/nn.py",3025,1313,"class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)
",python,selection_command
1181,7839778,"utils/nn.py",3025,1312,"class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)",python,selection_command
1182,7840105,"utils/nn.py",3025,0,"",python,selection_command
1183,7840303,"utils/nn.py",3057,0,"",python,selection_command
1184,7840631,"utils/nn.py",3076,0,"",python,selection_command
1185,7840637,"utils/nn.py",3093,0,"",python,selection_command
1186,7840640,"utils/nn.py",3110,0,"",python,selection_command
1187,7840642,"utils/nn.py",3130,0,"",python,selection_command
1188,7840662,"utils/nn.py",3149,0,"",python,selection_command
1189,7840688,"utils/nn.py",3168,0,"",python,selection_command
1190,7840722,"utils/nn.py",3195,0,"",python,selection_command
1191,7840828,"utils/nn.py",3216,0,"",python,selection_command
1192,7840830,"utils/nn.py",3246,0,"",python,selection_command
1193,7840833,"utils/nn.py",3247,0,"",python,selection_command
1194,7840855,"utils/nn.py",3263,0,"",python,selection_command
1195,7840889,"utils/nn.py",3314,0,"",python,selection_command
1196,7840938,"utils/nn.py",3341,0,"",python,selection_command
1197,7840961,"utils/nn.py",3355,0,"",python,selection_command
1198,7840989,"utils/nn.py",3385,0,"",python,selection_command
1199,7841023,"utils/nn.py",3435,0,"",python,selection_command
1200,7841056,"utils/nn.py",3473,0,"",python,selection_command
1201,7841155,"utils/nn.py",3492,0,"",python,selection_command
1202,7841157,"utils/nn.py",3533,0,"",python,selection_command
1203,7841160,"utils/nn.py",3579,0,"",python,selection_command
1204,7841192,"utils/nn.py",3613,0,"",python,selection_command
1205,7841226,"utils/nn.py",3632,0,"",python,selection_command
1206,7841262,"utils/nn.py",3662,0,"",python,selection_command
1207,7841351,"utils/nn.py",3712,0,"",python,selection_command
1208,7841354,"utils/nn.py",3750,0,"",python,selection_command
1209,7841362,"utils/nn.py",3769,0,"",python,selection_command
1210,7841395,"utils/nn.py",3783,0,"",python,selection_command
1211,7841429,"utils/nn.py",3796,0,"",python,selection_command
1212,7841468,"utils/nn.py",3837,0,"",python,selection_command
1213,7841499,"utils/nn.py",3862,0,"",python,selection_command
1214,7841537,"utils/nn.py",3898,0,"",python,selection_command
1215,7841575,"utils/nn.py",3936,0,"",python,selection_command
1216,7841682,"utils/nn.py",3978,0,"",python,selection_command
1217,7841685,"utils/nn.py",4016,0,"",python,selection_command
1218,7841688,"utils/nn.py",4062,0,"",python,selection_command
1219,7841701,"utils/nn.py",4096,0,"",python,selection_command
1220,7841734,"utils/nn.py",4158,0,"",python,selection_command
1221,7841769,"utils/nn.py",4175,0,"",python,selection_command
1222,7841875,"utils/nn.py",4197,0,"",python,selection_command
1223,7841878,"utils/nn.py",4223,0,"",python,selection_command
1224,7841880,"utils/nn.py",4265,0,"",python,selection_command
1225,7841905,"utils/nn.py",4295,0,"",python,selection_command
1226,7841934,"utils/nn.py",4308,0,"",python,selection_command
1227,7842928,"utils/nn.py",4337,0,"
        ",python,content
1228,7842938,"utils/nn.py",4342,4,"",python,content
1229,7842990,"utils/nn.py",4338,4,"",python,content
1230,7843337,"utils/nn.py",4338,0,"
",python,content
1231,7843774,"utils/nn.py",4339,0,"
class STTransformer(nn.Module):
    model_dim: int
    ffn_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                ffn_dim=self.ffn_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)",python,content
1232,7843780,"utils/nn.py",4340,0,"",python,selection_command
1233,7845870,"utils/nn.py",4359,0,"2",python,content
1234,7845871,"utils/nn.py",4360,0,"",python,selection_command
1235,7847968,"utils/nn.py",4390,0,"",python,selection_command
1236,7848167,"utils/nn.py",4407,0,"",python,selection_command
1237,7848170,"utils/nn.py",4424,0,"",python,selection_command
1238,7848196,"utils/nn.py",4444,0,"",python,selection_command
1239,7848222,"utils/nn.py",4463,0,"",python,selection_command
1240,7848266,"utils/nn.py",4482,0,"",python,selection_command
1241,7848291,"utils/nn.py",4504,0,"",python,selection_command
1242,7848322,"utils/nn.py",4530,0,"",python,selection_command
1243,7848357,"utils/nn.py",4552,0,"",python,selection_command
1244,7848388,"utils/nn.py",4562,0,"",python,selection_command
1245,7848496,"utils/nn.py",4577,0,"",python,selection_command
1246,7848498,"utils/nn.py",4599,0,"",python,selection_command
1247,7848501,"utils/nn.py",4650,0,"",python,selection_command
1248,7848528,"utils/nn.py",4669,0,"",python,selection_command
1249,7848556,"utils/nn.py",4691,0,"",python,selection_command
1250,7848591,"utils/nn.py",4721,0,"",python,selection_command
1251,7848689,"utils/nn.py",4771,0,"",python,selection_command
1252,7848691,"utils/nn.py",4806,0,"",python,selection_command
1253,7848693,"utils/nn.py",4828,0,"",python,selection_command
1254,7848724,"utils/nn.py",4869,0,"",python,selection_command
1255,7848763,"utils/nn.py",4915,0,"",python,selection_command
1256,7848798,"utils/nn.py",4946,0,"",python,selection_command
1257,7848828,"utils/nn.py",4968,0,"",python,selection_command
1258,7848867,"utils/nn.py",4998,0,"",python,selection_command
1259,7848895,"utils/nn.py",5048,0,"",python,selection_command
1260,7849016,"utils/nn.py",5083,0,"",python,selection_command
1261,7849018,"utils/nn.py",5097,0,"",python,selection_command
1262,7849020,"utils/nn.py",5110,0,"",python,selection_command
1263,7849029,"utils/nn.py",5132,0,"",python,selection_command
1264,7849060,"utils/nn.py",5173,0,"",python,selection_command
1265,7849094,"utils/nn.py",5198,0,"",python,selection_command
1266,7849541,"utils/nn.py",5173,0,"",python,selection_command
1267,7849655,"utils/nn.py",5174,0,"",python,selection_command
1268,7849969,"utils/nn.py",5175,0,"",python,selection_command
1269,7850599,"utils/nn.py",5176,0,"2",python,content
1270,7850600,"utils/nn.py",5177,0,"",python,selection_command
1271,7854109,"utils/nn.py",1018,0,"",python,selection_command
1272,7855137,"utils/nn.py",1012,25,"class STBlock(nn.Module):",python,selection_command
1273,7855310,"utils/nn.py",1012,172,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
",python,selection_command
1274,7855626,"utils/nn.py",1012,791,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z
",python,selection_command
1275,7856031,"utils/nn.py",1012,1548,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)
",python,selection_command
1276,7856378,"utils/nn.py",1012,1993,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z
",python,selection_command
1277,7857311,"utils/nn.py",1012,2010,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
1278,7857459,"utils/nn.py",1012,2011,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x
",python,selection_command
1279,7857822,"utils/nn.py",1012,2010,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,selection_command
1280,7858188,"utils/nn.py",3006,0,"",python,selection_command
1281,7858478,"utils/nn.py",3023,0,"",python,selection_command
1282,7858745,"utils/nn.py",3006,0,"",python,selection_command
1283,7858998,"utils/nn.py",3022,0,"
        ",python,content
1284,7859248,"utils/nn.py",3031,0,"
        ",python,content
1285,7859249,"utils/nn.py",3023,8,"",python,content
1286,7859514,"utils/nn.py",3028,4,"",python,content
1287,7859876,"utils/nn.py",3024,4,"",python,content
1288,7860314,"utils/nn.py",3024,0,"class STBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x",python,content
1289,7860649,"utils/nn.py",5033,0,"",python,selection_command
1290,7861833,"utils/nn.py",3890,0,"",python,selection_command
1291,7862454,"utils/nn.py",3861,0,"",python,selection_command
1292,7862610,"utils/nn.py",3824,0,"",python,selection_command
1293,7862629,"utils/nn.py",3815,0,"",python,selection_command
1294,7862661,"utils/nn.py",3805,0,"",python,selection_command
1295,7862725,"utils/nn.py",3792,0,"",python,selection_command
1296,7862749,"utils/nn.py",3696,0,"",python,selection_command
1297,7862764,"utils/nn.py",3666,0,"",python,selection_command
1298,7862800,"utils/nn.py",3624,0,"",python,selection_command
1299,7862833,"utils/nn.py",3585,0,"",python,selection_command
1300,7862867,"utils/nn.py",3550,0,"",python,selection_command
1301,7862900,"utils/nn.py",3512,0,"",python,selection_command
1302,7862934,"utils/nn.py",3477,0,"",python,selection_command
1303,7862967,"utils/nn.py",3464,0,"",python,selection_command
1304,7863000,"utils/nn.py",3434,0,"",python,selection_command
1305,7863034,"utils/nn.py",3392,0,"",python,selection_command
1306,7863071,"utils/nn.py",3366,0,"",python,selection_command
1307,7863102,"utils/nn.py",3322,0,"",python,selection_command
1308,7863134,"utils/nn.py",3286,0,"",python,selection_command
1309,7863168,"utils/nn.py",3235,0,"",python,selection_command
1310,7863201,"utils/nn.py",3219,0,"",python,selection_command
1311,7863240,"utils/nn.py",3205,0,"",python,selection_command
1312,7863275,"utils/nn.py",3196,0,"",python,selection_command
1313,7863301,"utils/nn.py",3174,0,"",python,selection_command
1314,7863334,"utils/nn.py",3153,0,"",python,selection_command
1315,7863367,"utils/nn.py",3126,0,"",python,selection_command
1316,7863407,"utils/nn.py",3107,0,"",python,selection_command
1317,7863442,"utils/nn.py",3088,0,"",python,selection_command
1318,7863478,"utils/nn.py",3071,0,"",python,selection_command
1319,7863508,"utils/nn.py",3058,0,"",python,selection_command
1320,7863778,"utils/nn.py",3032,0,"",python,selection_command
1321,7864095,"utils/nn.py",3036,0,"",python,selection_command
1322,7868127,"utils/nn.py",3037,0,"2",python,content
1323,7868128,"utils/nn.py",3038,0,"",python,selection_command
1324,7891988,"utils/nn.py",3062,0,"",python,selection_command
1325,7892219,"utils/nn.py",3078,0,"",python,selection_command
1326,7892253,"utils/nn.py",3095,0,"",python,selection_command
1327,7892286,"utils/nn.py",3114,0,"",python,selection_command
1328,7892319,"utils/nn.py",3133,0,"",python,selection_command
1329,7892354,"utils/nn.py",3160,0,"",python,selection_command
1330,7892387,"utils/nn.py",3181,0,"",python,selection_command
1331,7892419,"utils/nn.py",3197,0,"",python,selection_command
1332,7892456,"utils/nn.py",3210,0,"",python,selection_command
1333,7892590,"utils/nn.py",3226,0,"",python,selection_command
1334,7892857,"utils/nn.py",3242,0,"",python,selection_command
1335,7892887,"utils/nn.py",3293,0,"",python,selection_command
1336,7892905,"utils/nn.py",3329,0,"",python,selection_command
1337,7892941,"utils/nn.py",3373,0,"",python,selection_command
1338,7892975,"utils/nn.py",3399,0,"",python,selection_command
1339,7893008,"utils/nn.py",3441,0,"",python,selection_command
1340,7893042,"utils/nn.py",3468,0,"",python,selection_command
1341,7893076,"utils/nn.py",3484,0,"",python,selection_command
1342,7893108,"utils/nn.py",3519,0,"",python,selection_command
1343,7893141,"utils/nn.py",3557,0,"",python,selection_command
1344,7893175,"utils/nn.py",3592,0,"",python,selection_command
1345,7893210,"utils/nn.py",3631,0,"",python,selection_command
1346,7893247,"utils/nn.py",3673,0,"",python,selection_command
1347,7893278,"utils/nn.py",3703,0,"",python,selection_command
1348,7893310,"utils/nn.py",3796,0,"",python,selection_command
1349,7893344,"utils/nn.py",3812,0,"",python,selection_command
1350,7893383,"utils/nn.py",3816,0,"",python,selection_command
1351,7893416,"utils/nn.py",3831,0,"",python,selection_command
1352,7893450,"utils/nn.py",3868,0,"",python,selection_command
1353,7893481,"utils/nn.py",3897,0,"",python,selection_command
1354,7893514,"utils/nn.py",3941,0,"",python,selection_command
1355,7893549,"utils/nn.py",3967,0,"",python,selection_command
1356,7893583,"utils/nn.py",4009,0,"",python,selection_command
1357,7893615,"utils/nn.py",4036,0,"",python,selection_command
1358,7893648,"utils/nn.py",4052,0,"",python,selection_command
1359,7893680,"utils/nn.py",4095,0,"",python,selection_command
1360,7894114,"utils/nn.py",4130,0,"",python,selection_command
1361,7894364,"utils/nn.py",4168,0,"",python,selection_command
1362,7894385,"utils/nn.py",4203,0,"",python,selection_command
1363,7894417,"utils/nn.py",4242,0,"",python,selection_command
1364,7894445,"utils/nn.py",4284,0,"",python,selection_command
1365,7894483,"utils/nn.py",4314,0,"",python,selection_command
1366,7894517,"utils/nn.py",4409,0,"",python,selection_command
1367,7894553,"utils/nn.py",4509,0,"",python,selection_command
1368,7894584,"utils/nn.py",4540,0,"",python,selection_command
1369,7894618,"utils/nn.py",4558,0,"",python,selection_command
1370,7894649,"utils/nn.py",4573,0,"",python,selection_command
1371,7894683,"utils/nn.py",4588,0,"",python,selection_command
1372,7894716,"utils/nn.py",4618,0,"",python,selection_command
1373,7894971,"utils/nn.py",4644,0,"",python,selection_command
1374,7895226,"utils/nn.py",4686,0,"",python,selection_command
1375,7895249,"utils/nn.py",4713,0,"",python,selection_command
1376,7895277,"utils/nn.py",4729,0,"",python,selection_command
1377,7895311,"utils/nn.py",4751,0,"",python,selection_command
1378,7895343,"utils/nn.py",4777,0,"",python,selection_command
1379,7895520,"utils/nn.py",4819,0,"",python,selection_command
1380,7895740,"utils/nn.py",4846,0,"",python,selection_command
1381,7897098,"utils/nn.py",4835,12,"        )(z)",python,selection_command
1382,7897251,"utils/nn.py",4805,42,"            dtype=self.dtype,
        )(z)",python,selection_command
1383,7897428,"utils/nn.py",4763,84,"            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)",python,selection_command
1384,7897581,"utils/nn.py",4737,110,"            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)",python,selection_command
1385,7897777,"utils/nn.py",4715,132,"        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)",python,selection_command
1386,7901388,"utils/nn.py",4726,0,"",python,selection_command
1387,7901612,"utils/nn.py",4748,0,"",python,selection_command
1388,7901860,"utils/nn.py",4774,0,"",python,selection_command
1389,7901883,"utils/nn.py",4816,0,"",python,selection_command
1390,7901913,"utils/nn.py",4846,0,"",python,selection_command
1391,7901945,"utils/nn.py",4859,0,"",python,selection_command
1392,7902215,"utils/nn.py",4882,0,"",python,selection_command
1393,7902399,"utils/nn.py",4871,21,"        z = nn.Dense(",python,selection_command
1394,7902510,"utils/nn.py",4871,43,"        z = nn.Dense(
            self.dim,",python,selection_command
1395,7902660,"utils/nn.py",4871,85,"        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,",python,selection_command
1396,7902800,"utils/nn.py",4871,115,"        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,",python,selection_command
1397,7902949,"utils/nn.py",4871,128,"        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)",python,selection_command
1398,7903225,"utils/nn.py",4871,129,"",python,content
1399,7903239,"utils/nn.py",4879,0,"",python,selection_command
1400,7908176,"experiments/train_dynamics_lr_1e-4_mixed_precision.sh",0,0,"#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/%x_%j.log
#SBATCH --job-name=train_dynamics_modelsize_scaling_36M_2_node_mixed_precision_flash_attention

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraining-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-mixed-precision-flash-attention-$slurm_job_id \
    --tags dynamics mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
",shellscript,tab
1401,7910999,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
1402,7912172,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1403,7912807,"TERMINAL",0,0,"[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1404,7912926,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
1405,7913426,"TERMINAL",0,0,"

",,terminal_output
1406,7920534,"TERMINAL",0,0,"
(reverse-i-search)`': ",,terminal_output
1407,7920797,"TERMINAL",0,0,"t': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
1408,7921007,"TERMINAL",0,0,"
r': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
1409,7921239,"TERMINAL",0,0,"
a': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log
i': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log
n': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
1410,7921960,"TERMINAL",0,0,"

 ",,terminal_output
1411,7922396,"TERMINAL",0,0,"
sbatch experiments/train_dynamics_lr_1e-4_larger_ffn.sh

",,terminal_output
1412,7922639,"TERMINAL",0,0,"tokenizer_lr_1e-4_larger_ffn.sh",,terminal_output
1413,7925106,"TERMINAL",0,0,"
[tum_dbd0378@hkn0819 jafar]$ sbatch experiments/train_tokenizer_lr_1e-4_larger_ffn.shqueue",,terminal_output
1414,7925641,"TERMINAL",0,0,"sbatch experiments/train_dynamics_lr_1e-4_larger_ffn.sh",,terminal_output
1415,7926164,"TERMINAL",0,0,"
queue",,terminal_output
1416,7926614,"TERMINAL",0,0,"sbatch experiments/train_dynamics_lr_1e-4_larger_ffn.sh",,terminal_output
1417,7927414,"TERMINAL",0,0,"",,terminal_output
1418,7927861,"TERMINAL",0,0,"",,terminal_output
1419,7927966,"TERMINAL",0,0,"",,terminal_output
1420,7928072,"TERMINAL",0,0,"rger_ffn.sh",,terminal_output
1421,7930504,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1048,0,"",shellscript,selection_command
1422,7930809,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1080,0,"",shellscript,selection_command
1423,7931005,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1098,0,"",shellscript,selection_command
1424,7931102,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",1131,0,"",shellscript,selection_command
1425,7935307,"utils/nn.py",0,0,"",python,tab
1426,7937207,"experiments/train_dynamics_lr_1e-4_larger_ffn.sh",0,0,"",shellscript,tab
1427,7937862,"utils/nn.py",0,0,"",python,tab
1428,7938546,"utils/nn.py",4856,0,"",python,selection_command
1429,7938871,"utils/nn.py",4843,0,"",python,selection_command
1430,7938873,"utils/nn.py",4813,0,"",python,selection_command
1431,7938875,"utils/nn.py",4771,0,"",python,selection_command
1432,7938895,"utils/nn.py",4745,0,"",python,selection_command
1433,7938929,"utils/nn.py",4723,0,"",python,selection_command
1434,7938962,"utils/nn.py",4710,0,"",python,selection_command
1435,7938996,"utils/nn.py",4680,0,"",python,selection_command
1436,7939029,"utils/nn.py",4638,0,"",python,selection_command
1437,7939063,"utils/nn.py",4612,0,"",python,selection_command
1438,7939095,"utils/nn.py",4582,0,"",python,selection_command
1439,7939199,"utils/nn.py",4573,0,"",python,selection_command
1440,7939201,"utils/nn.py",4552,0,"",python,selection_command
1441,7939204,"utils/nn.py",4534,0,"",python,selection_command
1442,7939231,"utils/nn.py",4503,0,"",python,selection_command
1443,7939264,"utils/nn.py",4403,0,"",python,selection_command
1444,7939311,"utils/nn.py",4308,0,"",python,selection_command
1445,7939394,"utils/nn.py",4278,0,"",python,selection_command
1446,7939396,"utils/nn.py",4236,0,"",python,selection_command
1447,7939403,"utils/nn.py",4197,0,"",python,selection_command
1448,7939438,"utils/nn.py",4162,0,"",python,selection_command
1449,7939471,"utils/nn.py",4124,0,"",python,selection_command
1450,7939502,"utils/nn.py",4089,0,"",python,selection_command
1451,7939536,"utils/nn.py",4046,0,"",python,selection_command
1452,7939573,"utils/nn.py",4033,0,"",python,selection_command
1453,7939602,"utils/nn.py",4003,0,"",python,selection_command
1454,7939722,"utils/nn.py",3961,0,"",python,selection_command
1455,7939724,"utils/nn.py",3935,0,"",python,selection_command
1456,7939754,"utils/nn.py",3891,0,"",python,selection_command
1457,7939759,"utils/nn.py",3862,0,"",python,selection_command
1458,7939768,"utils/nn.py",3825,0,"",python,selection_command
1459,7939802,"utils/nn.py",3816,0,"",python,selection_command
1460,7939919,"utils/nn.py",3806,0,"",python,selection_command
1461,7939922,"utils/nn.py",3793,0,"",python,selection_command
1462,7939925,"utils/nn.py",3697,0,"",python,selection_command
1463,7939937,"utils/nn.py",3667,0,"",python,selection_command
1464,7939974,"utils/nn.py",3625,0,"",python,selection_command
1465,7940003,"utils/nn.py",3586,0,"",python,selection_command
1466,7940038,"utils/nn.py",3551,0,"",python,selection_command
1467,7940072,"utils/nn.py",3513,0,"",python,selection_command
1468,7940103,"utils/nn.py",3478,0,"",python,selection_command
1469,7940136,"utils/nn.py",3465,0,"",python,selection_command
1470,7940247,"utils/nn.py",3435,0,"",python,selection_command
1471,7940249,"utils/nn.py",3393,0,"",python,selection_command
1472,7940251,"utils/nn.py",3367,0,"",python,selection_command
1473,7940269,"utils/nn.py",3323,0,"",python,selection_command
1474,7940302,"utils/nn.py",3287,0,"",python,selection_command
1475,7940336,"utils/nn.py",3236,0,"",python,selection_command
1476,7940452,"utils/nn.py",3220,0,"",python,selection_command
1477,7940456,"utils/nn.py",3206,0,"",python,selection_command
1478,7940459,"utils/nn.py",3197,0,"",python,selection_command
1479,7940469,"utils/nn.py",3175,0,"",python,selection_command
1480,7940505,"utils/nn.py",3154,0,"",python,selection_command
1481,7940537,"utils/nn.py",3127,0,"",python,selection_command
1482,7940570,"utils/nn.py",3108,0,"",python,selection_command
1483,7940604,"utils/nn.py",3089,0,"",python,selection_command
1484,7940638,"utils/nn.py",3072,0,"",python,selection_command
1485,7940800,"utils/nn.py",3059,0,"",python,selection_command
1486,7940979,"utils/nn.py",3032,0,"",python,selection_command
1487,7941567,"utils/nn.py",7053,0,"",python,selection_command
1488,7943492,"utils/nn.py",7012,0,"",python,selection_command
1489,7943744,"utils/nn.py",6994,0,"",python,selection_command
1490,7943775,"utils/nn.py",6981,0,"",python,selection_command
1491,7943796,"utils/nn.py",6966,0,"",python,selection_command
1492,7943927,"utils/nn.py",6928,0,"",python,selection_command
1493,7943929,"utils/nn.py",6878,0,"",python,selection_command
1494,7943931,"utils/nn.py",6848,0,"",python,selection_command
1495,7943934,"utils/nn.py",6829,0,"",python,selection_command
1496,7943968,"utils/nn.py",6795,0,"",python,selection_command
1497,7944002,"utils/nn.py",6749,0,"",python,selection_command
1498,7944115,"utils/nn.py",6708,0,"",python,selection_command
1499,7944118,"utils/nn.py",6689,0,"",python,selection_command
1500,7944122,"utils/nn.py",6651,0,"",python,selection_command
1501,7944133,"utils/nn.py",6601,0,"",python,selection_command
1502,7944173,"utils/nn.py",6571,0,"",python,selection_command
1503,7944204,"utils/nn.py",6553,0,"",python,selection_command
1504,7944234,"utils/nn.py",6530,0,"",python,selection_command
1505,7944267,"utils/nn.py",6479,0,"",python,selection_command
1506,7944309,"utils/nn.py",6461,0,"",python,selection_command
1507,7944341,"utils/nn.py",6446,0,"",python,selection_command
1508,7944440,"utils/nn.py",6432,0,"",python,selection_command
1509,7944443,"utils/nn.py",6411,0,"",python,selection_command
1510,7944447,"utils/nn.py",6384,0,"",python,selection_command
1511,7944468,"utils/nn.py",6365,0,"",python,selection_command
1512,7944518,"utils/nn.py",6346,0,"",python,selection_command
1513,7944637,"utils/nn.py",6365,0,"",python,selection_command
1514,7944809,"utils/nn.py",6346,0,"",python,selection_command
1515,7945062,"utils/nn.py",6326,0,"",python,selection_command
1516,7945164,"utils/nn.py",6308,0,"",python,selection_command
1517,7945169,"utils/nn.py",6291,0,"",python,selection_command
1518,7945172,"utils/nn.py",6273,0,"",python,selection_command
1519,7945192,"utils/nn.py",6240,0,"",python,selection_command
1520,7945223,"utils/nn.py",6223,0,"",python,selection_command
1521,7945258,"utils/nn.py",6222,0,"",python,selection_command
1522,7945290,"utils/nn.py",6208,0,"",python,selection_command
1523,7945323,"utils/nn.py",6190,0,"",python,selection_command
1524,7945358,"utils/nn.py",6165,0,"",python,selection_command
1525,7945390,"utils/nn.py",6123,0,"",python,selection_command
1526,7945718,"utils/nn.py",6165,0,"",python,selection_command
1527,7946013,"utils/nn.py",6190,0,"",python,selection_command
1528,7946016,"utils/nn.py",6208,0,"",python,selection_command
1529,7946030,"utils/nn.py",6222,0,"",python,selection_command
1530,7946064,"utils/nn.py",6223,0,"",python,selection_command
1531,7946118,"utils/nn.py",6240,0,"",python,selection_command
1532,7947783,"utils/nn.py",6273,0,"",python,selection_command
1533,7947947,"utils/nn.py",6291,0,"",python,selection_command
1534,7947980,"utils/nn.py",6308,0,"",python,selection_command
1535,7948014,"utils/nn.py",6326,0,"",python,selection_command
1536,7948115,"utils/nn.py",6346,0,"",python,selection_command
1537,7948119,"utils/nn.py",6365,0,"",python,selection_command
1538,7948122,"utils/nn.py",6384,0,"",python,selection_command
1539,7948142,"utils/nn.py",6411,0,"",python,selection_command
1540,7948177,"utils/nn.py",6432,0,"",python,selection_command
1541,7948211,"utils/nn.py",6446,0,"",python,selection_command
1542,7948308,"utils/nn.py",6461,0,"",python,selection_command
1543,7948311,"utils/nn.py",6479,0,"",python,selection_command
1544,7948481,"utils/nn.py",6530,0,"",python,selection_command
1545,7948652,"utils/nn.py",6553,0,"",python,selection_command
1546,7948834,"utils/nn.py",6571,0,"",python,selection_command
1547,7949071,"utils/nn.py",6601,0,"",python,selection_command
1548,7949161,"utils/nn.py",6651,0,"",python,selection_command
1549,7949163,"utils/nn.py",6689,0,"",python,selection_command
1550,7949167,"utils/nn.py",6708,0,"",python,selection_command
1551,7949191,"utils/nn.py",6749,0,"",python,selection_command
1552,7949222,"utils/nn.py",6795,0,"",python,selection_command
1553,7949257,"utils/nn.py",6829,0,"",python,selection_command
1554,7949359,"utils/nn.py",6848,0,"",python,selection_command
1555,7949362,"utils/nn.py",6878,0,"",python,selection_command
1556,7949368,"utils/nn.py",6928,0,"",python,selection_command
1557,7949395,"utils/nn.py",6966,0,"",python,selection_command
1558,7949426,"utils/nn.py",6981,0,"",python,selection_command
1559,7949460,"utils/nn.py",6994,0,"",python,selection_command
1560,7949493,"utils/nn.py",7012,0,"",python,selection_command
1561,7949882,"utils/nn.py",7053,0,"",python,selection_command
1562,7950047,"utils/nn.py",3030,0,"",python,selection_command
1563,7952366,"utils/nn.py",7053,0,"",python,selection_command
1564,7952913,"utils/nn.py",7012,0,"",python,selection_command
1565,7953174,"utils/nn.py",6994,0,"",python,selection_command
1566,7953212,"utils/nn.py",6981,0,"",python,selection_command
1567,7953227,"utils/nn.py",6966,0,"",python,selection_command
1568,7953260,"utils/nn.py",6928,0,"",python,selection_command
1569,7953356,"utils/nn.py",6878,0,"",python,selection_command
1570,7953358,"utils/nn.py",6848,0,"",python,selection_command
1571,7953361,"utils/nn.py",6829,0,"",python,selection_command
1572,7953391,"utils/nn.py",6795,0,"",python,selection_command
1573,7953421,"utils/nn.py",6749,0,"",python,selection_command
1574,7953456,"utils/nn.py",6708,0,"",python,selection_command
1575,7953493,"utils/nn.py",6689,0,"",python,selection_command
1576,7953523,"utils/nn.py",6651,0,"",python,selection_command
1577,7953558,"utils/nn.py",6601,0,"",python,selection_command
1578,7953593,"utils/nn.py",6571,0,"",python,selection_command
1579,7953625,"utils/nn.py",6553,0,"",python,selection_command
1580,7953658,"utils/nn.py",6530,0,"",python,selection_command
1581,7953691,"utils/nn.py",6479,0,"",python,selection_command
1582,7953727,"utils/nn.py",6461,0,"",python,selection_command
1583,7953760,"utils/nn.py",6446,0,"",python,selection_command
1584,7953794,"utils/nn.py",6432,0,"",python,selection_command
1585,7953884,"utils/nn.py",6411,0,"",python,selection_command
1586,7953889,"utils/nn.py",6384,0,"",python,selection_command
1587,7954123,"utils/nn.py",6365,0,"",python,selection_command
1588,7954453,"utils/nn.py",6346,0,"",python,selection_command
1589,7954455,"utils/nn.py",6326,0,"",python,selection_command
1590,7954459,"utils/nn.py",6308,0,"",python,selection_command
1591,7954464,"utils/nn.py",6291,0,"",python,selection_command
1592,7954497,"utils/nn.py",6273,0,"",python,selection_command
1593,7954532,"utils/nn.py",6240,0,"",python,selection_command
1594,7954566,"utils/nn.py",6223,0,"",python,selection_command
1595,7954599,"utils/nn.py",6222,0,"",python,selection_command
1596,7954632,"utils/nn.py",6208,0,"",python,selection_command
1597,7954792,"utils/nn.py",6222,0,"",python,selection_command
1598,7954967,"utils/nn.py",6223,0,"",python,selection_command
1599,7955132,"utils/nn.py",6240,0,"",python,selection_command
1600,7955408,"utils/nn.py",6230,0,"",python,selection_command
1601,7962646,"models/tokenizer.py",0,0,"",python,tab
1602,7965723,"models/tokenizer.py",547,0,"",python,selection_command
1603,7965966,"models/tokenizer.py",527,0,"",python,selection_command
1604,7965994,"models/tokenizer.py",525,0,"",python,selection_command
1605,7966027,"models/tokenizer.py",495,0,"",python,selection_command
1606,7966062,"models/tokenizer.py",474,0,"",python,selection_command
1607,7966096,"models/tokenizer.py",447,0,"",python,selection_command
1608,7966129,"models/tokenizer.py",419,0,"",python,selection_command
1609,7966162,"models/tokenizer.py",400,0,"",python,selection_command
1610,7966198,"models/tokenizer.py",381,0,"",python,selection_command
1611,7966230,"models/tokenizer.py",361,0,"",python,selection_command
1612,7966268,"models/tokenizer.py",341,0,"",python,selection_command
1613,7966300,"models/tokenizer.py",320,0,"",python,selection_command
1614,7966336,"models/tokenizer.py",300,0,"",python,selection_command
1615,7966399,"models/tokenizer.py",283,0,"",python,selection_command
1616,7966414,"models/tokenizer.py",264,0,"",python,selection_command
1617,7966446,"models/tokenizer.py",249,0,"",python,selection_command
1618,7966479,"models/tokenizer.py",247,0,"",python,selection_command
1619,7966515,"models/tokenizer.py",221,0,"",python,selection_command
1620,7966548,"models/tokenizer.py",189,0,"",python,selection_command
1621,7966586,"models/tokenizer.py",188,0,"",python,selection_command
1622,7966854,"models/tokenizer.py",172,0,"",python,selection_command
1623,7968020,"models/tokenizer.py",187,0,"",python,selection_command
1624,7968197,"models/tokenizer.py",187,0,",",python,content
1625,7968199,"models/tokenizer.py",188,0,"",python,selection_keyboard
1626,7968384,"models/tokenizer.py",188,0," ",python,content
1627,7968384,"models/tokenizer.py",189,0,"",python,selection_keyboard
1628,7969020,"models/tokenizer.py",189,0,"STTransformer2",python,content
1629,7969228,"models/tokenizer.py",202,0,"",python,selection_command
1630,7970290,"models/tokenizer.py",204,0,"",python,selection_command
1631,7970531,"models/tokenizer.py",205,0,"",python,selection_command
1632,7970557,"models/tokenizer.py",237,0,"",python,selection_command
1633,7970589,"models/tokenizer.py",263,0,"",python,selection_command
1634,7970622,"models/tokenizer.py",265,0,"",python,selection_command
1635,7970659,"models/tokenizer.py",280,0,"",python,selection_command
1636,7970690,"models/tokenizer.py",299,0,"",python,selection_command
1637,7970723,"models/tokenizer.py",316,0,"",python,selection_command
1638,7970755,"models/tokenizer.py",336,0,"",python,selection_command
1639,7970795,"models/tokenizer.py",357,0,"",python,selection_command
1640,7970821,"models/tokenizer.py",377,0,"",python,selection_command
1641,7970855,"models/tokenizer.py",397,0,"",python,selection_command
1642,7970889,"models/tokenizer.py",416,0,"",python,selection_command
1643,7970937,"models/tokenizer.py",435,0,"",python,selection_command
1644,7970992,"models/tokenizer.py",463,0,"",python,selection_command
1645,7970995,"models/tokenizer.py",490,0,"",python,selection_command
1646,7971022,"models/tokenizer.py",511,0,"",python,selection_command
1647,7971058,"models/tokenizer.py",541,0,"",python,selection_command
1648,7971088,"models/tokenizer.py",543,0,"",python,selection_command
1649,7971123,"models/tokenizer.py",563,0,"",python,selection_command
1650,7971681,"models/tokenizer.py",602,0,"",python,selection_command
1651,7971933,"models/tokenizer.py",588,0,"",python,selection_command
1652,7973069,"models/tokenizer.py",565,0,"",python,selection_command
1653,7975718,"TERMINAL",0,0,"

Submitted batch job 3353419
[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1654,7977022,"TERMINAL",0,0,"q",,terminal_output
1655,7977105,"TERMINAL",0,0,"u",,terminal_output
1656,7977180,"TERMINAL",0,0,"e",,terminal_output
1657,7977244,"TERMINAL",0,0,"u",,terminal_output
1658,7977358,"TERMINAL",0,0,"e",,terminal_output
1659,7977492,"TERMINAL",0,0,"

(BEvery 1.0s: squeue --mehkn0819.localdomain: Thu Jul 17 13:53:42 2025JOBID PARTITION     NAME     USER ST	TIME  NODES NODELIST(REASON)3353419 accelerat train_dy tum_dbd0  R	0:00	  2 hkn[0426,0431]3353272 accelerat interact tum_dbd0  R    1:19:03	  2 hkn[0819,0821]",,terminal_output
1660,7978514,"TERMINAL",0,0,"314",,terminal_output
1661,7979484,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1662,7980137,"TERMINAL",0,0,"i",,terminal_output
1663,7980193,"TERMINAL",0,0,"d",,terminal_output
1664,7980314,"TERMINAL",0,0,"l",,terminal_output
1665,7980378,"TERMINAL",0,0,"e",,terminal_output
1666,7980545,"TERMINAL",0,0,"

Partition dev_cpuonly             :     11 nodes idle
Partition cpuonly                 :      3 nodes idle
Partition dev_accelerated         :      2 nodes idle
Partition accelerated             :     56 nodes idle
Partition dev_accelerated-h100    :      0 nodes idle
Partition accelerated-h100        :      0 nodes idle
Partition large                   :      8 nodes idle
[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1667,7981855,"TERMINAL",0,0,"^C


[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1668,7984397,"TERMINAL",0,0,"q",,terminal_output
1669,7984508,"TERMINAL",0,0,"u",,terminal_output
1670,7984587,"TERMINAL",0,0,"e",,terminal_output
1671,7984668,"TERMINAL",0,0,"u",,terminal_output
1672,7984741,"TERMINAL",0,0,"e",,terminal_output
1673,7984866,"TERMINAL",0,0,"

(BEvery 1.0s: squeue --mehkn0819.localdomain: Thu Jul 17 13:53:50 2025JOBID PARTITION     NAME     USER ST	TIME  NODES NODELIST(REASON)3353419 accelerat train_dy tum_dbd0  R	0:08	  2 hkn[0426,0431]3353272 accelerat interact tum_dbd0  R    1:19:11	  2 hkn[0819,0821]",,terminal_output
1674,7985890,"TERMINAL",0,0,"192",,terminal_output
1675,7986622,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1676,8031768,"TERMINAL",0,0,"queue

(BEvery 1.0s: squeue --mehkn0819.localdomain: Thu Jul 17 13:54:37 2025JOBID PARTITION     NAME     USER ST	TIME  NODES NODELIST(REASON)3353419 accelerat train_dy tum_dbd0  R	0:55	  2 hkn[0426,0431]3353272 accelerat interact tum_dbd0  R    1:19:58	  2 hkn[0819,0821]",,terminal_output
1677,8032716,"TERMINAL",0,0,"869",,terminal_output
1678,8033718,"TERMINAL",0,0,"9720:00",,terminal_output
1679,8034737,"TERMINAL",0,0,"4081",,terminal_output
1680,8035754,"TERMINAL",0,0,"192",,terminal_output
1681,8036776,"TERMINAL",0,0,"21:003",,terminal_output
1682,8037831,"TERMINAL",0,0,"314",,terminal_output
1683,8038874,"TERMINAL",0,0,"425",,terminal_output
1684,8039865,"TERMINAL",0,0,"536",,terminal_output
1685,8040915,"TERMINAL",0,0,"647",,terminal_output
1686,8041860,"TERMINAL",0,0,"758",,terminal_output
1687,8042886,"TERMINAL",0,0,"869",,terminal_output
1688,8043893,"TERMINAL",0,0,"9710",,terminal_output
1689,8044909,"TERMINAL",0,0,"5081",,terminal_output
1690,8045934,"TERMINAL",0,0,"192",,terminal_output
1691,8046943,"TERMINAL",0,0,"2103",,terminal_output
1692,8048064,"TERMINAL",0,0,"314",,terminal_output
1693,8049090,"TERMINAL",0,0,"425",,terminal_output
1694,8050100,"TERMINAL",0,0,"536",,terminal_output
1695,8051119,"TERMINAL",0,0,"647",,terminal_output
1696,8052148,"TERMINAL",0,0,"758",,terminal_output
1697,8053154,"TERMINAL",0,0,"869",,terminal_output
1698,8054169,"TERMINAL",0,0,"9720",,terminal_output
1699,8055189,"TERMINAL",0,0,"5:0081",,terminal_output
1700,8056207,"TERMINAL",0,0,"192",,terminal_output
1701,8057224,"TERMINAL",0,0,"2203",,terminal_output
1702,8058245,"TERMINAL",0,0,"314",,terminal_output
1703,8059260,"TERMINAL",0,0,"425",,terminal_output
1704,8060274,"TERMINAL",0,0,"536",,terminal_output
1705,8061294,"TERMINAL",0,0,"647",,terminal_output
1706,8062315,"TERMINAL",0,0,"758",,terminal_output
1707,8063323,"TERMINAL",0,0,"869",,terminal_output
1708,8064347,"TERMINAL",0,0,"9730",,terminal_output
1709,8065366,"TERMINAL",0,0,"1081",,terminal_output
1710,8066380,"TERMINAL",0,0,"192",,terminal_output
1711,8067399,"TERMINAL",0,0,"2303",,terminal_output
1712,8068414,"TERMINAL",0,0,"314",,terminal_output
1713,8069425,"TERMINAL",0,0,"425",,terminal_output
1714,8069660,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1715,8070476,"TERMINAL",0,0,"
(reverse-i-search)`': ",,terminal_output
1716,8070719,"TERMINAL",0,0,"l': idlee': idle",,terminal_output
1717,8070771,"TERMINAL",0,0,"s': less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353271.log",,terminal_output
1718,8071325,"TERMINAL",0,0,"
[tum_dbd0378@hkn0819 jafar]$ less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/trai_dynamics_lr_1e-4_larger_ffn_3353271.log

",,terminal_output
1719,8071640,"TERMINAL",0,0,"",,terminal_output
1720,8072157,"TERMINAL",0,0,"",,terminal_output
1721,8072290,"TERMINAL",0,0,"",,terminal_output
1722,8072565,"TERMINAL",0,0,"3353419",,terminal_output
1723,8072768,"TERMINAL",0,0,"3353419.log ",,terminal_output
1724,8072928,"TERMINAL",0,0,"

=
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrec ords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraini ng-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr- scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
SLURM_STEP_NUM_TASKS=1
SLURM_JOB_USER=tum_dbd0378
SLURM_TASKS_PER_NODE=4(x2)
SLURM_JOB_UID=996262
SLURM_TASK_PID=2999816
SLURM_JOB_GPUS=0,1,2,3
SLURM_LOCALID=0
SLURM_SUBMIT_DIR=/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar
SLURMD_NODENAME=hkn0426
SLURM_JOB_START_TIME=1752753222
SLURM_STEP_NODELIST=hkn0819
SLURM_CLUSTER_NAME=hk
SLURM_JOB_END_TIME=1752839622
SLURM_PMI2_SRUN_PORT=43595
SLURM_CPUS_ON_NODE=24
SLURM_JOB_CPUS_PER_NODE=24(x2)
SLURM_GPUS_ON_NODE=4
SLURM_GTIDS=0
SLURM_JOB_PARTITION=accelerated
SLURM_TRES_PER_TASK=cpu=5
SLURM_OOM_KILL_STEP=0
SLURM_JOB_NUM_NODES=2
SLURM_STEPID=4294967290
SLURM_JOBID=3353419
SLURM_PTY_PORT=34217
SLURM_JOB_QOS=normal
SLURM_LAUNCH_NODE_IPADDR=10.0.7.198
SLURM_PTY_WIN_ROW=93
SLURM_PMI2_PROC_MAPPING=(vector,(0,1,1))
SLURMD_DEBUG=2
SLURM_PROCID=0
SLURM_CPUS_PER_TASK=5
SLURM_NTASKS=8
SLURM_TOPOLOGY_ADDR=hkibb.hkibbi1.hkibbi1e13.hkn0426
SLURM_TOPOLOGY_ADDR_PATTERN=switch.switch.switch.node
SLURM_SRUN_COMM_HOST=10.0.7.198
SLURM_SCRIPT_CONTEXT=prolog_task
/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353419.log",,terminal_output
1725,8073843,"TERMINAL",0,0,"
SLURM_PTY_WIN_COL=109
SLURM_NODELIST=hkn[0426,0431]
SLURM_SRUN_COMM_PORT=37059
SLURM_STEP_ID=4294967290
SLURM_JOB_ACCOUNT=hk-project-p0023960
SLURM_PRIO_PROCESS=0
SLURM_NPROCS=8
SLURM_NNODES=2
SLURM_SUBMIT_HOST=hkn0819.localdomain
SLURM_JOB_ID=3353419
SLURM_NODEID=0
SLURM_STEP_NUM_NODES=1
SLURM_STEP_TASKS_PER_NODE=1
SLURM_MPI_TYPE=pmi2
SLURM_PMI2_STEP_NODES=hkn0819
SLURM_CONF=/etc/slurm/slurm.conf
SLURM_JOB_NAME=train_dynamics_lr_1e-4_larger_ffn
SLURM_NTASKS_PER_NODE=4
SLURM_STEP_LAUNCHER_PORT=37059
SLURM_JOB_GID=502226
SLURM_JOB_NODELIST=hkn[0426,0431]
GpuFreq=control_disabled
GpuFreq=control_disabled
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
:",,terminal_output
1726,8074350,"TERMINAL",0,0,"
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default val ue `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to h andle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: :
347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py: 347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<c lass 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle  this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
Running on 8 devices.
2025-07-17 13:54:38.806808: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.810442: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.814469: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.814471: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.815413: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.820262: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.821425: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:54:38.821516: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:04.813982: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:05.002835: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
:
2025-07-17 13:55:05.067907: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:05.177993: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:07.060701: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:07.576120: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:07.747293: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:08.610838: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:11.448876: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:11.664369: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:11.840530: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:11.919849: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 17229792, 'dynamics': 26555904, 'total': 81775312}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 17229792, 'dynamics': 26555904, 'total': 81775312}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 17229792, 'dynamics': 26555904, 'total': 81775312}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 17229792, 'dynamics': 26555904, 'total': 81775312}
2025-07-17 13:55:15.100659: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:15.212729: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo :
d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:15.480954: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
2025-07-17 13:55:15.951489: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs we re filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a goo d representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
:

:

:

:

:

:",,terminal_output
1727,8074429,"TERMINAL",0,0,"

:

:",,terminal_output
1728,8075447,"TERMINAL",0,0,"
>[tum_dbd0378@hkn0819 jafar]$ ",,terminal_output
1729,8076155,"TERMINAL",0,0,"less /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dynamics_lr_1e-4_larger_ffn_3353419.log ",,terminal_output
1730,8076247,"TERMINAL",0,0,"",,terminal_output
1731,8076962,"TERMINAL",0,0,"",,terminal_output
1732,8077201,"TERMINAL",0,0,"",,terminal_output
1733,8077365,"TERMINAL",0,0,"",,terminal_output
1734,8077498,"TERMINAL",0,0,"",,terminal_output
1735,8077664,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dy",,terminal_output
1736,8077915,"TERMINAL",0,0,"t /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_dy",,terminal_output
1737,8078017,"TERMINAL",0,0,"a /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_d",,terminal_output
1738,8078159,"TERMINAL",0,0,"i /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train_l /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/train",,terminal_output
1739,8078298,"TERMINAL",0,0," /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/trai",,terminal_output
1740,8078448,"TERMINAL",0,0,"- /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tra",,terminal_output
1741,8078526,"TERMINAL",0,0,"f /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tr",,terminal_output
1742,8078700,"TERMINAL",0,0,"

",,terminal_output
1743,8078922,"TERMINAL",0,0,"

Parameter counts:
{'tokenizer': 37989616, 'lam': 17229792, 'dynamics': 26555904, 'total': 81775312}
wandb: Tracking run with wandb version 0.19.11
wandb: Run data is saved locally in /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/wandb/run-20250717_135518-427ls6he
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run dynamics-lr-1e-4-larger-ffn
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/427ls6he
Parameter counts:
{'tokenizer': 37989616, 'lam': 17229792, 'dynamics': 26555904, 'total': 81775312}
",,terminal_output
1744,8093566,"utils/nn.py",0,0,"",python,tab
1745,8093568,"utils/nn.py",1193,5,"remat",python,selection_command
1746,8095240,"utils/nn.py",3206,5,"remat",python,selection_command
1747,8096237,"utils/nn.py",1193,5,"remat",python,selection_command
1748,8097533,"utils/nn.py",1198,0,"",python,selection_mouse
1749,8097536,"utils/nn.py",1197,0,"",python,selection_command