absternator commited on Jun 18

Commit

8bb5775

verified ·

1 Parent(s): 2134d25

remove cylical and seed

Browse files

Files changed (29) hide show

.gitattributes +3 -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/b50490e3b3a737de52244f0b9a9b4ff7 +0 -0
cases/checkpoint/0/model/manifest.ocdbt +0 -0
cases/checkpoint/0/model/ocdbt.process_0/d/0ae6916d82a72162e3c6f18653cd1db6 +3 -0
cases/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt +0 -0
cases/model_config.json +2 -3
cases/preprocessing_config.json +0 -1
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/f6960c2bc59a59b2de53c3fd977595e0 +0 -0
prevalence/checkpoint/0/model/manifest.ocdbt +0 -0
prevalence/checkpoint/0/model/ocdbt.process_0/d/a9cf1f7e373fed53a7fa062bf6ff0a9f +3 -0
prevalence/checkpoint/0/model/ocdbt.process_0/manifest.ocdbt +0 -0
prevalence/model_config.json +2 -3
prevalence/preprocessing_config.json +0 -1
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/6f028d877d6db7a49f9ed03f4e30e165 +0 -0
prevalence/prevalence/checkpoint/0/model/manifest.ocdbt +0 -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 +15 -0
prevalence/prevalence/preprocessing_config.json +55 -0

.gitattributes CHANGED Viewed

@@ -35,3 +35,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 cases/checkpoint/0/model/ocdbt.process_0/d/aa4916a5cf33130352a5a1b62eaca42b filter=lfs diff=lfs merge=lfs -text
 prevalence/checkpoint/0/model/ocdbt.process_0/d/cfaf63ebed00aa0884aeab11dad58eec filter=lfs diff=lfs merge=lfs -text

 *tfevents* filter=lfs diff=lfs merge=lfs -text
 cases/checkpoint/0/model/ocdbt.process_0/d/aa4916a5cf33130352a5a1b62eaca42b filter=lfs diff=lfs merge=lfs -text
 prevalence/checkpoint/0/model/ocdbt.process_0/d/cfaf63ebed00aa0884aeab11dad58eec filter=lfs diff=lfs merge=lfs -text
+cases/checkpoint/0/model/ocdbt.process_0/d/0ae6916d82a72162e3c6f18653cd1db6 filter=lfs diff=lfs merge=lfs -text
+prevalence/checkpoint/0/model/ocdbt.process_0/d/a9cf1f7e373fed53a7fa062bf6ff0a9f filter=lfs diff=lfs merge=lfs -text
+prevalence/prevalence/checkpoint/0/model/ocdbt.process_0/d/cfaf63ebed00aa0884aeab11dad58eec 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": ~~1781690581776149939~~, "commit_timestamp_nsecs": ~~1781690581904832544~~, "custom_metadata": {}}


1	+ {"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1781769025394800066, "commit_timestamp_nsecs": 1781769025540291272, "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": [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": [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}

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": [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": [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}}]}

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

Binary file (7.37 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/0ae6916d82a72162e3c6f18653cd1db6 ADDED Viewed

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

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

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": 16,
   "d_model": 256,
   "n_layers": 4,
   "d_state": 128,
@@ -10,6 +10,5 @@
   "head_dim": 64,
   "chunk_size": 256,
   "output_dim": 1,
-  "dropout": 0.3,
-  "seed": 42
 }

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

cases/preprocessing_config.json CHANGED Viewed

@@ -21,7 +21,6 @@
     "routine"
   ],
   "intervention_day": 3285,
-  "use_cyclical_time": true,
   "predictor": "cases",
   "eps_prevalence": 1e-05,
   "scaler_mean": [

     "routine"
   ],
   "intervention_day": 3285,
   "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": ~~1781690568898697676~~, "commit_timestamp_nsecs": ~~1781690569030953030~~, "custom_metadata": {}}


1	+ {"item_handlers": {"model": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1781768948148563972, "commit_timestamp_nsecs": 1781768948281337790, "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/f6960c2bc59a59b2de53c3fd977595e0 ADDED Viewed

Binary file (7 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/a9cf1f7e373fed53a7fa062bf6ff0a9f ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:0b9e2b371bf8617a3eb470d8693a36b684ade6a7c894e928021e44b694f3ca3f
+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": 16,
   "d_model": 256,
   "n_layers": 4,
   "d_state": 128,
@@ -10,6 +10,5 @@
   "head_dim": 64,
   "chunk_size": 256,
   "output_dim": 1,
-  "dropout": 0.3,
-  "seed": 42
 }

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

prevalence/preprocessing_config.json CHANGED Viewed

@@ -21,7 +21,6 @@
     "routine"
   ],
   "intervention_day": 3285,
-  "use_cyclical_time": true,
   "predictor": "prevalence",
   "eps_prevalence": 1e-05,
   "scaler_mean": [

     "routine"
   ],
   "intervention_day": 3285,
   "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": 1781690568898697676, "commit_timestamp_nsecs": 1781690569030953030, "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": [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}

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": [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}}]}

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

Binary file (5.91 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/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,15 @@

+{
+  "model_type": "Mamba2Regressor",
+  "predictor": "prevalence",
+  "input_size": 16,
+  "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,
+  "seed": 42
+}

prevalence/prevalence/preprocessing_config.json ADDED Viewed

	@@ -0,0 +1,55 @@

+{
+  "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,
+  "use_cyclical_time": true,
+  "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
+  ]
+}