diff --git a/.gitattributes b/.gitattributes
index 860031b3f606629db25a5ac1ec44f0ac3aa24f14..2e717cc1c1f0da03d7a407a9e96f6e6262b76d75 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -59,3 +59,22 @@ simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/b732753
 simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f filter=lfs diff=lfs merge=lfs -text
 simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e filter=lfs diff=lfs merge=lfs -text
 simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 filter=lfs diff=lfs merge=lfs -text
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/assets/robomme/norm_stats.json b/simulation/robocasa/full_attention/stage2/step-150000/assets/robomme/norm_stats.json
new file mode 100644
index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/assets/robomme/norm_stats.json
@@ -0,0 +1,136 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        0.005957702174782753,
+        -0.009433974511921406,
+        -0.027537286281585693,
+        -0.0002329433336853981,
+        -0.0019665092695504427,
+        0.0010746258776634932,
+        -0.17213615775108337,
+        0.0022201647516340017,
+        -0.0015274320030584931,
+        0.0009210301213897765,
+        -0.000019904651708202437,
+        -0.7345612645149231
+      ],
+      "std": [
+        0.2822888493537903,
+        0.31111952662467957,
+        0.30976882576942444,
+        0.060622796416282654,
+        0.08188065886497498,
+        0.08514650166034698,
+        0.9420629143714905,
+        0.06350194662809372,
+        0.1455080211162567,
+        0.05478019639849663,
+        0.0006542060291394591,
+        0.6832558512687683
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6785714030265808,
+        -0.1971428543329239,
+        -0.27142858505249023,
+        -0.30571427941322327,
+        -1.0,
+        -0.07321428507566452,
+        -0.5982142686843872,
+        -0.24081368744373322,
+        0.0,
+        -1.0
+      ],
+      "q99": [
+        0.6482142806053162,
+        1.0,
+        1.0,
+        0.1971428543329239,
+        0.2514285743236542,
+        0.3028571307659149,
+        1.0,
+        0.19107143580913544,
+        0.612500011920929,
+        0.24883341789245605,
+        0.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        0.5936015248298645,
+        -0.20600375533103943,
+        0.4474135637283325,
+        -0.30925947427749634,
+        -0.007771184667944908,
+        -0.04887167364358902,
+        0.1312701404094696,
+        2.4705684185028076,
+        -1.8421560525894165,
+        0.6735693216323853,
+        -0.018410686403512955,
+        -0.002943282015621662,
+        0.05312279239296913,
+        0.595700740814209,
+        0.03035394847393036,
+        -0.030480381101369858
+      ],
+      "std": [
+        0.8089187145233154,
+        0.6286742687225342,
+        0.18121899664402008,
+        0.8043732047080994,
+        0.2596682906150818,
+        0.349939227104187,
+        0.21797646582126617,
+        1.5330910682678223,
+        1.3586890697479248,
+        0.09442085027694702,
+        0.2916492521762848,
+        0.08581095188856125,
+        0.6641682982444763,
+        0.32726120948791504,
+        0.011618087068200111,
+        0.011510285548865795
+      ],
+      "q01": [
+        0.042294517159461975,
+        -3.4426209926605225,
+        0.0796390026807785,
+        -0.9991932511329651,
+        -0.7087220549583435,
+        -1.0,
+        8.079589264298193e-8,
+        -3.263056755065918,
+        -5.687641620635986,
+        0.23626811802387238,
+        -0.9920241236686707,
+        -0.3793386220932007,
+        -1.0,
+        9.924601727107074e-7,
+        0.0007141401874832809,
+        -0.04061261937022209
+      ],
+      "q99": [
+        4.271139144897461,
+        0.4542047083377838,
+        0.8608905673027039,
+        0.9988747239112854,
+        0.813831627368927,
+        1.0,
+        1.0,
+        6.125694274902344,
+        0.5118050575256348,
+        0.7076432704925537,
+        0.9900418519973755,
+        0.32602807879447937,
+        1.0,
+        1.0,
+        0.04024730622768402,
+        -0.0001981566019821912
+      ]
+    }
+  }
+}
\ No newline at end of file
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/_METADATA b/simulation/robocasa/full_attention/stage2/step-150000/params/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..3174ecb3aadf31e14f1f7c8b230a9cbe3b80c0dc
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 4304]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288, 2048]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "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]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "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, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 256, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 256, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1, 256, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 512, 16384]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 256, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256, 3072]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_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]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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": [256, 32]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "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]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "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]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "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, 1024]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "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]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "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]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "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]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "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]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "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]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "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]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "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, 1024]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "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]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "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]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "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]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "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]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "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, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "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, 8, 128]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "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, 2048]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "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]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "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]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "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, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "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]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "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]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "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, 1024]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "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]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "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, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/_sharding b/simulation/robocasa/full_attention/stage2/step-150000/params/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..8e8a2723950e24f11d6c7d7ef67a79f8ccf59c4d
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/_sharding
@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}"}
\ No newline at end of file
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/array_metadatas/process_0 b/simulation/robocasa/full_attention/stage2/step-150000/params/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..5eacdbf35307909fee5f510e0873b2c738cafdbf
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 288, 4304], "chunk_shape": [27, 288, 4304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 4, 72, 1152], "chunk_shape": [27, 4, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [288, 2048], "chunk_shape": [288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [256, 3072], "chunk_shape": [256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 2, 256, 2048], "chunk_shape": [18, 2, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 2, 256, 1024], "chunk_shape": [18, 2, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 2, 2048, 256], "chunk_shape": [18, 2, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 2, 1024, 256], "chunk_shape": [18, 2, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 1, 256, 1024], "chunk_shape": [18, 1, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 1, 1024, 256], "chunk_shape": [18, 1, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 256, 2048], "chunk_shape": [18, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 512, 16384], "chunk_shape": [18, 2, 512, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 256, 4096], "chunk_shape": [18, 2, 256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 256, 3072], "chunk_shape": [18, 256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 256, 3072], "chunk_shape": [18, 256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [3, 512], "chunk_shape": [3, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 b/simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9
new file mode 100644
index 0000000000000000000000000000000000000000..c2829ba644b3c90f69e2a24dbb19c2579eee130e
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:476ff74dc38c756c9c195faec2152f34b052f4263d9e19899bc40360de746aae
+size 158450
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/manifest.ocdbt b/simulation/robocasa/full_attention/stage2/step-150000/params/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..f9a67dc627e4ec7ffe90cae0ae33f8049c68fea5
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/manifest.ocdbt differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709
new file mode 100644
index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709 differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9
new file mode 100644
index 0000000000000000000000000000000000000000..dba5303318fbf97dda54ca5bc7d3aa75ac8b3f99
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:824eaea25c59bac7143776b1c37425d47a9426986d0246d389cc6a5eca912ed8
+size 606012846
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60
new file mode 100644
index 0000000000000000000000000000000000000000..236c8b96073f50c3551cf32f34d6f98c5882e494
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2290868f5d3f5fb15c3e9342cc3fa0adeb4e0517ea2ba667c00e1f9555b32be0
+size 258762
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8
new file mode 100644
index 0000000000000000000000000000000000000000..3375329894cd2666f0875e243ab0b76f6482beea
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7088b323898c94fe320df837059a255f9ac31db7a06b98cb350a373b4e7e855f
+size 1048614538
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5
new file mode 100644
index 0000000000000000000000000000000000000000..e9b55d519ab04cfff1769c24d2c9b5154df4bd03
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3855446dea4243997fe5ad2e90d0e7a9f76d000c414f3e6998606716384078db
+size 2239695988
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1
new file mode 100644
index 0000000000000000000000000000000000000000..80fb627928c156305708bd239c9cb1e017f26eec
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed6ad005edb24d33eb04226b8bee988a0ef5c44e3ef94ec84b56765001326be4
+size 2239686238
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129
new file mode 100644
index 0000000000000000000000000000000000000000..ad055f7a7f18fe1bccc605d0ae354067803db860
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129 differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891
new file mode 100644
index 0000000000000000000000000000000000000000..db849424a23ca20f90e975d1112ba1634ff4cbb3
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60b237dba5c9d4a74f511cb55156e8a59b605339a72cbfb12a6c3fab1b1a6f8
+size 1120293104
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524
new file mode 100644
index 0000000000000000000000000000000000000000..0d52b5cf70d83b1810119f630b224e340320c7a6
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524 differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6
new file mode 100644
index 0000000000000000000000000000000000000000..eb22654e4958d0d465397b732958faf33fe3dac0
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6 differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853
new file mode 100644
index 0000000000000000000000000000000000000000..0277dc69c3e76357fded92f1415e96dd146bf0c7
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853 differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3
new file mode 100644
index 0000000000000000000000000000000000000000..1a2ea0ca6e8386bbc6eb3e79fdfea23edef271b3
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3 differ
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065
new file mode 100644
index 0000000000000000000000000000000000000000..004d36c4434232db9f660e7a9f77e3cef88c328e
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d58075ff02f957a5bfc1f9892c8956b19df5dfd13968cea4cb51530ded6b1137
+size 158430
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761
new file mode 100644
index 0000000000000000000000000000000000000000..d484b7f2caa187a89b614304a4f65e483bc496e5
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b42d860c3b83040a1bf3da8ae1e73497e29fa83d4ab6d7c7bf2ee1f9a9bac22c
+size 2587810029
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25
new file mode 100644
index 0000000000000000000000000000000000000000..af09e1794f91285055749415128ca909d6fe5040
--- /dev/null
+++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f3e1d8d11d6d0f2d58a47be5d1c47d533b9c07d892e41057f1acc6cb2be23a9
+size 2117122051
diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..b3c0725729cac774cfbd49806482419fa6f9405b
Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt differ
diff --git a/simulation/robocasa/recurrent/.gitattributes b/simulation/robocasa/recurrent/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..8df4ba5a44a559b25144859e1653e7f8f309c3af
--- /dev/null
+++ b/simulation/robocasa/recurrent/.gitattributes
@@ -0,0 +1,44 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 filter=lfs diff=lfs merge=lfs -text
+stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 filter=lfs diff=lfs merge=lfs -text
diff --git a/simulation/robocasa/recurrent/LICENSE_GEMMA.txt b/simulation/robocasa/recurrent/LICENSE_GEMMA.txt
new file mode 100644
index 0000000000000000000000000000000000000000..e0f8930d83fe9eef924c0eed50ee0e572eb471fa
--- /dev/null
+++ b/simulation/robocasa/recurrent/LICENSE_GEMMA.txt
@@ -0,0 +1,77 @@
+Gemma Terms of Use 
+
+Last modified: February 21, 2024
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement.
+
+Section 1: DEFINITIONS
+1.1 Definitions
+(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference.
+
+(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service").
+
+(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from.
+
+(d) "Google" means Google LLC.
+
+(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives.
+
+(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service.
+
+1.2
+As used in this Agreement, "including" means "including without limitation".
+
+Section 2: ELIGIBILITY AND USAGE
+2.1 Eligibility
+You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity.
+
+2.2 Use
+You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement.
+
+Section 3: DISTRIBUTION AND RESTRICTIONS
+3.1 Distribution and Redistribution
+You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions:
+
+You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2.
+You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement.
+You must cause any modified files to carry prominent notices stating that you modified the files.
+All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms".
+You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement.
+
+3.2 Use Restrictions
+You must not use any of the Gemma Services:
+
+for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or
+in violation of applicable laws and regulations.
+To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement.
+
+3.3 Generated Output
+Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses.
+
+Section 4: ADDITIONAL PROVISIONS
+4.1 Updates
+Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma.
+
+4.2 Trademarks
+Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein.
+
+4.3 DISCLAIMER OF WARRANTY
+UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT.
+
+4.4 LIMITATION OF LIABILITY
+TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+4.5 Term, Termination, and Survival
+The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement.
+
+4.6 Governing Law and Jurisdiction
+This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+4.7 Severability
+If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
+
+4.8 Entire Agreement
+This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter.
+
+4.9 No Waiver
+Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement.
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/NOTICE b/simulation/robocasa/recurrent/NOTICE
new file mode 100644
index 0000000000000000000000000000000000000000..1b9cb6541d7135c4c78a542ac4a3e04a971c96db
--- /dev/null
+++ b/simulation/robocasa/recurrent/NOTICE
@@ -0,0 +1,11 @@
+IMPRINT
+Copyright 2026
+
+This distribution contains modifications to the OpenPI project:
+https://github.com/Physical-Intelligence/openpi
+
+OpenPI is distributed under the Apache License 2.0. Gemma components and
+derivatives are subject to the Gemma Terms of Use in LICENSE_GEMMA.txt.
+
+Gemma is provided under and subject to the Gemma Terms of Use found at
+ai.google.dev/gemma/terms.
diff --git a/simulation/robocasa/recurrent/README.md b/simulation/robocasa/recurrent/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..7d64b6dee24b77f06ef7f2a5f55a3ba9504289f2
--- /dev/null
+++ b/simulation/robocasa/recurrent/README.md
@@ -0,0 +1,38 @@
+---
+license: gemma
+library_name: jax
+tags:
+  - robotics
+  - vision-language-action
+  - memory
+  - robocasa
+  - dinov3
+  - recurrent
+---
+
+# IMPRINT RoboCasa Recurrent Checkpoint
+
+Canonical Pi0.5 IMPRINT checkpoint for the customized 12-task RoboCasa
+benchmark using the recurrent gated-delta memory encoder.
+
+- Stage 1: trained 180,000 updates, selected at 160,000, W&B `far-wandb/Robocasa/ag6ytzrm`
+- Stage 2: 150,000 updates, W&B `far-wandb/Robocasa/1ozs9qy8`
+- memory tokens: 256 x 1024
+- policy horizon: 20
+- training scan: chunkwise WY
+- inference update: incremental recurrent
+- Stage 2: memory encoder finetuned; image tower frozen
+
+The Stage-2 model-level auxiliary weight was zero. The archived history YAML
+retains a legacy `aux_weight: 0.25` preload flag, which caused auxiliary
+targets to be read but did not contribute an auxiliary loss.
+
+Use `yilin-wu/imprint-robocasa-600` for training and
+`yilin-wu/imprint-robocasa-environment` for the exact patched RoboCasa 1.0.1
+simulator source. This release contains the complete Stage-2 Orbax parameter
+checkpoint, including the PaliGemma/Gemma backbone. It is a Gemma model
+derivative and is distributed subject to the included Gemma Terms of Use.
+Large simulator assets are not duplicated.
+
+The matching training store, patched simulator source, and reproducibility
+code are publicly available under the IMPRINT release.
diff --git a/simulation/robocasa/recurrent/SHA256SUMS b/simulation/robocasa/recurrent/SHA256SUMS
new file mode 100644
index 0000000000000000000000000000000000000000..181d534d8e76656e3b938d7257ebceb4a7a1bf91
--- /dev/null
+++ b/simulation/robocasa/recurrent/SHA256SUMS
@@ -0,0 +1,27 @@
+3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132  ./LICENSE_GEMMA.txt
+5b21d7116ac42a0ab6141488cb654c2411557ac977e2276accb88ef774f28c19  ./NOTICE
+1c4d7719baef48359e81dc02e729ec9361c51acf5ba3bd788f21dcd3268b5277  ./README.md
+3f5ee9416db5224d6effaea03e799540023817301a0ba2677922a31709216b1d  ./config/history_config.yaml
+4143882de7d4e41230b03bd95a7001745139f6a757e64b4db3afa2857ec5b9e7  ./config/input_provenance.json
+39a8f9a475d18c9167a3872c7679a348b81a165776fe3e16c391f354c7ad8514  ./config/norm_stats.json
+7567aca14012be88045a10203432c7117c128b4b1455718143c429cea8155fc1  ./stage1/stage1_step160000.npz
+fe7ad82eeb9e1096c90f433a50fb4b2fabe000779a6d14bfe8ee317f6aea9c72  ./stage2/step-150000/_CHECKPOINT_METADATA
+39a8f9a475d18c9167a3872c7679a348b81a165776fe3e16c391f354c7ad8514  ./stage2/step-150000/assets/robomme/norm_stats.json
+68305898798e4aa166cbeff6f5375e02d51d18c84920784c0a7b34159292f481  ./stage2/step-150000/params/_METADATA
+306825e3332fd927d96181a73228f3480a3ceb705c7db27fe49a2563204ebb8f  ./stage2/step-150000/params/_sharding
+d86f5bf955a3a93391b23797e2c7e1cd146341f989074e96b0678b8351067c88  ./stage2/step-150000/params/array_metadatas/process_0
+1d455b9c7cd4010bac48990f5b286f3d4764e677b577361e584d5de51702d72d  ./stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1
+c56bbba464e09062f252686a9b61c77d21f4f6c9f5176c0ce0e07f5ccab96860  ./stage2/step-150000/params/manifest.ocdbt
+2f44abe9cc4c9b57417a84eccbc4cfaacbf78138b754f58deb3663e8f0b9cd0c  ./stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b
+78c760f9903cc094536408330040c07125539871571bcd8f66addfeff3489211  ./stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075
+ede3e60d71ddb1726e98428f57b9be80a113cc3101a18c8c1a0362224951497a  ./stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb
+57a8cccbbd6268bbe3cd7892014c27da7f4bc9e734cea8dd64881c144d15a9f5  ./stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2
+d5c5bbdb19cc2144308820ef54b0382b2eb07092c2ed0ea67a251914752735ae  ./stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9
+81594df81b1b9b7ae76963f695bff5e4ed1da3a0d7e29dea42475d84e432e6b9  ./stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966
+8b3d3995e3aef785f5244dae3e6588ddf1f0386876778f95303b68f3b41d71ef  ./stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f
+d86dfa44bb73ca13d774185159e069a4d7aa6585af4a569f9f7569b9460b3392  ./stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191
+2144a693e034a9dab2595f94cfaf7f29f9a92384fdb8827e7556ecd3b8d42f53  ./stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01
+034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7  ./stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c
+af709070d5c1a4b648d7799923f503c658d2403dfffed9a54fc301a89d18a1f9  ./stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19
+3614c36b08b05532184f2b34cd514d9d2dcc3f0be7038e846356371d62e52e54  ./stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20
+7c124c19a20c79bc57bde2839dee54753c2290f0fe166070dcf59258e3222f1b  ./stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt
diff --git a/simulation/robocasa/recurrent/config/history_config.yaml b/simulation/robocasa/recurrent/config/history_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e3d972ab3e3d527bbefd8115c06108707d3eda0b
--- /dev/null
+++ b/simulation/robocasa/recurrent/config/history_config.yaml
@@ -0,0 +1,89 @@
+budget: 256
+num_views: 3
+streaming_obs_horizon: 16
+integration_type: modulation
+memory_token_dim: 1024
+representation_type: cosmo_bottleneck
+memory_feature:
+  img:
+    net: identity
+    input_dim: 768
+  pos:
+    input_dim: 768
+    hidden_dim: 768
+  state:
+    input_dim: 16
+    hidden_dim: 512
+cosmo_bottleneck_memory:
+  memory_encoder_type: gated_delta_hybrid
+  memory_architecture_version: gated_delta_hybrid_v1
+  memory_compute_dtype: bfloat16
+  visual_feature_stream: dinov3
+  latent_source: offline
+  n_variants: 8
+  num_query_samples: 16
+  variant: Wan2.2-TI2V-5B
+  temporal_factor: 4
+  downsample: 2
+  num_latents: 256
+  num_heads: 8
+  cosmos_latent_dim: 48
+  siglip_dim: 768
+  proprio_in_dim: 16
+  proprio_encoding: raw
+  action_in_dim: 12
+  action_encoding: raw
+  lang_emb_dim: 2048
+  action_tokenizer_hidden: 512
+  proprio_tokenizer_hidden: 512
+  pool_hw: 2
+  wan_h: 256
+  wan_w: 256
+  latent_cap: 4096
+  latent_buckets:
+    - 256
+    - 512
+    - 1024
+    - 2048
+    - 4096
+  max_recent_tokens: 768
+  max_step_tokens: 1776
+  max_stream_steps: 1776
+  delta_num_layers: 2
+  delta_num_heads: 8
+  delta_head_dim: 128
+  delta_conv_kernel: 4
+  delta_ffn_inner: 2048
+  delta_decay_bias: -4.0
+  delta_beta_bias: -1.0
+  delta_scan_checkpoint_size: 20
+  delta_full_scan_mode: chunkwise_wy
+  delta_wy_matmul_dtype: bfloat16
+  delta_recent_fusion_blocks: 1
+  delta_max_chunks: 384
+  delta_event_order: action_proprio_then_finalized_visual_v1
+cosmo_decoder:
+  num_heads: 8
+  num_blocks: 2
+  align_dim: 768
+  align_num_patches: 192
+  ffn_inner: 8192
+  action_dim: 12
+  proprio_dim: 16
+  max_query_steps: 2048
+frame_jitter:
+  enabled: false
+  start_jitter: 3
+  stride_choices:
+    - 1
+    - 2
+    - 3
+  interval_jitter: 1
+  canonical_downsample: 2
+stream_aug:
+  enabled: false
+  gauss_sigma: 0.05
+  stream_drop_p: 0.1
+  feat_drop_p: 0.0
+aux_weight: 0.25
+action_horizon: 20
diff --git a/simulation/robocasa/recurrent/config/input_provenance.json b/simulation/robocasa/recurrent/config/input_provenance.json
new file mode 100644
index 0000000000000000000000000000000000000000..59e7878ede6291481d9539c527bd4e7be4b106d7
--- /dev/null
+++ b/simulation/robocasa/recurrent/config/input_provenance.json
@@ -0,0 +1,28 @@
+{
+  "action_causality": "strictly_causal",
+  "architecture": "gated_delta_hybrid_v1",
+  "base_policy": "pi0.5",
+  "checkpoint_metadata_sha256": "fe7ad82eeb9e1096c90f433a50fb4b2fabe000779a6d14bfe8ee317f6aea9c72",
+  "checkpoint_path": "stage2/step-150000",
+  "checkpoint_step": 150000,
+  "dataset": "yilin-wu/imprint-robocasa-600",
+  "eval_consolidate_all": 1,
+  "execution_horizon": 20,
+  "freeze_variant": "image_only",
+  "latent_cap": 4096,
+  "legacy_history_aux_preload_flag": 0.25,
+  "max_recent_tokens": 768,
+  "max_step_tokens": 1776,
+  "memory_inference_mode": "incremental_recurrent",
+  "memory_language_source": "live",
+  "memory_training_scan": "chunkwise_wy",
+  "num_latents": 256,
+  "prediction_horizon": 20,
+  "stage1_checkpoint": "stage1/stage1_step160000.npz",
+  "stage1_sha256": "7567aca14012be88045a10203432c7117c128b4b1455718143c429cea8155fc1",
+  "stage1_wandb_run": "far-wandb/Robocasa/ag6ytzrm",
+  "stage2_auxiliary_loss_weight": 0.0,
+  "stage2_memory_encoder": "finetuned",
+  "stage2_wandb_run": "far-wandb/Robocasa/1ozs9qy8",
+  "visual_feature_stream": "dinov3"
+}
diff --git a/simulation/robocasa/recurrent/config/norm_stats.json b/simulation/robocasa/recurrent/config/norm_stats.json
new file mode 100644
index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46
--- /dev/null
+++ b/simulation/robocasa/recurrent/config/norm_stats.json
@@ -0,0 +1,136 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        0.005957702174782753,
+        -0.009433974511921406,
+        -0.027537286281585693,
+        -0.0002329433336853981,
+        -0.0019665092695504427,
+        0.0010746258776634932,
+        -0.17213615775108337,
+        0.0022201647516340017,
+        -0.0015274320030584931,
+        0.0009210301213897765,
+        -0.000019904651708202437,
+        -0.7345612645149231
+      ],
+      "std": [
+        0.2822888493537903,
+        0.31111952662467957,
+        0.30976882576942444,
+        0.060622796416282654,
+        0.08188065886497498,
+        0.08514650166034698,
+        0.9420629143714905,
+        0.06350194662809372,
+        0.1455080211162567,
+        0.05478019639849663,
+        0.0006542060291394591,
+        0.6832558512687683
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6785714030265808,
+        -0.1971428543329239,
+        -0.27142858505249023,
+        -0.30571427941322327,
+        -1.0,
+        -0.07321428507566452,
+        -0.5982142686843872,
+        -0.24081368744373322,
+        0.0,
+        -1.0
+      ],
+      "q99": [
+        0.6482142806053162,
+        1.0,
+        1.0,
+        0.1971428543329239,
+        0.2514285743236542,
+        0.3028571307659149,
+        1.0,
+        0.19107143580913544,
+        0.612500011920929,
+        0.24883341789245605,
+        0.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        0.5936015248298645,
+        -0.20600375533103943,
+        0.4474135637283325,
+        -0.30925947427749634,
+        -0.007771184667944908,
+        -0.04887167364358902,
+        0.1312701404094696,
+        2.4705684185028076,
+        -1.8421560525894165,
+        0.6735693216323853,
+        -0.018410686403512955,
+        -0.002943282015621662,
+        0.05312279239296913,
+        0.595700740814209,
+        0.03035394847393036,
+        -0.030480381101369858
+      ],
+      "std": [
+        0.8089187145233154,
+        0.6286742687225342,
+        0.18121899664402008,
+        0.8043732047080994,
+        0.2596682906150818,
+        0.349939227104187,
+        0.21797646582126617,
+        1.5330910682678223,
+        1.3586890697479248,
+        0.09442085027694702,
+        0.2916492521762848,
+        0.08581095188856125,
+        0.6641682982444763,
+        0.32726120948791504,
+        0.011618087068200111,
+        0.011510285548865795
+      ],
+      "q01": [
+        0.042294517159461975,
+        -3.4426209926605225,
+        0.0796390026807785,
+        -0.9991932511329651,
+        -0.7087220549583435,
+        -1.0,
+        8.079589264298193e-8,
+        -3.263056755065918,
+        -5.687641620635986,
+        0.23626811802387238,
+        -0.9920241236686707,
+        -0.3793386220932007,
+        -1.0,
+        9.924601727107074e-7,
+        0.0007141401874832809,
+        -0.04061261937022209
+      ],
+      "q99": [
+        4.271139144897461,
+        0.4542047083377838,
+        0.8608905673027039,
+        0.9988747239112854,
+        0.813831627368927,
+        1.0,
+        1.0,
+        6.125694274902344,
+        0.5118050575256348,
+        0.7076432704925537,
+        0.9900418519973755,
+        0.32602807879447937,
+        1.0,
+        1.0,
+        0.04024730622768402,
+        -0.0001981566019821912
+      ]
+    }
+  }
+}
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/stage1/stage1_step160000.npz b/simulation/robocasa/recurrent/stage1/stage1_step160000.npz
new file mode 100644
index 0000000000000000000000000000000000000000..3a3e7c32bac553285e3444b4950470e53a2efb12
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage1/stage1_step160000.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7567aca14012be88045a10203432c7117c128b4b1455718143c429cea8155fc1
+size 391011410
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/_CHECKPOINT_METADATA b/simulation/robocasa/recurrent/stage2/step-150000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..6a4b614b88361e569f91e3b944e3c966ee7b2974
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1789085274077480773, "commit_timestamp_nsecs": 1789085308428284253, "custom_metadata": {}}
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/assets/robomme/norm_stats.json b/simulation/robocasa/recurrent/stage2/step-150000/assets/robomme/norm_stats.json
new file mode 100644
index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/assets/robomme/norm_stats.json
@@ -0,0 +1,136 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        0.005957702174782753,
+        -0.009433974511921406,
+        -0.027537286281585693,
+        -0.0002329433336853981,
+        -0.0019665092695504427,
+        0.0010746258776634932,
+        -0.17213615775108337,
+        0.0022201647516340017,
+        -0.0015274320030584931,
+        0.0009210301213897765,
+        -0.000019904651708202437,
+        -0.7345612645149231
+      ],
+      "std": [
+        0.2822888493537903,
+        0.31111952662467957,
+        0.30976882576942444,
+        0.060622796416282654,
+        0.08188065886497498,
+        0.08514650166034698,
+        0.9420629143714905,
+        0.06350194662809372,
+        0.1455080211162567,
+        0.05478019639849663,
+        0.0006542060291394591,
+        0.6832558512687683
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6785714030265808,
+        -0.1971428543329239,
+        -0.27142858505249023,
+        -0.30571427941322327,
+        -1.0,
+        -0.07321428507566452,
+        -0.5982142686843872,
+        -0.24081368744373322,
+        0.0,
+        -1.0
+      ],
+      "q99": [
+        0.6482142806053162,
+        1.0,
+        1.0,
+        0.1971428543329239,
+        0.2514285743236542,
+        0.3028571307659149,
+        1.0,
+        0.19107143580913544,
+        0.612500011920929,
+        0.24883341789245605,
+        0.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        0.5936015248298645,
+        -0.20600375533103943,
+        0.4474135637283325,
+        -0.30925947427749634,
+        -0.007771184667944908,
+        -0.04887167364358902,
+        0.1312701404094696,
+        2.4705684185028076,
+        -1.8421560525894165,
+        0.6735693216323853,
+        -0.018410686403512955,
+        -0.002943282015621662,
+        0.05312279239296913,
+        0.595700740814209,
+        0.03035394847393036,
+        -0.030480381101369858
+      ],
+      "std": [
+        0.8089187145233154,
+        0.6286742687225342,
+        0.18121899664402008,
+        0.8043732047080994,
+        0.2596682906150818,
+        0.349939227104187,
+        0.21797646582126617,
+        1.5330910682678223,
+        1.3586890697479248,
+        0.09442085027694702,
+        0.2916492521762848,
+        0.08581095188856125,
+        0.6641682982444763,
+        0.32726120948791504,
+        0.011618087068200111,
+        0.011510285548865795
+      ],
+      "q01": [
+        0.042294517159461975,
+        -3.4426209926605225,
+        0.0796390026807785,
+        -0.9991932511329651,
+        -0.7087220549583435,
+        -1.0,
+        8.079589264298193e-8,
+        -3.263056755065918,
+        -5.687641620635986,
+        0.23626811802387238,
+        -0.9920241236686707,
+        -0.3793386220932007,
+        -1.0,
+        9.924601727107074e-7,
+        0.0007141401874832809,
+        -0.04061261937022209
+      ],
+      "q99": [
+        4.271139144897461,
+        0.4542047083377838,
+        0.8608905673027039,
+        0.9988747239112854,
+        0.813831627368927,
+        1.0,
+        1.0,
+        6.125694274902344,
+        0.5118050575256348,
+        0.7076432704925537,
+        0.9900418519973755,
+        0.32602807879447937,
+        1.0,
+        1.0,
+        0.04024730622768402,
+        -0.0001981566019821912
+      ]
+    }
+  }
+}
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/_METADATA b/simulation/robocasa/recurrent/stage2/step-150000/params/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..636ae950c6862f3ae9d4c1a0b4e40f216af0afd0
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "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]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_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]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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": [256, 32]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'beta_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "beta_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": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'beta_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "beta_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, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'decay_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "decay_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": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'decay_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "decay_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, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "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, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "k_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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "k_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'norm_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_in", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'norm_in', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_gate', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_gate", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_gate', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_norm", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "q_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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "q_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "v_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]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "v_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'beta_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "beta_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": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'beta_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "beta_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, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'decay_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "decay_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": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'decay_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "decay_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, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_out", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_out", "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, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "k_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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "k_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'norm_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_in", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'norm_in', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_gate', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_gate", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_gate', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_norm", "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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "q_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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "q_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "v_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]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "v_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": [1024, 256]}}, "('params', 'mem_encoder', 'event_slot', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "event_slot", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('params', 'mem_encoder', 'event_type', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "event_type", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "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]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'read_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_norm", "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]}}, "('params', 'mem_encoder', 'read_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_out", "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]}}, "('params', 'mem_encoder', 'read_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'read_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_q", "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]}}, "('params', 'mem_encoder', 'read_q', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_q", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "key", "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, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "out", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "query", "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, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "value", "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, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_kv_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_kv_norm", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_kv_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_kv_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_q_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_q_norm", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_q_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_q_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "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, 1024]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "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, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "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, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "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, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_norm", "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]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_null', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_null", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 256]}}, "('params', 'mem_encoder', 'time_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "time_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]}}, "('params', 'mem_encoder', 'time_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "time_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": [1, 256]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "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]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "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, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "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]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "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]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "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]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/_sharding b/simulation/robocasa/recurrent/stage2/step-150000/params/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..4a6e96dbc26510168ad386a01cd30f863dfac7c5
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/_sharding
@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm5vcm1faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm5vcm1faW4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9nYXRlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9nYXRlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmJldGFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmJldGFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmNvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmNvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmRlY2F5X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmRlY2F5X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm5vcm1faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm5vcm1faW4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9nYXRlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9nYXRlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmJldGFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmJldGFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmNvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmNvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmRlY2F5X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmRlY2F5X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmV2ZW50X3Nsb3QudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmV2ZW50X3R5cGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19rdl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19rdl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19xX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19xX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5faW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9udWxsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfcS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRpbWVfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRpbWVfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"}
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/array_metadatas/process_0 b/simulation/robocasa/recurrent/stage2/step-150000/params/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..1674d2dff173c2f90742beebf091c36816411bbe
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [1024, 768], "chunk_shape": [1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.beta_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.beta_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.conv.kernel.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.decay_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.decay_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.k_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.norm_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.norm_in.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_gate.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_gate.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.q_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.v_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.beta_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.beta_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.conv.kernel.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.decay_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.decay_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.k_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.norm_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.norm_in.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_gate.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_gate.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.q_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.v_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.event_slot.value", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.event_type.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_q.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_null.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.time_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.time_proj.kernel.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [3, 512], "chunk_shape": [3, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 b/simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1
new file mode 100644
index 0000000000000000000000000000000000000000..e27154cde53ee258ce7f1c5abe7a4473a84c77d3
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d455b9c7cd4010bac48990f5b286f3d4764e677b577361e584d5de51702d72d
+size 196861
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/manifest.ocdbt b/simulation/robocasa/recurrent/stage2/step-150000/params/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..a1a37be397ad0662fac95ccf7f00b4b7936da6ec
Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/manifest.ocdbt differ
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b
new file mode 100644
index 0000000000000000000000000000000000000000..cd287190512c3f96a737bc04bdb8b504e4e9ef59
Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b differ
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075
new file mode 100644
index 0000000000000000000000000000000000000000..2a43b49d5ce574d513aa1404b5b0c95e29748d27
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:78c760f9903cc094536408330040c07125539871571bcd8f66addfeff3489211
+size 2030582856
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb
new file mode 100644
index 0000000000000000000000000000000000000000..258f02dc60b38c6e34b3e0def1c9855a70f1ff07
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ede3e60d71ddb1726e98428f57b9be80a113cc3101a18c8c1a0362224951497a
+size 196857
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2
new file mode 100644
index 0000000000000000000000000000000000000000..6887e66c412833401819a5005c6060f358e1302c
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57a8cccbbd6268bbe3cd7892014c27da7f4bc9e734cea8dd64881c144d15a9f5
+size 2240054650
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9
new file mode 100644
index 0000000000000000000000000000000000000000..252e72af5487357dd2d528fa1fd2d7bdb006a9ba
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5c5bbdb19cc2144308820ef54b0382b2eb07092c2ed0ea67a251914752735ae
+size 79040671
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966
new file mode 100644
index 0000000000000000000000000000000000000000..49ff9c74e5e88dcd4713c142d6a0599d79afc148
Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966 differ
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f
new file mode 100644
index 0000000000000000000000000000000000000000..55ca384744eb81310252f6fce617225cb731dae0
Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f differ
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191
new file mode 100644
index 0000000000000000000000000000000000000000..4e1fc196beb146f21ef9a34c84078a1564a9d77b
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d86dfa44bb73ca13d774185159e069a4d7aa6585af4a569f9f7569b9460b3392
+size 1227340570
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01
new file mode 100644
index 0000000000000000000000000000000000000000..57d53d046386f84a55550d8bb651e1e608b39d4c
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2144a693e034a9dab2595f94cfaf7f29f9a92384fdb8827e7556ecd3b8d42f53
+size 1680405616
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c
new file mode 100644
index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243
Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c differ
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19
new file mode 100644
index 0000000000000000000000000000000000000000..126dffad8b23049ab59381236b4fe8259cf3ef1a
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af709070d5c1a4b648d7799923f503c658d2403dfffed9a54fc301a89d18a1f9
+size 2240068387
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20
new file mode 100644
index 0000000000000000000000000000000000000000..a92395dff6dcdfb3eb3be9c27515f4a00f4107e3
--- /dev/null
+++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3614c36b08b05532184f2b34cd514d9d2dcc3f0be7038e846356371d62e52e54
+size 2517086701
diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..2d77683f5468fb8cdd93e9db50f808e3b24a7ae1
Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt differ