| {"tree_metadata": {"('Q_head', 'layers', '1', 'bias', 'value')": {"key_metadata": [{"key": "Q_head", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "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]}}, "('Q_head', 'layers', '1', 'kernel', 'value')": {"key_metadata": [{"key": "Q_head", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1]}}, "('init_y', 'scale')": {"key_metadata": [{"key": "init_y", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('init_z', 'scale')": {"key_metadata": [{"key": "init_z", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('input_embedding', 'embedding', 'value')": {"key_metadata": [{"key": "input_embedding", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10, 128]}}, "('net', 'net', 'layers', '0', 'd_mixer', 'W1', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "d_mixer", "key_type": 2}, {"key": "W1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1536]}}, "('net', 'net', 'layers', '0', 'd_mixer', 'W2', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "d_mixer", "key_type": 2}, {"key": "W2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 512]}}, "('net', 'net', 'layers', '0', 'd_mixer', 'W3', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "d_mixer", "key_type": 2}, {"key": "W3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1536]}}, "('net', 'net', 'layers', '0', 'l_mixer', 'W1', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "l_mixer", "key_type": 2}, {"key": "W1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [81, 64]}}, "('net', 'net', 'layers', '0', 'l_mixer', 'W2', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "l_mixer", "key_type": 2}, {"key": "W2", "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, 81]}}, "('net', 'net', 'layers', '0', 'l_mixer', 'W3', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "l_mixer", "key_type": 2}, {"key": "W3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [81, 64]}}, "('net', 'net', 'layers', '1', 'd_mixer', 'W1', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "d_mixer", "key_type": 2}, {"key": "W1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1536]}}, "('net', 'net', 'layers', '1', 'd_mixer', 'W2', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "d_mixer", "key_type": 2}, {"key": "W2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 512]}}, "('net', 'net', 'layers', '1', 'd_mixer', 'W3', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "d_mixer", "key_type": 2}, {"key": "W3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1536]}}, "('net', 'net', 'layers', '1', 'l_mixer', 'W1', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "l_mixer", "key_type": 2}, {"key": "W1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [81, 64]}}, "('net', 'net', 'layers', '1', 'l_mixer', 'W2', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "l_mixer", "key_type": 2}, {"key": "W2", "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, 81]}}, "('net', 'net', 'layers', '1', 'l_mixer', 'W3', 'kernel', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "l_mixer", "key_type": 2}, {"key": "W3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [81, 64]}}, "('net', 'net', 'layers', '2', 'scale', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "net", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('net', 'x_norm', 'scale', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "x_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('net', 'y_norm', 'scale', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "y_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('net', 'z_norm', 'scale', 'value')": {"key_metadata": [{"key": "net", "key_type": 2}, {"key": "z_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('output_head', 'bias', 'value')": {"key_metadata": [{"key": "output_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('output_head', 'kernel', 'value')": {"key_metadata": [{"key": "output_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 10]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} |