absternator commited on Jun 19

Commit

c7f4514

verified ·

1 Parent(s): 080e885

revertto using cyclical

Browse files

Files changed (35) hide show

.gitattributes +5 -0
cases/checkpoint/0/_CHECKPOINT_METADATA +1 -1
cases/checkpoint/0/model/_METADATA +1 -1
cases/checkpoint/0/model/array_metadatas/process_0 +1 -1
cases/checkpoint/0/model/d/ac157745cac15cd5f74b263800640aef +0 -0
cases/checkpoint/0/model/manifest.ocdbt +0 -0
cases/checkpoint/0/model/ocdbt.process_0/d/94295e7c41f00840659b71cce2db8478 +3 -0
cases/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt +0 -0
cases/model_config.json +1 -1
cases/preprocessing_config.json +1 -0
prevalence/checkpoint/0/_CHECKPOINT_METADATA +1 -1
prevalence/checkpoint/0/model/_METADATA +1 -1
prevalence/checkpoint/0/model/array_metadatas/process_0 +1 -1
prevalence/checkpoint/0/model/d/938add3a2e0a0f3da0f55daf07751030 +0 -0
prevalence/checkpoint/0/model/manifest.ocdbt +0 -0
prevalence/checkpoint/0/model/ocdbt.process_0/d/d10fca500f1eab46932837dc869c5494 +3 -0
prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt +0 -0
prevalence/model_config.json +1 -1
prevalence/preprocessing_config.json +1 -0
prevalence/prevalence/checkpoint/0/_CHECKPOINT_METADATA +1 -0
prevalence/prevalence/checkpoint/0/model/_METADATA +1 -0
prevalence/prevalence/checkpoint/0/model/_sharding +1 -0
prevalence/prevalence/checkpoint/0/model/array_metadatas/process_0 +1 -0
prevalence/prevalence/checkpoint/0/model/d/1379fa36a2cdae35737569f8a80705b2 +0 -0
prevalence/prevalence/checkpoint/0/model/d/6f028d877d6db7a49f9ed03f4e30e165 +0 -0
prevalence/prevalence/checkpoint/0/model/d/c67f084fb225decf0d17b56a3669291d +0 -0
prevalence/prevalence/checkpoint/0/model/d/f6960c2bc59a59b2de53c3fd977595e0 +0 -0
prevalence/prevalence/checkpoint/0/model/manifest.ocdbt +0 -0
prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/0049b6ca029ee329b7a1cf4e458f246a +3 -0
prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/30e8b5517ee25c5a106ead3ea0ef1ac7 +3 -0
prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/a9cf1f7e373fed53a7fa062bf6ff0a9f +3 -0
prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/cfaf63ebed00aa0884aeab11dad58eec +3 -0
prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt +0 -0
prevalence/prevalence/model_config.json +14 -0
prevalence/prevalence/preprocessing_config.json +57 -0

.gitattributes CHANGED Viewed

@@ -42,3 +42,8 @@ cases/checkpoint/0/model/ocdbt.process_0/d/c1650d46cf961976b0c0ad5e546afc83 filt
 prevalence/checkpoint/0/model/ocdbt.process_0/d/30e8b5517ee25c5a106ead3ea0ef1ac7 filter=lfs diff=lfs merge=lfs -text
 cases/checkpoint/0/model/ocdbt.process_0/d/607444173de9ed3857284a8fb132fbd0 filter=lfs diff=lfs merge=lfs -text
 prevalence/checkpoint/0/model/ocdbt.process_0/d/0049b6ca029ee329b7a1cf4e458f246a filter=lfs diff=lfs merge=lfs -text

 prevalence/checkpoint/0/model/ocdbt.process_0/d/30e8b5517ee25c5a106ead3ea0ef1ac7 filter=lfs diff=lfs merge=lfs -text
 cases/checkpoint/0/model/ocdbt.process_0/d/607444173de9ed3857284a8fb132fbd0 filter=lfs diff=lfs merge=lfs -text
 prevalence/checkpoint/0/model/ocdbt.process_0/d/0049b6ca029ee329b7a1cf4e458f246a filter=lfs diff=lfs merge=lfs -text
+cases/checkpoint/0/model/ocdbt.process_0/d/94295e7c41f00840659b71cce2db8478 filter=lfs diff=lfs merge=lfs -text
+prevalence/checkpoint/0/model/ocdbt.process_0/d/d10fca500f1eab46932837dc869c5494 filter=lfs diff=lfs merge=lfs -text
+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/_CHECKPOINT_METADATA CHANGED Viewed

	@@ -1 +1 @@
