2 layer runs

.gitattributes

@@ -47,3 +47,5 @@ prevalence/checkpoint/0/model/ocdbt.process_0/d/d10fca500f1eab46932837dc869c5494
 prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/0049b6ca029ee329b7a1cf4e458f246a filter=lfs diff=lfs merge=lfs -text
 prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/30e8b5517ee25c5a106ead3ea0ef1ac7 filter=lfs diff=lfs merge=lfs -text
 prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/a9cf1f7e373fed53a7fa062bf6ff0a9f filter=lfs diff=lfs merge=lfs -text
+cases/checkpoint/0/model/ocdbt.process_0/d/a06a0ca7b21e4d14424d2d0a4a424908 filter=lfs diff=lfs merge=lfs -text
+prevalence/checkpoint/0/model/ocdbt.process_0/d/edf4f438b62977b3496d99bc6407482d filter=lfs diff=lfs merge=lfs -text

cases/checkpoint/0/_CHECKPOINT_METADATA

@@ -1 +1 @@
-{"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1781869089346365988, "commit_timestamp_nsecs": 1781869089478125398, "custom_metadata": {}}
+{"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1782201840742047273, "commit_timestamp_nsecs": 1782201840979433869, "custom_metadata": {}}

cases/checkpoint/0/model/_METADATA

@@ -1 +1 @@
-{"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]}}, "('mamba2', 'layers', '2', 'mixer', 'A_log', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'D', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'dt_bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '3', 'mixer', 'A_log', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'D', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'dt_bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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}
+{"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}

cases/checkpoint/0/model/_sharding

@@ -1 +1 @@
-{"ZHJvcG91dC5ybmdzLmNvdW50LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHJvcG91dC5ybmdzLmtleS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmVtYmVkZGVyLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmZpbmFsX25vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
+{"ZHJvcG91dC5ybmdzLmNvdW50LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHJvcG91dC5ybmdzLmtleS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmVtYmVkZGVyLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmZpbmFsX25vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}

cases/checkpoint/0/model/array_metadatas/process_0

@@ -1 +1 @@
-{"array_metadatas": [{"array_metadata": {"param_name": "dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "input_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "input_proj.kernel.value", "write_shape": [16, 256], "chunk_shape": [16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.embedder.embedding.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.final_norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.bias.value", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.kernel.value", "write_shape": [256, 1], "chunk_shape": [256, 1], "ext_metadata": null}}]}
+{"array_metadatas": [{"array_metadata": {"param_name": "dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "input_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "input_proj.kernel.value", "write_shape": [16, 256], "chunk_shape": [16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.embedder.embedding.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.final_norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.bias.value", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.kernel.value", "write_shape": [256, 1], "chunk_shape": [256, 1], "ext_metadata": null}}]}

cases/checkpoint/0/model/ocdbt.process_0/d/a06a0ca7b21e4d14424d2d0a4a424908

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c29335426a5c931c70d98fd9dbebef4e14809a6a3d4dbe80bfeccb3ca15ec595
+size 3488111

cases/model_config.json

@@ -3,12 +3,12 @@
   "predictor": "cases",
   "input_size": 16,
   "d_model": 256,
-  "n_layers": 4,
+  "n_layers": 2,
   "d_state": 128,
   "d_conv": 4,
   "expand": 2,
   "head_dim": 64,
   "chunk_size": 256,
   "output_dim": 1,
-  "dropout": 0.3
+  "dropout": 0.57
 }

prevalence/checkpoint/0/_CHECKPOINT_METADATA

@@ -1 +1 @@
-{"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1781869029705347980, "commit_timestamp_nsecs": 1781869029846940211, "custom_metadata": {}}
+{"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1782201591094655013, "commit_timestamp_nsecs": 1782201591315196000, "custom_metadata": {}}

prevalence/checkpoint/0/model/_METADATA

@@ -1 +1 @@
-{"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]}}, "('mamba2', 'layers', '2', 'mixer', 'A_log', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'D', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'dt_bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'mixer', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '2', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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', '3', 'mixer', 'A_log', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'D', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'dt_bias', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'mixer', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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', '3', 'norm', 'weight', 'value')": {"key_metadata": [{"key": "mamba2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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}
+{"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}

prevalence/checkpoint/0/model/_sharding

@@ -1 +1 @@
-{"ZHJvcG91dC5ybmdzLmNvdW50LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHJvcG91dC5ybmdzLmtleS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmVtYmVkZGVyLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmZpbmFsX25vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4yLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4zLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
+{"ZHJvcG91dC5ybmdzLmNvdW50LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHJvcG91dC5ybmdzLmtleS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","aW5wdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3V0cHV0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmVtYmVkZGVyLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmZpbmFsX25vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4wLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkFfbG9nLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLkQudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmR0X2JpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm1peGVyLmluX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","bWFtYmEyLmxheWVycy4xLm5vcm0ud2VpZ2h0LnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}

prevalence/checkpoint/0/model/array_metadatas/process_0

@@ -1 +1 @@
-{"array_metadatas": [{"array_metadata": {"param_name": "dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "input_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "input_proj.kernel.value", "write_shape": [16, 256], "chunk_shape": [16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.embedder.embedding.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.final_norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.2.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.3.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.bias.value", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.kernel.value", "write_shape": [256, 1], "chunk_shape": [256, 1], "ext_metadata": null}}]}
+{"array_metadatas": [{"array_metadata": {"param_name": "dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "input_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "input_proj.kernel.value", "write_shape": [16, 256], "chunk_shape": [16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.embedder.embedding.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.final_norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.0.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.A_log.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.D.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.conv1d.conv.kernel.value", "write_shape": [4, 1, 768], "chunk_shape": [4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.dt_bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.in_proj.kernel.value", "write_shape": [256, 1288], "chunk_shape": [256, 1288], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.norm.weight.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.mixer.out_proj.kernel.value", "write_shape": [512, 256], "chunk_shape": [512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "mamba2.layers.1.norm.weight.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.bias.value", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.kernel.value", "write_shape": [256, 1], "chunk_shape": [256, 1], "ext_metadata": null}}]}

prevalence/checkpoint/0/model/ocdbt.process_0/d/edf4f438b62977b3496d99bc6407482d

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8829ca874d251e74ddde94b4c520b4d3b6266d867352131aa70b49aab5e0b0b0
+size 3456690

prevalence/model_config.json

@@ -3,12 +3,12 @@
   "predictor": "prevalence",
   "input_size": 16,
   "d_model": 256,
-  "n_layers": 4,
+  "n_layers": 2,
   "d_state": 128,
   "d_conv": 4,
   "expand": 2,
   "head_dim": 64,
   "chunk_size": 256,
   "output_dim": 1,
-  "dropout": 0.3
+  "dropout": 0.24
 }