checkpoint/default/_METADATA

{"tree_metadata": {"('dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'count')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'key')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 128]}}, "('encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 128]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('encoder', 'temporal_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 64]}}, "('encoder', 'type_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 64]}}, "('fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 64]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 64]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 8, 8]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 7]}}, "('policy_query',)": {"key_metadata": [{"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 64]}}, "('reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1, 64]}}, "('reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1, 64]}}, "('reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1, 64]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [71, 64]}}, "('reward_emb', 'pos_emb')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 64]}}, "('scaler', 'scaler_params', 'actions', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'actions', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'commands', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'commands', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'proprio_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'proprio_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'rewards', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'rewards', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1, 64]}}, "('value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1, 64]}}, "('value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1, 64]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('value_emb', 'pos_emb')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 64]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}