[tab] temp.out (plaintext)
Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/sample.py", line 141, in <module>
    restored = checkpoint_manager.restore(
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/checkpoint_manager.py", line 1647, in restore
    restored = self._checkpointer.restore(restore_directory, args=args)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/async_checkpointer.py", line 562, in restore
    return super().restore(directory, *args, **kwargs)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 304, in restore
    restored = self._restore(directory, args=ckpt_args)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 323, in _restore
    return self._handler.restore(directory, args=args)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py", line 859, in restore
    restored[item_name] = handler.restore(
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py", line 816, in restore
    return self._handler_impl.restore(directory, args=args)
  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py", line 792, in restore
    raise ValueError(
ValueError: User-provided restore item and on-disk value metadata tree structures do not match: {'model': {'dynamics': {'transformer': Diff(lhs={'blocks': {'0': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}}, '1': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}}, '2': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}}, '3': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}}, '4': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}}, '5': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_pos_enc': {'pe': {'value': ShapeDtypeStruct(shape=(5000, 512), dtype=float32)}}}}, 'input_dense': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 512), dtype=float32)}}, 'input_norm1': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'input_norm2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'output_dense': {'bias': {'value': ShapeDtypeStruct(shape=(1024,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 1024), dtype=float32)}}}, rhs=None), 'dynamics': Diff(lhs=None, rhs={'blocks': {'0': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}}, '1': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}}, '2': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}}, '3': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}}, '4': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}}, '5': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_pos_enc': {'pe': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(5000, 512))}}}}, 'input_dense': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 512))}}, 'input_norm1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'input_norm2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'output_dense': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1024,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 1024))}}})}}, 'opt_state': {'0': {'mu': {'dynamics': {'transformer': Diff(lhs={'blocks': {'0': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=bfloat16)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=bfloat16)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}}, '1': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=bfloat16)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=bfloat16)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}}, '2': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=bfloat16)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=bfloat16)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}}, '3': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=bfloat16)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=bfloat16)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}}, '4': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=bfloat16)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=bfloat16)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}}, '5': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=bfloat16)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=bfloat16)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=bfloat16)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=bfloat16)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}}}, 'input_dense': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 512), dtype=bfloat16)}}, 'input_norm1': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'input_norm2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=bfloat16)}}, 'output_dense': {'bias': {'value': ShapeDtypeStruct(shape=(1024,), dtype=bfloat16)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 1024), dtype=bfloat16)}}}, rhs=None), 'dynamics': Diff(lhs=None, rhs={'blocks': {'0': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '1': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '2': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '3': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '4': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '5': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}}, 'input_dense': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 512))}}, 'input_norm1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'input_norm2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'output_dense': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1024,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 1024))}}})}}, 'nu': {'dynamics': {'transformer': Diff(lhs={'blocks': {'0': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}}, '1': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}}, '2': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}}, '3': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}}, '4': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}}, '5': {'ffn_dense1': {'bias': {'value': ShapeDtypeStruct(shape=(2048,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 2048), dtype=float32)}}, 'ffn_dense2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(2048, 512), dtype=float32)}}, 'ffn_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'spatial_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'spatial_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'temporal_attention': {'key': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'out': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32)}}, 'query': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}, 'value': {'bias': {'value': ShapeDtypeStruct(shape=(8, 64), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32)}}}, 'temporal_norm': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}}}, 'input_dense': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 512), dtype=float32)}}, 'input_norm1': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'input_norm2': {'bias': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}, 'scale': {'value': ShapeDtypeStruct(shape=(512,), dtype=float32)}}, 'output_dense': {'bias': {'value': ShapeDtypeStruct(shape=(1024,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(512, 1024), dtype=float32)}}}, rhs=None), 'dynamics': Diff(lhs=None, rhs={'blocks': {'0': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '1': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '2': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '3': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '4': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}, '5': {'ffn_dense1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 2048))}}, 'ffn_dense2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(2048, 512))}}, 'ffn_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'spatial_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'spatial_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'temporal_attention': {'key': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'out': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64, 512))}}, 'query': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}, 'value': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(8, 64))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 8, 64))}}}, 'temporal_norm': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}}}, 'input_dense': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 512))}}, 'input_norm1': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'input_norm2': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}, 'scale': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512,))}}, 'output_dense': {'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1024,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(512, 1024))}}})}}}}}
srun: error: hai002: task 0: Exited with exit code 1