{"tree_metadata": {"('embedding_layer', 'pos_emb', 'embedding', 'value')": {"key_metadata": [{"key": "embedding_layer", "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": [64, 128]}}, "('embedding_layer', 'token_emb', 'embedding', 'value')": {"key_metadata": [{"key": "embedding_layer", "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": [50258, 32]}}, "('output_layer', 'bias', 'value')": {"key_metadata": [{"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [25129]}}, "('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": [64, 25129]}}, "('transformer_blocks', '0', 'layer_norm1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '0', 'layer_norm1', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '0', 'layer_norm2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '0', 'layer_norm2', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '0', 'linear1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '0', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '0', 'linear2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '0', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '0', 'mha', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '0', 'mha', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '0', 'mha', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [32]}}, "('transformer_blocks', '0', 'mha', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [2, 16, 256]}}, "('transformer_blocks', '0', 'mha', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '0', 'mha', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '0', 'mha', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '0', 'mha', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '1', 'layer_norm1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '1', 'layer_norm1', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '1', 'layer_norm2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '1', 'layer_norm2', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '1', 'linear1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '1', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '1', 'linear2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '1', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '1', 'mha', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '1', 'mha', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '1', 'mha', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [32]}}, "('transformer_blocks', '1', 'mha', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [2, 16, 256]}}, "('transformer_blocks', '1', 'mha', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '1', 'mha', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '1', 'mha', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '1', 'mha', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '2', 'layer_norm1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '2', 'layer_norm1', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '2', 'layer_norm2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '2', 'layer_norm2', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '2', 'linear1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '2', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '2', 'linear2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '2', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '2', 'mha', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '2', 'mha', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '2', 'mha', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [32]}}, "('transformer_blocks', '2', 'mha', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [2, 16, 256]}}, "('transformer_blocks', '2', 'mha', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '2', 'mha', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '2', 'mha', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '2', 'mha', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '3', 'layer_norm1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '3', 'layer_norm1', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '3', 'layer_norm2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '3', 'layer_norm2', 'scale', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '3', 'linear1', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '3', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '3', 'linear2', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('transformer_blocks', '3', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('transformer_blocks', '3', 'mha', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '3', 'mha', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '3', 'mha', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [32]}}, "('transformer_blocks', '3', 'mha', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [2, 16, 256]}}, "('transformer_blocks', '3', 'mha', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '3', 'mha', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}, "('transformer_blocks', '3', 'mha', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [1, 32]}}, "('transformer_blocks', '3', 'mha', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer_blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mha", "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": [64, 4, 32]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}