ckpt_best/params/_METADATA

{"tree_metadata": {"('params', 'block_0', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_0', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_0', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_0', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_0', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_0', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_0', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_0', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_1', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_1', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_1', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_1', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_1', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_1', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_1', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_1', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_2', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_2', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_2', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_2', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_2', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_2', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_2', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_2', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_3', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_3', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_3', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_3', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_3', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_3', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_3', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_3', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_4', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_4', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_4', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_4', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_4', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_4', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_4', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_4', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_5', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_5', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_5', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_5', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_5', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_5', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_5', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_5', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_6', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_6', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_6', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_6', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_6', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_6', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_6', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_6', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_7', 'attn', 'k_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_7', 'attn', 'out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'block_7', 'attn', 'q_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'block_7', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'block_7', 'ln1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_7', 'ln2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'block_7', 'mlp', 'down', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'block_7', 'mlp', 'up', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "block_7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'embed', 'embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [186, 512]}}, "('params', 'ln_f', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "ln_f", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'ln_x0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "ln_x0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'resid_lambda_0')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_1')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_2')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_3')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_4')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_5')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_6')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'resid_lambda_7')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "resid_lambda_7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_0')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_1')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_2')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_3')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_4')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_5')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_6')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'x0_lambda_7')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "x0_lambda_7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}