[tab] temp.out (plaintext)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
2025-07-22 23:07:35.460903: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-22 23:07:41.388920: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-22 23:07:46.669343: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-22 23:07:54.175116: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/sample.py", line 125, 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/standard_checkpoint_handler.py", line 259, in restore
    return self._impl.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: {'params': {'params': {'dynamics': {'dynamics': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_10': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_11': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_8': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_9': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_10': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_11': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_8': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_9': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}, 'lam': {'encoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}, 'tokenizer': {'decoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}, 'encoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}}}, 'opt_state': [{'mu': {'params': {'dynamics': {'dynamics': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_10': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_11': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_8': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_9': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_10': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_11': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64)srun: error: hai006: task 0: Exited with exit code 1
)}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_8': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_9': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}, 'lam': {'encoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}, 'tokenizer': {'decoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}, 'encoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}}}, 'nu': {'params': {'dynamics': {'dynamics': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_10': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_11': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_8': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_9': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_10': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_11': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_8': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_9': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}, 'lam': {'encoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}, 'tokenizer': {'decoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}, 'encoder': {'st_block_0': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_1': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_2': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_3': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_4': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_5': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_6': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'st_block_7': Diff(lhs={'Dense_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_0': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_1': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'LayerNorm_2': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'scale': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'spatial_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}, 'temporal_attn': {'key': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'out': {'bias': ShapeDtypeStruct(shape=(512,), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(8, 64, 512), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'query': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}, 'value': {'bias': ShapeDtypeStruct(shape=(8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device)), 'kernel': ShapeDtypeStruct(shape=(512, 8, 64), dtype=float32, sharding=SingleDeviceSharding(device=CudaDevice(id=0), memory_kind=device))}}}, rhs=None), 'STBlock_0': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_1': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_2': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_3': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_4': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_5': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_6': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}}), 'STBlock_7': Diff(lhs=None, rhs={'Dense_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 512))}, 'LayerNorm_0': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_1': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'LayerNorm_2': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'scale': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,))}, 'MultiHeadAttention_0': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}, 'MultiHeadAttention_1': {'key': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'out': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64,)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64, 512))}, 'query': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}, 'value': {'bias': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(1, 64)), 'kernel': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(64, 8, 64))}}})}}}}}, None, None]}