{"tree_metadata": {"('embedding', 'pos_emb', 'embedding', 'value')": {"key_metadata": [{"key": "embedding", "key_type": 2}, {"key": "pos_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('embedding', 'token_emb', 'embedding', 'value')": {"key_metadata": [{"key": "embedding", "key_type": 2}, {"key": "token_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [50257, 192]}}, "('output_layer', 'kernel', 'value')": {"key_metadata": [{"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 50257]}}, "('transformer_blocks', '0', 'attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '0', 'attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '0', 'attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('transformer_blocks', '0', 'attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('transformer_blocks', '0', 'attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '0', 'attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '0', 'attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '0', 'attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '1', 'attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '1', 'attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '1', 'attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('transformer_blocks', '1', 'attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('transformer_blocks', '1', 'attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '1', 'attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '1', 'attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '1', 'attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '2', 'attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '2', 'attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '2', 'attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('transformer_blocks', '2', 'attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('transformer_blocks', '2', 'attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '2', 'attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '2', 'attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '2', 'attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '3', 'attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '3', 'attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '3', 'attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('transformer_blocks', '3', 'attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('transformer_blocks', '3', 'attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '3', 'attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '3', 'attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '3', 'attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '4', 'attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '4', 'attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '4', 'attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('transformer_blocks', '4', 'attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('transformer_blocks', '4', 'attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '4', 'attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '4', 'attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '4', 'attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '5', 'attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '5', 'attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '5', 'attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('transformer_blocks', '5', 'attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('transformer_blocks', '5', 'attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '5', 'attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('transformer_blocks', '5', 'attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('transformer_blocks', '5', 'attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}