Add jax/plant-starrseq-proto-enhancer

.gitattributes

@@ -69,3 +69,5 @@ jax/plant-starrseq-leaf-enhancer/stage2/checkpoint/ocdbt.process_0/d/b6881d2941c
 jax/plant-starrseq-leaf-promoter_only/stage2/checkpoint/ocdbt.process_0/d/69fef76905902a3d617f0d014399e556 filter=lfs diff=lfs merge=lfs -text
 jax/plant-starrseq-proto-combined/stage2/checkpoint/ocdbt.process_0/d/75d43ca0f69abb39feba538ee11e6807 filter=lfs diff=lfs merge=lfs -text
 jax/plant-starrseq-proto-combined/stage2/checkpoint/ocdbt.process_0/d/dfb4644fa4310c1f8df6dad0a31f1267 filter=lfs diff=lfs merge=lfs -text
+jax/plant-starrseq-proto-enhancer/stage2/checkpoint/ocdbt.process_0/d/127e6b0037d7d21fd02b3f9ea50c50b9 filter=lfs diff=lfs merge=lfs -text
+jax/plant-starrseq-proto-enhancer/stage2/checkpoint/ocdbt.process_0/d/a506a428b5a6a0c45ea29d5337e862f8 filter=lfs diff=lfs merge=lfs -text

jax/plant-starrseq-proto-enhancer/stage1/checkpoint/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1784051912729643600, "commit_timestamp_nsecs": 1784051913204842809, "custom_metadata": {}}

jax/plant-starrseq-proto-enhancer/stage1/checkpoint/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('lam',)": {"key_metadata": [{"key": "lam", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('w',)": {"key_metadata": [{"key": "w", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('xb',)": {"key_metadata": [{"key": "xb", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('yb',)": {"key_metadata": [{"key": "yb", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

jax/plant-starrseq-proto-enhancer/stage1/config.json

@@ -0,0 +1,17 @@
+{
+  "custom_heads": [
+    "plant_probe_ridge"
+  ],
+  "head_configs": {
+    "plant_probe_ridge": {
+      "source": "ridge",
+      "note": "Closed-form ridge on mean-pooled frozen AlphaGenome encoder features. Not an alphagenome_ft model checkpoint \u2014 the encoder is the unmodified pretrained all_folds backbone.",
+      "encoder_dim": 1536,
+      "lambda": 1.0,
+      "predict": "(X - xb) @ w + yb, X = encoder_output.mean(axis=1)"
+    }
+  },
+  "save_full_model": false,
+  "save_minimal_model": false,
+  "use_encoder_output": true
+}

jax/plant-starrseq-proto-enhancer/stage2/checkpoint/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1784051961531650962, "commit_timestamp_nsecs": 1784051962244913073, "custom_metadata": {}}

jax/plant-starrseq-proto-enhancer/stage2/checkpoint/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('0', 'alphagenome/sequence_encoder/dna_embedder/conv1_d', 'b')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv1_d", "key_type": 2}, {"key": "b", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'alphagenome/sequence_encoder/dna_embedder/conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [15, 4, 768]}}, "('0', 'alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('0', 'alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('0', 'alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('0', 'alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 768, 768]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 768, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 896, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 896, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1152, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1152, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1280, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1280, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1408, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1408, 1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d', 'bias')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1536]}}, "('0', 'alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1536, 1536]}}, "('0', 'head/mpra_head/~predict/layer_norm', 'offset')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "head/mpra_head/~predict/layer_norm", "key_type": 2}, {"key": "offset", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('0', 'head/mpra_head/~predict/layer_norm', 'scale')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "head/mpra_head/~predict/layer_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('0', 'head/mpra_head/~predict/linear', 'b')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "head/mpra_head/~predict/linear", "key_type": 2}, {"key": "b", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('0', 'head/mpra_head/~predict/linear', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "head/mpra_head/~predict/linear", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6144, 1024]}}, "('0', 'head/mpra_head/~predict/linear_1', 'b')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "head/mpra_head/~predict/linear_1", "key_type": 2}, {"key": "b", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'head/mpra_head/~predict/linear_1', 'w')": {"key_metadata": [{"key": "0", "key_type": 1}, {"key": "head/mpra_head/~predict/linear_1", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('1', 'alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('1', 'alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 768]}}, "('1', 'alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('1', 'alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 896]}}, "('1', 'alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('1', 'alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1024]}}, "('1', 'alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('1', 'alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1152]}}, "('1', 'alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('1', 'alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1280]}}, "('1', 'alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('1', 'alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1408]}}, "('1', 'alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm', 'var_ema')": {"key_metadata": [{"key": "1", "key_type": 1}, {"key": "alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm", "key_type": 2}, {"key": "var_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1536]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

jax/plant-starrseq-proto-enhancer/stage2/checkpoint/_sharding

@@ -0,0 +1 @@
+{"MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252MV9kLmI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2QuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Quc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Qudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0ub2Zmc2V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2QuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Quc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Qudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0ub2Zmc2V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2QuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Quc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Qudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0ub2Zmc2V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2QuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Quc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Qudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0ub2Zmc2V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2QuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Quc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Qudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0ub2Zmc2V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrL3N0YW5kYXJkaXplZF9jb252MV9kLnc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2QuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Quc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrXzEvc3RhbmRhcmRpemVkX2NvbnYxX2Qudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0ub2Zmc2V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5oZWFkL21wcmFfaGVhZC9-cHJlZGljdC9sYXllcl9ub3JtLm9mZnNldA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5oZWFkL21wcmFfaGVhZC9-cHJlZGljdC9sYXllcl9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5oZWFkL21wcmFfaGVhZC9-cHJlZGljdC9saW5lYXIuYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5oZWFkL21wcmFfaGVhZC9-cHJlZGljdC9saW5lYXIudw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5oZWFkL21wcmFfaGVhZC9-cHJlZGljdC9saW5lYXJfMS53":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MC5oZWFkL21wcmFfaGVhZC9-cHJlZGljdC9saW5lYXJfMS5i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2RuYV9lbWJlZGRlci9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMC9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0udmFyX2VtYQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMS9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0udmFyX2VtYQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMi9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0udmFyX2VtYQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfMy9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0udmFyX2VtYQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNC9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0udmFyX2VtYQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrL3Jtc19iYXRjaF9ub3JtLnZhcl9lbWE=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","MS5hbHBoYWdlbm9tZS9zZXF1ZW5jZV9lbmNvZGVyL2Rvd25yZXNfYmxvY2tfNS9jb252X2Jsb2NrXzEvcm1zX2JhdGNoX25vcm0udmFyX2VtYQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}

