[tab] temp.out (plaintext)
Running on 1 devices.
Counting all components: ['action_in', 'action_up', 'decoder', 'encoder', 'patch_up', 'vq']
Parameter counts:
{'action_in': 768, 'action_up': 16896, 'decoder': 17474816, 'encoder': 17228832, 'patch_up': 393728, 'vq': 192, 'total': 35115232}
Starting training from step 0...
2025-07-24 19:34:02.159619: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-24 19:34:02.161376: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-24 19:34:02.162216: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
Step 0, loss: 0.2764054834842682
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_lam.py", line 310, in <module>
    loss, recon, action_last_active, metrics = train_step(
  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_lam.py", line 119, 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: Mismatch custom node data: ('action_in', 'action_up', 'decoder', 'encoder', 'patch_up') != ('action_in', 'action_up', 'decoder', 'encoder', 'patch_up', 'vq'); value: State({
  '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 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '