{"tree_metadata": {"('model', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'dynamics', 'diffusion_transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'dynamics', 'diffusion_transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'dynamics', 'diffusion_transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'diffusion_transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'dynamics', 'diffusion_transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('model', 'dynamics', 'timestep_embed', 'embedding', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "timestep_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [25, 32]}}, "('model', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('model', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('model', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 768]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('model', 'tokenizer', 'mask_patch', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "mask_patch", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'dynamics', 'diffusion_transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('opt_state', '0', 'mu', 'dynamics', 'timestep_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "timestep_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [25, 32]}}, "('opt_state', '0', 'mu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'mask_patch', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "mask_patch", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'dynamics', 'diffusion_transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "diffusion_transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('opt_state', '0', 'nu', 'dynamics', 'timestep_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "timestep_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [25, 32]}}, "('opt_state', '0', 'nu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'mask_patch', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "mask_patch", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '2', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step', 'value')": {"key_metadata": [{"key": "step", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}