[tab] temp.out (plaintext)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1256: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.
  warnings.warn(
Running on 1 devices.
Counting all components: ['dynamics', 'lam', 'tokenizer']
Parameter counts:
{'dynamics': 26555904, 'lam': 35115232, 'tokenizer': 33750256, 'total': 95421392}
Starting training from step 0...
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 362, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 126, in __call__
    out = self.f(*args, **kwargs)
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 146, in train_step
    optimizer.update(grads)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/training/optimizer.py", line 271, in update
    updates, new_opt_state = self.tx.update(grads, opt_state, params, **kwargs)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/optax/transforms/_combining.py", line 89, in update_fn
    updates, new_s = fn(updates, s, params, **extra_args)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/optax/_src/base.py", line 335, in update
    return tx.update(updates, state, params)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/optax/_src/transform.py", line 283, in update_fn
    mu = optax.tree.update_moment(updates, state.mu, b1, 1)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/optax/tree_utils/_tree_math.py", line 332, in tree_update_moment
    return jax.tree.map(
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
ValueError: Dict key mismatch; expected keys: ['action_in', 'action_up', 'encoder', 'patch_up', 'vq']; dict: {'action_in': VariableState( # 768 (1.5 KB)
  type=Param,
  value=Traced<bfloat16[1,1,1,768]>with<DynamicJaxprTrace>
), 'action_up': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 16,384 (32.8 KB)
  type=Param,
  value=Traced<bfloat16[32,512]>with<DynamicJaxprTrace>
)}, 'decoder': {'blocks': {0: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}, 1: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}, 2: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}, 3: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}}, 'input_dense': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,512]>with<DynamicJaxprTrace>
)}, 'input_norm1': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'input_norm2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'output_dense': {'bias': VariableState( # 768 (1.5 KB)
  type=Param,
  value=Traced<bfloat16[768]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 393,216 (786.4 KB)
  type=Param,
  value=Traced<bfloat16[512,768]>with<DynamicJaxprTrace>
)}}, 'encoder': {'blocks': {0: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}, 1: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}, 2: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}, 3: {'ffn_dense1': {'bias': VariableState( # 2,048 (4.1 KB)
  type=Param,
  value=Traced<bfloat16[2048]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[512,2048]>with<DynamicJaxprTrace>
)}, 'ffn_dense2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 1,048,576 (2.1 MB)
  type=Param,
  value=Traced<bfloat16[2048,512]>with<DynamicJaxprTrace>
)}, 'ffn_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'spatial_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'spatial_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'temporal_attention': {'key': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'out': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[8,64,512]>with<DynamicJaxprTrace>
)}, 'query': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}, 'value': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[8,64]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 262,144 (524.3 KB)
  type=Param,
  value=Traced<bfloat16[512,8,64]>with<DynamicJaxprTrace>
)}}, 'temporal_norm': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}}}, 'input_dense': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 393,216 (786.4 KB)
  type=Param,
  value=Traced<bfloat16[768,512]>with<DynamicJaxprTrace>
)}, 'input_norm1': {'bias': VariableState( # 768 (1.5 KB)
  type=Param,
  value=Traced<bfloat16[768]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 768 (1.5 KB)
  type=Param,
  value=Traced<bfloat16[768]>with<DynamicJaxprTrace>
)}, 'input_norm2': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'scale': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
)}, 'output_dense': {'bias': VariableState( # 32 (64 B)
  type=Param,
  value=Traced<bfloat16[32]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 16,384 (32.8 KB)
  type=Param,
  value=Traced<bfloat16[512,32]>with<DynamicJaxprTrace>
)}}, 'patch_up': {'bias': VariableState( # 512 (1.0 KB)
  type=Param,
  value=Traced<bfloat16[512]>with<DynamicJaxprTrace>
), 'kernel': VariableState( # 393,216 (786.4 KB)
  type=Param,
  value=Traced<bfloat16[768,512]>with<DynamicJaxprTrace>
)}, 'vq': {'codebook': VariableState( # 192 (384 B)
  type=Param,
  value=Traced<bfloat16[6,32]>with<DynamicJaxprTrace>
)}}.
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 7 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
srun: error: hai002: task 0: Exited with exit code 1