jax/plant-starrseq-proto-enhancer/stage2/checkpoint/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv1_d.b", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv1_d.w", "write_shape": [15, 4, 768], "chunk_shape": [15, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/dna_embedder/conv_block/standardized_conv1_d.w", "write_shape": [5, 768, 768], "chunk_shape": [5, 768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d.bias", "write_shape": [896], "chunk_shape": [896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block/standardized_conv1_d.w", "write_shape": [5, 768, 896], "chunk_shape": [5, 768, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm.offset", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm.scale", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d.bias", "write_shape": [896], "chunk_shape": [896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d.scale", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_0/conv_block_1/standardized_conv1_d.w", "write_shape": [5, 896, 896], "chunk_shape": [5, 896, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block/standardized_conv1_d.w", "write_shape": [5, 896, 1024], "chunk_shape": [5, 896, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm.offset", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm.scale", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d.scale", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_1/conv_block_1/standardized_conv1_d.w", "write_shape": [5, 1024, 1024], "chunk_shape": [5, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block/standardized_conv1_d.w", "write_shape": [5, 1024, 1152], "chunk_shape": [5, 1024, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm.offset", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm.scale", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d.scale", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_2/conv_block_1/standardized_conv1_d.w", "write_shape": [5, 1152, 1152], "chunk_shape": [5, 1152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d.bias", "write_shape": [1280], "chunk_shape": [1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block/standardized_conv1_d.w", "write_shape": [5, 1152, 1280], "chunk_shape": [5, 1152, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm.offset", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm.scale", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d.bias", "write_shape": [1280], "chunk_shape": [1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d.scale", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_3/conv_block_1/standardized_conv1_d.w", "write_shape": [5, 1280, 1280], "chunk_shape": [5, 1280, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d.bias", "write_shape": [1408], "chunk_shape": [1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block/standardized_conv1_d.w", "write_shape": [5, 1280, 1408], "chunk_shape": [5, 1280, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm.offset", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm.scale", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d.bias", "write_shape": [1408], "chunk_shape": [1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d.scale", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_4/conv_block_1/standardized_conv1_d.w", "write_shape": [5, 1408, 1408], "chunk_shape": [5, 1408, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm.offset", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm.scale", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d.bias", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d.scale", "write_shape": [1, 1, 1536], "chunk_shape": [1, 1, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block/standardized_conv1_d.w", "write_shape": [5, 1408, 1536], "chunk_shape": [5, 1408, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm.offset", "write_shape": [1, 1, 1536], "chunk_shape": [1, 1, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm.scale", "write_shape": [1, 1, 1536], "chunk_shape": [1, 1, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d.bias", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d.scale", "write_shape": [1, 1, 1536], "chunk_shape": [1, 1, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.alphagenome/sequence_encoder/downres_block_5/conv_block_1/standardized_conv1_d.w", "write_shape": [5, 1536, 1536], "chunk_shape": [5, 1536, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.head/mpra_head/~predict/layer_norm.offset", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.head/mpra_head/~predict/layer_norm.scale", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.head/mpra_head/~predict/linear.b", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.head/mpra_head/~predict/linear.w", "write_shape": [6144, 1024], "chunk_shape": [6144, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.head/mpra_head/~predict/linear_1.b", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.head/mpra_head/~predict/linear_1.w", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/dna_embedder/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_0/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 768], "chunk_shape": [1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_0/conv_block_1/rms_batch_norm.var_ema", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_1/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 896], "chunk_shape": [1, 1, 896], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_1/conv_block_1/rms_batch_norm.var_ema", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_2/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 1024], "chunk_shape": [1, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_2/conv_block_1/rms_batch_norm.var_ema", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_3/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 1152], "chunk_shape": [1, 1, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_3/conv_block_1/rms_batch_norm.var_ema", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_4/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 1280], "chunk_shape": [1, 1, 1280], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_4/conv_block_1/rms_batch_norm.var_ema", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_5/conv_block/rms_batch_norm.var_ema", "write_shape": [1, 1, 1408], "chunk_shape": [1, 1, 1408], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.alphagenome/sequence_encoder/downres_block_5/conv_block_1/rms_batch_norm.var_ema", "write_shape": [1, 1, 1536], "chunk_shape": [1, 1, 1536], "ext_metadata": null}}]}

jax/plant-starrseq-proto-enhancer/stage2/checkpoint/ocdbt.process_0/d/127e6b0037d7d21fd02b3f9ea50c50b9

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

jax/plant-starrseq-proto-enhancer/stage2/checkpoint/ocdbt.process_0/d/a506a428b5a6a0c45ea29d5337e862f8

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

jax/plant-starrseq-proto-enhancer/stage2/config.json

@@ -0,0 +1,23 @@
+{
+  "custom_heads": [
+    "mpra_head"
+  ],
+  "head_configs": {
+    "mpra_head": {
+      "source": "custom",
+      "type": "GENOME_TRACKS",
+      "output_type": "RNA_SEQ",
+      "num_tracks": 1,
+      "name": "mpra_head",
+      "metadata": {
+        "nl_size": 1024,
+        "pooling_type": "flatten",
+        "activation": "relu",
+        "do": 0.1
+      }
+    }
+  },
+  "save_full_model": false,
+  "save_minimal_model": true,
+  "use_encoder_output": false
+}

jax/plant-starrseq-proto-enhancer/stage2/metrics.json

@@ -0,0 +1,10 @@
+{
+  "val_pearson": 0.7869499921798706,
+  "test_pearson": 0.8035923838615417,
+  "source_config": {
+    "tissue": "proto",
+    "mode": "enhancer",
+    "stage2_enabled": true
+  },
+  "source_pickle": "/grid/koo/home/shared/models/alphagenome_encoder/jax/plant_starrseq_proto_enhancer/finetuned_encoder.pkl"
+}