| {"tree_metadata": {"('dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('input_proj', 'bias', 'value')": {"key_metadata": [{"key": "input_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('input_proj', 'kernel', 'value')": {"key_metadata": [{"key": "input_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 256]}}, "('mamba2', 'embedder', 'embedding', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 256]}}, "('mamba2', 'final_norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('mamba2', 'layers', '0', 'mixer', 'A_log', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "A_log", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('mamba2', 'layers', '0', 'mixer', 'D', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "D", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('mamba2', 'layers', '0', 'mixer', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('mamba2', 'layers', '0', 'mixer', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 1, 768]}}, "('mamba2', 'layers', '0', 'mixer', 'dt_bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "dt_bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('mamba2', 'layers', '0', 'mixer', 'in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1288]}}, "('mamba2', 'layers', '0', 'mixer', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('mamba2', 'layers', '0', 'mixer', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 256]}}, "('mamba2', 'layers', '0', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('mamba2', 'layers', '1', 'mixer', 'A_log', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "A_log", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('mamba2', 'layers', '1', 'mixer', 'D', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "D", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('mamba2', 'layers', '1', 'mixer', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('mamba2', 'layers', '1', 'mixer', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 1, 768]}}, "('mamba2', 'layers', '1', 'mixer', 'dt_bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "dt_bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('mamba2', 'layers', '1', 'mixer', 'in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1288]}}, "('mamba2', 'layers', '1', 'mixer', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('mamba2', 'layers', '1', 'mixer', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mixer", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 256]}}, "('mamba2', 'layers', '1', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('output_proj', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "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]}}, "('output_proj', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1]}}}, "use_ocdbt": true, "use_zarr3": true, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} |