1	- {"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": ~~1781856517936657603~~, "commit_timestamp_nsecs": ~~1781856518072237388~~, "custom_metadata": {}}


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": {}}

cases/checkpoint/0/model/_METADATA CHANGED Viewed

@@ -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": [15, 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]}}, "('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}

cases/checkpoint/0/model/array_metadatas/process_0 CHANGED Viewed

@@ -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": [15, 256], "chunk_shape": [15, 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": "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}}]}

cases/checkpoint/0/model/d/ac157745cac15cd5f74b263800640aef ADDED Viewed

Binary file (7.44 kB). View file

cases/checkpoint/0/model/manifest.ocdbt CHANGED Viewed

Binary files a/cases/checkpoint/0/model/manifest.ocdbt and b/cases/checkpoint/0/model/manifest.ocdbt differ

cases/checkpoint/0/model/ocdbt.process_0/d/94295e7c41f00840659b71cce2db8478 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f6c4ea2e116d195af4a4fd5d7136c87c3f432566d7b09fb5dc912f4de2a2cc91
+size 6938624

cases/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt CHANGED Viewed

Binary files a/cases/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt and b/cases/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt differ

cases/model_config.json CHANGED Viewed

@@ -1,7 +1,7 @@
 {
   "model_type": "Mamba2Regressor",
   "predictor": "cases",
-  "input_size": 15,
   "d_model": 256,
   "n_layers": 4,
   "d_state": 128,

 {
   "model_type": "Mamba2Regressor",
   "predictor": "cases",
+  "input_size": 16,
   "d_model": 256,
   "n_layers": 4,
   "d_state": 128,

cases/preprocessing_config.json CHANGED Viewed

@@ -24,6 +24,7 @@
   "n_steps": 157,
   "burnin_day": 2190,
   "window_size": 14,
   "predictor": "cases",
   "eps_prevalence": 1e-05,
   "scaler_mean": [

   "n_steps": 157,
   "burnin_day": 2190,
   "window_size": 14,
+  "use_cyclical_time": true,
   "predictor": "cases",
   "eps_prevalence": 1e-05,
   "scaler_mean": [

prevalence/checkpoint/0/_CHECKPOINT_METADATA CHANGED Viewed

	@@ -1 +1 @@
1	- {"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": ~~1781856479050278438~~, "commit_timestamp_nsecs": ~~1781856479190601137~~, "custom_metadata": {}}


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": {}}

prevalence/checkpoint/0/model/_METADATA CHANGED Viewed

@@ -1 +1 @@

prevalence/checkpoint/0/model/array_metadatas/process_0 CHANGED Viewed

@@ -1 +1 @@

prevalence/checkpoint/0/model/d/938add3a2e0a0f3da0f55daf07751030 ADDED Viewed

Binary file (5.91 kB). View file

prevalence/checkpoint/0/model/manifest.ocdbt CHANGED Viewed

Binary files a/prevalence/checkpoint/0/model/manifest.ocdbt and b/prevalence/checkpoint/0/model/manifest.ocdbt differ

prevalence/checkpoint/0/model/ocdbt.process_0/d/d10fca500f1eab46932837dc869c5494 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:41d82142e8044b9e994de624eefd8937dd6e2e2aebaf1fc23e34396656ce4717
+size 6893568

prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt CHANGED Viewed

Binary files a/prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt and b/prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt differ

prevalence/model_config.json CHANGED Viewed

@@ -1,7 +1,7 @@
 {
   "model_type": "Mamba2Regressor",
   "predictor": "prevalence",
-  "input_size": 15,
   "d_model": 256,
   "n_layers": 4,
   "d_state": 128,

 {
   "model_type": "Mamba2Regressor",
   "predictor": "prevalence",
+  "input_size": 16,
   "d_model": 256,
   "n_layers": 4,
   "d_state": 128,

prevalence/preprocessing_config.json CHANGED Viewed

@@ -24,6 +24,7 @@
   "n_steps": 157,
   "burnin_day": 2190,
   "window_size": 14,
   "predictor": "prevalence",
   "eps_prevalence": 1e-05,
   "scaler_mean": [

   "n_steps": 157,
   "burnin_day": 2190,
   "window_size": 14,
+  "use_cyclical_time": true,
   "predictor": "prevalence",
   "eps_prevalence": 1e-05,
   "scaler_mean": [

prevalence/prevalence/checkpoint/0/_CHECKPOINT_METADATA ADDED Viewed

	@@ -0,0 +1 @@


1	+ {"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1781856479050278438, "commit_timestamp_nsecs": 1781856479190601137, "custom_metadata": {}}

prevalence/prevalence/checkpoint/0/model/_METADATA ADDED Viewed

	@@ -0,0 +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": [15, 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}

prevalence/prevalence/checkpoint/0/model/_sharding ADDED Viewed

	@@ -0,0 +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\"}"}

prevalence/prevalence/checkpoint/0/model/array_metadatas/process_0 ADDED Viewed

	@@ -0,0 +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": [15, 256], "chunk_shape": [15, 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}}]}

prevalence/prevalence/checkpoint/0/model/d/1379fa36a2cdae35737569f8a80705b2 ADDED Viewed

Binary file (7 kB). View file

prevalence/prevalence/checkpoint/0/model/d/6f028d877d6db7a49f9ed03f4e30e165 ADDED Viewed

Binary file (5.91 kB). View file

prevalence/prevalence/checkpoint/0/model/d/c67f084fb225decf0d17b56a3669291d ADDED Viewed

Binary file (7.01 kB). View file

prevalence/prevalence/checkpoint/0/model/d/f6960c2bc59a59b2de53c3fd977595e0 ADDED Viewed

Binary file (7 kB). View file

prevalence/prevalence/checkpoint/0/model/manifest.ocdbt ADDED Viewed

Binary file (116 Bytes). View file

prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/0049b6ca029ee329b7a1cf4e458f246a ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:61d57d2fbeb54c2df5783521cf05251a9052a1db93cbbcfd92d340c4eb167d01
+size 6893568

prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/30e8b5517ee25c5a106ead3ea0ef1ac7 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6e31bb65b0bcb86ff20fa2cbeae25a1b7cfc8d912b32cc214a3dac97e7543f56
+size 6893568

prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/a9cf1f7e373fed53a7fa062bf6ff0a9f ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:0b9e2b371bf8617a3eb470d8693a36b684ade6a7c894e928021e44b694f3ca3f
+size 6893568

prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/cfaf63ebed00aa0884aeab11dad58eec ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:9e1a33012b23046d23e38ab02e0dad0e6d376b853871a910d53240a70c58ca0d
+size 6893568

prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt ADDED Viewed

Binary file (119 Bytes). View file

prevalence/prevalence/model_config.json ADDED Viewed

	@@ -0,0 +1,14 @@

+{
+  "model_type": "Mamba2Regressor",
+  "predictor": "prevalence",
+  "input_size": 15,
+  "d_model": 256,
+  "n_layers": 4,
+  "d_state": 128,
+  "d_conv": 4,
+  "expand": 2,
+  "head_dim": 64,
+  "chunk_size": 256,
+  "output_dim": 1,
+  "dropout": 0.3
+}

prevalence/prevalence/preprocessing_config.json ADDED Viewed

	@@ -0,0 +1,57 @@

+{
+  "static_covars": [
+    "eir",
+    "dn0_use",
+    "dn0_future",
+    "Q0",
+    "phi_bednets",
+    "seasonal",
+    "routine",
+    "itn_use",
+    "irs_use",
+    "itn_future",
+    "irs_future",
+    "lsm"
+  ],
+  "after_intervention": [
+    "dn0_future",
+    "itn_future",
+    "irs_future",
+    "lsm",
+    "routine"
+  ],
+  "intervention_day": 3285,
+  "n_steps": 157,
+  "burnin_day": 2190,
+  "window_size": 14,
+  "predictor": "prevalence",
+  "eps_prevalence": 1e-05,
+  "scaler_mean": [
+    58.74008560180664,
+    0.31604108214378357,
+    0.3175233006477356,
+    0.7978994250297546,
+    0.6636720895767212,
+    0.482501357793808,
+    0.48787713050842285,
+    0.25194787979125977,
+    0.15634511411190033,
+    0.25312861800193787,
+    0.17558598518371582,
+    0.24548254907131195
+  ],
+  "scaler_scale": [
+    104.59117126464844,
+    0.16008581221103668,
+    0.15946531295776367,
+    0.11743186414241791,
+    0.16026772558689117,
+    0.49969369173049927,
+    0.4998529851436615,
+    0.32460200786590576,
+    0.26269596815109253,
+    0.32312172651290894,
+    0.2886512279510498,
+    0.32189667224884033
+  ]
+}