diff --git a/config.json b/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..81919cba361a24d921f6d7ff2adada227186a951
--- /dev/null
+++ b/config.json
@@ -0,0 +1,31 @@
+{
+  "_name_or_path": "./",
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 16,
+  "d_kv": 6,
+  "d_model": 8,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "silu",
+  "dropout_rate": 0.0,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-silu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "num_decoder_layers": 3,
+  "num_heads": 4,
+  "num_layers": 3,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 64,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.26.0",
+  "use_cache": true,
+  "vocab_size": 256
+}
diff --git a/t5x/README.md b/t5x/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..3c7c6a10d0d6f251612d9fbc86faab06d66fd918
--- /dev/null
+++ b/t5x/README.md
@@ -0,0 +1,3 @@
+---
+license: cc-by-nc-sa-4.0
+---
diff --git a/t5x/checkpoint_0/checkpoint b/t5x/checkpoint_0/checkpoint
new file mode 100644
index 0000000000000000000000000000000000000000..1d8f53ed2e8a5b5cfe18229ade865f0dfd95445a
Binary files /dev/null and b/t5x/checkpoint_0/checkpoint differ
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_cross_attention_layer_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_cross_attention_layer_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_cross_attention_layer_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_cross_attention_layer_norm.scale.v/0.0 b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_cross_attention_layer_norm.scale.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..7a541332966b437f9fee4acf09513affb498403b
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_cross_attention_layer_norm.scale.v/0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_mlp_layer_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_mlp_layer_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_mlp_layer_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_mlp_layer_norm.scale.v/0.0 b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_mlp_layer_norm.scale.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..7a541332966b437f9fee4acf09513affb498403b
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_mlp_layer_norm.scale.v/0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_self_attention_layer_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_self_attention_layer_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_self_attention_layer_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_self_attention_layer_norm.scale.v/0.0 b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_self_attention_layer_norm.scale.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..7a541332966b437f9fee4acf09513affb498403b
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.decoder.decoder.pre_self_attention_layer_norm.scale.v/0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.relpos_bias.rel_embedding.v/.zarray b/t5x/checkpoint_0/state.param_states.decoder.decoder.relpos_bias.rel_embedding.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..1b229d8ada2624c8191ad9134dc6f9822d81da64
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.decoder.decoder.relpos_bias.rel_embedding.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,64],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,64],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder.relpos_bias.rel_embedding.v/0.0.0 b/t5x/checkpoint_0/state.param_states.decoder.decoder.relpos_bias.rel_embedding.v/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..9ff54768c5fb2aa1ead672604380c5e8f8f401f9
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.decoder.decoder.relpos_bias.rel_embedding.v/0.0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.decoder.decoder_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..01925abad1a84572db2d6d5531c6c7a7707617a0
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.decoder.decoder_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.decoder.decoder_norm.scale.v/0 b/t5x/checkpoint_0/state.param_states.decoder.decoder_norm.scale.v/0
new file mode 100644
index 0000000000000000000000000000000000000000..b06a1e2feb61c1ee0a97573e27060a5596d83bfe
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.decoder.decoder_norm.scale.v/0 differ
diff --git a/t5x/checkpoint_0/state.param_states.decoder.logits_dense.kernel.v/.zarray b/t5x/checkpoint_0/state.param_states.decoder.logits_dense.kernel.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..592900ffa2fdcdbf28e3d1f5df2d48c1718b7bf1
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.decoder.logits_dense.kernel.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,256],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,256],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.decoder.logits_dense.kernel.v/0.0 b/t5x/checkpoint_0/state.param_states.decoder.logits_dense.kernel.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..cb1feb018d0678c54bb055f4bbf7d2eec94020b3
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.decoder.logits_dense.kernel.v/0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_attention_layer_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_attention_layer_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_attention_layer_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_attention_layer_norm.scale.v/0.0 b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_attention_layer_norm.scale.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..7a541332966b437f9fee4acf09513affb498403b
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_attention_layer_norm.scale.v/0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_mlp_layer_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_mlp_layer_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_mlp_layer_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_mlp_layer_norm.scale.v/0.0 b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_mlp_layer_norm.scale.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..7a541332966b437f9fee4acf09513affb498403b
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.encoder.encoder.pre_mlp_layer_norm.scale.v/0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder.relpos_bias.rel_embedding.v/.zarray b/t5x/checkpoint_0/state.param_states.encoder.encoder.relpos_bias.rel_embedding.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..1b229d8ada2624c8191ad9134dc6f9822d81da64
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.encoder.encoder.relpos_bias.rel_embedding.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,64],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,64],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder.relpos_bias.rel_embedding.v/0.0.0 b/t5x/checkpoint_0/state.param_states.encoder.encoder.relpos_bias.rel_embedding.v/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..9ff54768c5fb2aa1ead672604380c5e8f8f401f9
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.encoder.encoder.relpos_bias.rel_embedding.v/0.0.0 differ
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder_norm.scale.v/.zarray b/t5x/checkpoint_0/state.param_states.encoder.encoder_norm.scale.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..01925abad1a84572db2d6d5531c6c7a7707617a0
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.encoder.encoder_norm.scale.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.encoder.encoder_norm.scale.v/0 b/t5x/checkpoint_0/state.param_states.encoder.encoder_norm.scale.v/0
new file mode 100644
index 0000000000000000000000000000000000000000..b06a1e2feb61c1ee0a97573e27060a5596d83bfe
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.encoder.encoder_norm.scale.v/0 differ
diff --git a/t5x/checkpoint_0/state.param_states.token_embedder.embedding.v/.zarray b/t5x/checkpoint_0/state.param_states.token_embedder.embedding.v/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..4e8729ee90c3e6de528fe0eec0f1b0a994ebe0d3
--- /dev/null
+++ b/t5x/checkpoint_0/state.param_states.token_embedder.embedding.v/.zarray
@@ -0,0 +1 @@
+{"chunks":[256,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[256,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/state.param_states.token_embedder.embedding.v/0.0 b/t5x/checkpoint_0/state.param_states.token_embedder.embedding.v/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..cb1feb018d0678c54bb055f4bbf7d2eec94020b3
Binary files /dev/null and b/t5x/checkpoint_0/state.param_states.token_embedder.embedding.v/0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.key.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.key.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.key.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.key.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.key.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..17ec3c567c896c7cb3728c488477b3cb08a75ffa
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.key.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.out.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.out.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..afa6346653cc293ffb96cd1bc743cc6910218cd0
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.out.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,6,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,6,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.out.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.out.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..3cf7f96f83c3c1f10e3e839e52cb16d6cc45ba51
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.out.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.query.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.query.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.query.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.query.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.query.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..ca2f73daea7169cd53d3161c398977fb7ccf6225
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.query.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.value.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.value.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.value.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.value.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.value.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..e64df1e5e0925523369fc774f1871ca4ef39b5f2
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.encoder_decoder_attention.value.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_0.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_0.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..cd07a9e840865497f6ff379cb68d73af7e37b719
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_0.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,16],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,16],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_0.kernel/0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_0.kernel/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..0660258bd071a22d7df2545a9cc643041ebb75f6
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_0.kernel/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_1.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_1.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..cd07a9e840865497f6ff379cb68d73af7e37b719
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_1.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,16],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,16],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_1.kernel/0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_1.kernel/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..f80444ad894b27fc05e073ac81fed88340906fc8
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.mlp.wi_1.kernel/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.mlp.wo.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.mlp.wo.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..fb0b55cbaa4249ea483455d5e1565c95cc50bfb7
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.mlp.wo.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[16,3,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[16,3,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.mlp.wo.kernel/0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.mlp.wo.kernel/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..9e25498368ce0b12237fb2bbac6e921f70264c1f
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.mlp.wo.kernel/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.pre_cross_attention_layer_norm.scale/.zarray b/t5x/checkpoint_0/target.decoder.decoder.pre_cross_attention_layer_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.pre_cross_attention_layer_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.pre_cross_attention_layer_norm.scale/0.0 b/t5x/checkpoint_0/target.decoder.decoder.pre_cross_attention_layer_norm.scale/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..d97f8728d629c9ec08b533f3adfbb8bda505a038
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.pre_cross_attention_layer_norm.scale/0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.pre_mlp_layer_norm.scale/.zarray b/t5x/checkpoint_0/target.decoder.decoder.pre_mlp_layer_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.pre_mlp_layer_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.pre_mlp_layer_norm.scale/0.0 b/t5x/checkpoint_0/target.decoder.decoder.pre_mlp_layer_norm.scale/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..d97f8728d629c9ec08b533f3adfbb8bda505a038
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.pre_mlp_layer_norm.scale/0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.pre_self_attention_layer_norm.scale/.zarray b/t5x/checkpoint_0/target.decoder.decoder.pre_self_attention_layer_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.pre_self_attention_layer_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.pre_self_attention_layer_norm.scale/0.0 b/t5x/checkpoint_0/target.decoder.decoder.pre_self_attention_layer_norm.scale/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..d97f8728d629c9ec08b533f3adfbb8bda505a038
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.pre_self_attention_layer_norm.scale/0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.relpos_bias.rel_embedding/.zarray b/t5x/checkpoint_0/target.decoder.decoder.relpos_bias.rel_embedding/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..1b229d8ada2624c8191ad9134dc6f9822d81da64
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.relpos_bias.rel_embedding/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,64],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,64],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.relpos_bias.rel_embedding/0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.relpos_bias.rel_embedding/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..8ae421e3a3768343e78b3b15832a7d8f2d30f1db
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.relpos_bias.rel_embedding/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.key.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.self_attention.key.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.self_attention.key.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.key.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.self_attention.key.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..49ae29c23e4c2d6e5244ae71f7e1e6da67d862bf
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.self_attention.key.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.out.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.self_attention.out.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..afa6346653cc293ffb96cd1bc743cc6910218cd0
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.self_attention.out.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,6,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,6,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.out.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.self_attention.out.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..497a6089c0ba8ab5b33d9e1905cb11f61ea61aba
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.self_attention.out.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.query.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.self_attention.query.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.self_attention.query.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.query.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.self_attention.query.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..b5c05d1bda33f594263bf41a4bebd7e00b18b254
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.self_attention.query.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.value.kernel/.zarray b/t5x/checkpoint_0/target.decoder.decoder.self_attention.value.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder.self_attention.value.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder.self_attention.value.kernel/0.0.0.0 b/t5x/checkpoint_0/target.decoder.decoder.self_attention.value.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..eaaf72407eb0a2dbcd59c096942414f7edc1cc3a
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder.self_attention.value.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.decoder.decoder_norm.scale/.zarray b/t5x/checkpoint_0/target.decoder.decoder_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..01925abad1a84572db2d6d5531c6c7a7707617a0
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.decoder_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.decoder_norm.scale/0 b/t5x/checkpoint_0/target.decoder.decoder_norm.scale/0
new file mode 100644
index 0000000000000000000000000000000000000000..034c3e7dac64ee11390a2a9b44a622e413f12dd6
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.decoder_norm.scale/0 differ
diff --git a/t5x/checkpoint_0/target.decoder.logits_dense.kernel/.zarray b/t5x/checkpoint_0/target.decoder.logits_dense.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..592900ffa2fdcdbf28e3d1f5df2d48c1718b7bf1
--- /dev/null
+++ b/t5x/checkpoint_0/target.decoder.logits_dense.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,256],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,256],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.decoder.logits_dense.kernel/0.0 b/t5x/checkpoint_0/target.decoder.logits_dense.kernel/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..aeeeeb50f8850a4387b6cf911891c0318892d156
Binary files /dev/null and b/t5x/checkpoint_0/target.decoder.logits_dense.kernel/0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.key.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.attention.key.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.attention.key.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.key.kernel/0.0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.attention.key.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..aee6d03a7901dd76a3300f06e006cefab6276fd7
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.attention.key.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.out.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.attention.out.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..afa6346653cc293ffb96cd1bc743cc6910218cd0
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.attention.out.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,6,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,6,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.out.kernel/0.0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.attention.out.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..4899a404a7825dc4b98ec0f022832e5822628e9f
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.attention.out.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.query.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.attention.query.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.attention.query.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.query.kernel/0.0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.attention.query.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..f789508add8fa88faf8ebea7ac108fe6fa89d23a
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.attention.query.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.value.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.attention.value.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..26e2fa9670e96fce101b8f6ac74a04f960def453
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.attention.value.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,4,6],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,4,6],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.attention.value.kernel/0.0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.attention.value.kernel/0.0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..9fe1fc8c8b6efa0803f5d58cd2850c290355181e
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.attention.value.kernel/0.0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_0.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_0.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..cd07a9e840865497f6ff379cb68d73af7e37b719
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_0.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,16],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,16],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_0.kernel/0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_0.kernel/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..0660258bd071a22d7df2545a9cc643041ebb75f6
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_0.kernel/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_1.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_1.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..cd07a9e840865497f6ff379cb68d73af7e37b719
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_1.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3,16],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3,16],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_1.kernel/0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_1.kernel/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..f80444ad894b27fc05e073ac81fed88340906fc8
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.mlp.wi_1.kernel/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.mlp.wo.kernel/.zarray b/t5x/checkpoint_0/target.encoder.encoder.mlp.wo.kernel/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..fb0b55cbaa4249ea483455d5e1565c95cc50bfb7
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.mlp.wo.kernel/.zarray
@@ -0,0 +1 @@
+{"chunks":[16,3,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[16,3,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.mlp.wo.kernel/0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.mlp.wo.kernel/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..9e25498368ce0b12237fb2bbac6e921f70264c1f
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.mlp.wo.kernel/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.pre_attention_layer_norm.scale/.zarray b/t5x/checkpoint_0/target.encoder.encoder.pre_attention_layer_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.pre_attention_layer_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.pre_attention_layer_norm.scale/0.0 b/t5x/checkpoint_0/target.encoder.encoder.pre_attention_layer_norm.scale/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..d97f8728d629c9ec08b533f3adfbb8bda505a038
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.pre_attention_layer_norm.scale/0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.pre_mlp_layer_norm.scale/.zarray b/t5x/checkpoint_0/target.encoder.encoder.pre_mlp_layer_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..94cd03e354db3f7611f0b444f71a3086c0e3a23a
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.pre_mlp_layer_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8,3],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8,3],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.pre_mlp_layer_norm.scale/0.0 b/t5x/checkpoint_0/target.encoder.encoder.pre_mlp_layer_norm.scale/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..d97f8728d629c9ec08b533f3adfbb8bda505a038
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.pre_mlp_layer_norm.scale/0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder.relpos_bias.rel_embedding/.zarray b/t5x/checkpoint_0/target.encoder.encoder.relpos_bias.rel_embedding/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..1b229d8ada2624c8191ad9134dc6f9822d81da64
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder.relpos_bias.rel_embedding/.zarray
@@ -0,0 +1 @@
+{"chunks":[4,3,64],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[4,3,64],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder.relpos_bias.rel_embedding/0.0.0 b/t5x/checkpoint_0/target.encoder.encoder.relpos_bias.rel_embedding/0.0.0
new file mode 100644
index 0000000000000000000000000000000000000000..8ae421e3a3768343e78b3b15832a7d8f2d30f1db
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder.relpos_bias.rel_embedding/0.0.0 differ
diff --git a/t5x/checkpoint_0/target.encoder.encoder_norm.scale/.zarray b/t5x/checkpoint_0/target.encoder.encoder_norm.scale/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..01925abad1a84572db2d6d5531c6c7a7707617a0
--- /dev/null
+++ b/t5x/checkpoint_0/target.encoder.encoder_norm.scale/.zarray
@@ -0,0 +1 @@
+{"chunks":[8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.encoder.encoder_norm.scale/0 b/t5x/checkpoint_0/target.encoder.encoder_norm.scale/0
new file mode 100644
index 0000000000000000000000000000000000000000..034c3e7dac64ee11390a2a9b44a622e413f12dd6
Binary files /dev/null and b/t5x/checkpoint_0/target.encoder.encoder_norm.scale/0 differ
diff --git a/t5x/checkpoint_0/target.token_embedder.embedding/.zarray b/t5x/checkpoint_0/target.token_embedder.embedding/.zarray
new file mode 100644
index 0000000000000000000000000000000000000000..4e8729ee90c3e6de528fe0eec0f1b0a994ebe0d3
--- /dev/null
+++ b/t5x/checkpoint_0/target.token_embedder.embedding/.zarray
@@ -0,0 +1 @@
+{"chunks":[256,8],"compressor":{"id":"gzip","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[256,8],"zarr_format":2}
\ No newline at end of file
diff --git a/t5x/checkpoint_0/target.token_embedder.embedding/0.0 b/t5x/checkpoint_0/target.token_embedder.embedding/0.0
new file mode 100644
index 0000000000000000000000000000000000000000..ceaac054c3b41a2efaa851c03414ab80439cc23a
Binary files /dev/null and b/t5x/checkpoint_0/target.token_embedder.embedding/0.0 differ
diff --git a/t5x/config.gin b/t5x/config.gin
new file mode 100644
index 0000000000000000000000000000000000000000..a94f879b23f6d74c7de78487ceb9cd71add4aaf4
--- /dev/null
+++ b/t5x/config.gin
@@ -0,0 +1,155 @@
+from __gin__ import dynamic_registration
+import __main__ as train_script
+import prot_gin.scalable_t5.scalablet5network as network
+import seqio
+from t5x import adafactor
+from t5x import gin_utils
+from t5x import models
+from t5x import partitioning
+from t5x import trainer
+from t5x import utils
+import tasks
+
+# Macros:
+# ==============================================================================
+BATCH_SIZE = 2
+DROPOUT_RATE = 0.0
+LABEL_SMOOTHING = 0.0
+LOSS_NORMALIZING_FACTOR = None
+MIXTURE_OR_TASK_MODULE = None
+MIXTURE_OR_TASK_NAME = 'test_task'
+MODEL = @models.EncoderDecoderModel()
+MODEL_DIR = \
+    '/tmp/t5x/scalable_tiny_test_pretrain/'
+OPTIMIZER = @adafactor.Adafactor()
+RANDOM_SEED = None
+SHUFFLE_TRAIN_EXAMPLES = True
+TASK_FEATURE_LENGTHS = {'inputs': 512, 'targets': 512}
+TRAIN_STEPS = 6000000
+USE_CACHED_TASKS = False
+USE_HARDWARE_RNG = False
+VOCABULARY = @seqio.SentencePieceVocabulary()
+Z_LOSS = 0.0001
+
+# Parameters for adafactor.Adafactor:
+# ==============================================================================
+adafactor.Adafactor.decay_rate = 0.8
+adafactor.Adafactor.logical_factor_rules = \
+    @adafactor.standard_logical_factor_rules()
+adafactor.Adafactor.skip_nan_updates = True
+adafactor.Adafactor.step_offset = 0
+
+# Parameters for utils.CheckpointConfig:
+# ==============================================================================
+utils.CheckpointConfig.restore = @utils.RestoreCheckpointConfig()
+utils.CheckpointConfig.save = @utils.SaveCheckpointConfig()
+
+# Parameters for utils.create_learning_rate_scheduler:
+# ==============================================================================
+utils.create_learning_rate_scheduler.base_learning_rate = 1.0
+utils.create_learning_rate_scheduler.factors = 'constant * rsqrt_decay'
+utils.create_learning_rate_scheduler.warmup_steps = 10000
+
+# Parameters for train/utils.DatasetConfig:
+# ==============================================================================
+train/utils.DatasetConfig.batch_size = %BATCH_SIZE
+train/utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+train/utils.DatasetConfig.module = %MIXTURE_OR_TASK_MODULE
+train/utils.DatasetConfig.pack = False
+train/utils.DatasetConfig.seed = None
+train/utils.DatasetConfig.shuffle = %SHUFFLE_TRAIN_EXAMPLES
+train/utils.DatasetConfig.split = 'train'
+train/utils.DatasetConfig.task_feature_lengths = %TASK_FEATURE_LENGTHS
+train/utils.DatasetConfig.use_cached = %USE_CACHED_TASKS
+
+# Parameters for train_eval/utils.DatasetConfig:
+# ==============================================================================
+train_eval/utils.DatasetConfig.batch_size = %BATCH_SIZE
+train_eval/utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+train_eval/utils.DatasetConfig.module = %MIXTURE_OR_TASK_MODULE
+train_eval/utils.DatasetConfig.pack = False
+train_eval/utils.DatasetConfig.seed = 42
+train_eval/utils.DatasetConfig.shuffle = False
+train_eval/utils.DatasetConfig.split = 'validation'
+train_eval/utils.DatasetConfig.task_feature_lengths = %TASK_FEATURE_LENGTHS
+train_eval/utils.DatasetConfig.use_cached = %USE_CACHED_TASKS
+
+# Parameters for models.EncoderDecoderModel:
+# ==============================================================================
+models.EncoderDecoderModel.input_vocabulary = %VOCABULARY
+models.EncoderDecoderModel.label_smoothing = %LABEL_SMOOTHING
+models.EncoderDecoderModel.loss_normalizing_factor = %LOSS_NORMALIZING_FACTOR
+models.EncoderDecoderModel.module = @network.Transformer()
+models.EncoderDecoderModel.optimizer_def = %OPTIMIZER
+models.EncoderDecoderModel.output_vocabulary = %VOCABULARY
+models.EncoderDecoderModel.z_loss = %Z_LOSS
+
+# Parameters for partitioning.PjitPartitioner:
+# ==============================================================================
+partitioning.PjitPartitioner.logical_axis_rules = \
+    @partitioning.standard_logical_axis_rules()
+partitioning.PjitPartitioner.model_parallel_submesh = None
+partitioning.PjitPartitioner.num_partitions = 1
+
+# Parameters for utils.RestoreCheckpointConfig:
+# ==============================================================================
+utils.RestoreCheckpointConfig.path = []
+
+# Parameters for utils.SaveCheckpointConfig:
+# ==============================================================================
+utils.SaveCheckpointConfig.dtype = 'float32'
+utils.SaveCheckpointConfig.keep = 3
+utils.SaveCheckpointConfig.period = 40000
+utils.SaveCheckpointConfig.save_dataset = False
+
+# Parameters for seqio.SentencePieceVocabulary:
+# ==============================================================================
+seqio.SentencePieceVocabulary.extra_ids = 225
+seqio.SentencePieceVocabulary.sentencepiece_model_file = 'm.model'
+
+# Parameters for partitioning.standard_logical_axis_rules:
+# ==============================================================================
+partitioning.standard_logical_axis_rules.activation_partitioning_dims = 1
+partitioning.standard_logical_axis_rules.parameter_partitioning_dims = 1
+
+# Parameters for network.T5Config:
+# ==============================================================================
+network.T5Config.dropout_rate = %DROPOUT_RATE
+network.T5Config.dtype = 'bfloat16'
+network.T5Config.emb_dim = 8
+network.T5Config.head_dim = 6
+network.T5Config.logits_via_embedding = False
+network.T5Config.mlp_activations = ('silu', 'linear')
+network.T5Config.mlp_dim = 16
+network.T5Config.num_decoder_layers = 3
+network.T5Config.num_encoder_layers = 3
+network.T5Config.num_heads = 4
+network.T5Config.remat_policy = 'minimal'
+network.T5Config.scan_layers = True
+network.T5Config.vocab_size = 256
+
+# Parameters for train_script.train:
+# ==============================================================================
+train_script.train.checkpoint_cfg = @utils.CheckpointConfig()
+train_script.train.eval_period = 20000
+train_script.train.eval_steps = 200
+train_script.train.infer_eval_dataset_cfg = None
+train_script.train.model = %MODEL
+train_script.train.model_dir = %MODEL_DIR
+train_script.train.partitioner = @partitioning.PjitPartitioner()
+train_script.train.random_seed = %RANDOM_SEED
+train_script.train.summarize_config_fn = @gin_utils.summarize_gin_config
+train_script.train.total_steps = %TRAIN_STEPS
+train_script.train.train_dataset_cfg = @train/utils.DatasetConfig()
+train_script.train.train_eval_dataset_cfg = @train_eval/utils.DatasetConfig()
+train_script.train.trainer_cls = @trainer.Trainer
+train_script.train.use_hardware_rng = %USE_HARDWARE_RNG
+
+# Parameters for trainer.Trainer:
+# ==============================================================================
+trainer.Trainer.learning_rate_fn = @utils.create_learning_rate_scheduler()
+trainer.Trainer.num_microbatches = 1
+
+# Parameters for network.Transformer:
+# ==============================================================================
+network.Transformer.config = @network.T5Config()
diff --git a/t5x/model-info.txt b/t5x/model-info.txt
new file mode 100644
index 0000000000000000000000000000000000000000..2d5f7bbe3c9470a2250a328148dfe19a7f7513ef
--- /dev/null
+++ b/t5x/model-info.txt
@@ -0,0 +1,148 @@
+Variable decoder/decoder/encoder_decoder_attention/key/kernel                             size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable decoder/decoder/encoder_decoder_attention/out/kernel                             size 576          shape (heads=4, layers=3, kv=6, embed=8)       partition spec ('model', None, None, None)
+Variable decoder/decoder/encoder_decoder_attention/query/kernel                           size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable decoder/decoder/encoder_decoder_attention/value/kernel                           size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable decoder/decoder/mlp/wi_0/kernel                                                  size 384          shape (embed=8, layers=3, mlp=16)              partition spec (None, None, 'model')
+Variable decoder/decoder/mlp/wi_1/kernel                                                  size 384          shape (embed=8, layers=3, mlp=16)              partition spec (None, None, 'model')
+Variable decoder/decoder/mlp/wo/kernel                                                    size 384          shape (mlp=16, layers=3, embed=8)              partition spec ('model', None, None)
+Variable decoder/decoder/pre_cross_attention_layer_norm/scale                             size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable decoder/decoder/pre_mlp_layer_norm/scale                                         size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable decoder/decoder/pre_self_attention_layer_norm/scale                              size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable decoder/decoder/relpos_bias/rel_embedding                                        size 768          shape (heads=4, layers=3, relpos_buckets=64)   partition spec ('model', None, None)
+Variable decoder/decoder/self_attention/key/kernel                                        size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable decoder/decoder/self_attention/out/kernel                                        size 576          shape (heads=4, layers=3, kv=6, embed=8)       partition spec ('model', None, None, None)
+Variable decoder/decoder/self_attention/query/kernel                                      size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable decoder/decoder/self_attention/value/kernel                                      size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable decoder/decoder_norm/scale                                                       size 8            shape (embed=8)                                partition spec (None,)
+Variable decoder/logits_dense/kernel                                                      size 2048         shape (embed=8, vocab=256)                     partition spec (None, 'model')
+Variable encoder/encoder/attention/key/kernel                                             size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable encoder/encoder/attention/out/kernel                                             size 576          shape (heads=4, layers=3, kv=6, embed=8)       partition spec ('model', None, None, None)
+Variable encoder/encoder/attention/query/kernel                                           size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable encoder/encoder/attention/value/kernel                                           size 576          shape (embed=8, layers=3, heads=4, kv=6)       partition spec (None, None, 'model', None)
+Variable encoder/encoder/mlp/wi_0/kernel                                                  size 384          shape (embed=8, layers=3, mlp=16)              partition spec (None, None, 'model')
+Variable encoder/encoder/mlp/wi_1/kernel                                                  size 384          shape (embed=8, layers=3, mlp=16)              partition spec (None, None, 'model')
+Variable encoder/encoder/mlp/wo/kernel                                                    size 384          shape (mlp=16, layers=3, embed=8)              partition spec ('model', None, None)
+Variable encoder/encoder/pre_attention_layer_norm/scale                                   size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable encoder/encoder/pre_mlp_layer_norm/scale                                         size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable encoder/encoder/relpos_bias/rel_embedding                                        size 768          shape (heads=4, layers=3, relpos_buckets=64)   partition spec ('model', None, None)
+Variable encoder/encoder_norm/scale                                                       size 8            shape (embed=8)                                partition spec (None,)
+Variable token_embedder/embedding                                                         size 2048         shape (vocab=256, embed=8)                     partition spec ('model', None)
+Total number of parameters: 14984
+
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/m              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/v              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/v_col          size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/v_row          size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/m              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/v              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/v_col          size 72           shape (4, 3, 6)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/v_row          size 24           shape (3, 8)                                   partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/m            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/v            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/v_col        size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/v_row        size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/m            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/v            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/v_col        size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/v_row        size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/v_col                               size 48           shape (3, 16)                                  partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/v_row                               size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/v_col                               size 48           shape (3, 16)                                  partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/v_row                               size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/m                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/v                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/v_col                                 size 48           shape (16, 3)                                  partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/v_row                                 size 24           shape (3, 8)                                   partition spec None
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/m              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/v              size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/v_col          size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/v_row          size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/m                          size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/v                          size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/v_col                      size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/v_row                      size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/m               size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/v               size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/v_col           size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/v_row           size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/v                         size 768          shape (heads=4, layers=3, relpos_buckets=64)   partition spec ('model', None, None)
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/v_col                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/v_row                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/v                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/v_col                     size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/v_row                     size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/v                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/v_col                     size 72           shape (4, 3, 6)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/v_row                     size 24           shape (3, 8)                                   partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/m                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/v                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/v_col                   size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/v_row                   size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/m                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/v                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/v_col                   size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/v_row                   size 24           shape (8, 3)                                   partition spec None
+Variable param_states/decoder/decoder_norm/scale/m                                        size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder_norm/scale/v                                        size 8            shape (embed=8)                                partition spec (None,)
+Variable param_states/decoder/decoder_norm/scale/v_col                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder_norm/scale/v_row                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/logits_dense/kernel/m                                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/logits_dense/kernel/v                                       size 2048         shape (embed=8, vocab=256)                     partition spec (None, 'model')
+Variable param_states/decoder/logits_dense/kernel/v_col                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/logits_dense/kernel/v_row                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/m                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/v                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/v_col                          size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/v_row                          size 24           shape (8, 3)                                   partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/m                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/v                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/v_col                          size 72           shape (4, 3, 6)                                partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/v_row                          size 24           shape (3, 8)                                   partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/m                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/v                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/v_col                        size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/v_row                        size 24           shape (8, 3)                                   partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/m                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/v                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/v_col                        size 72           shape (3, 4, 6)                                partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/v_row                        size 24           shape (8, 3)                                   partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/v_col                               size 48           shape (3, 16)                                  partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/v_row                               size 24           shape (8, 3)                                   partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/v_col                               size 48           shape (3, 16)                                  partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/v_row                               size 24           shape (8, 3)                                   partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/m                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/v                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/v_col                                 size 48           shape (16, 3)                                  partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/v_row                                 size 24           shape (3, 8)                                   partition spec None
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/m                    size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/v                    size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/v_col                size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/v_row                size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/m                          size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/v                          size 24           shape (embed=8, layers=3)                      partition spec (None, None)
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/v_col                      size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/v_row                      size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/v                         size 768          shape (heads=4, layers=3, relpos_buckets=64)   partition spec ('model', None, None)
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/v_col                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/v_row                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder_norm/scale/m                                        size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder_norm/scale/v                                        size 8            shape (embed=8)                                partition spec (None,)
+Variable param_states/encoder/encoder_norm/scale/v_col                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder_norm/scale/v_row                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/token_embedder/embedding/m                                          size 1            shape (1,)                                     partition spec None
+Variable param_states/token_embedder/embedding/v                                          size 2048         shape (vocab=256, embed=8)                     partition spec ('model', None)
+Variable param_states/token_embedder/embedding/v_col                                      size 1            shape (1,)                                     partition spec None
+Variable param_states/token_embedder/embedding/v_row                                      size 1            shape (1,)                                     partition spec None
+Variable step                                                                             size 1            shape ()                                       partition spec None
diff --git a/t5x/train/events.out.tfevents.1680673723.AMC-C02G81B6MD6R.30566.0.v2 b/t5x/train/events.out.tfevents.1680673723.AMC-C02G81B6MD6R.30566.0.v2
new file mode 100644
index 0000000000000000000000000000000000000000..2bfd2e57cd08de419656993b3491ab2468a4f8b7
--- /dev/null
+++ b/t5x/train/events.out.tfevents.1680673723.AMC-C02G81B6MD6R.30566.0.v2
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8340acca06eb772b28d360e2243c5bbe47645026450ce990c56365130877f504
+size 6487
diff --git a/t5x/training_eval/uniref50_pretrain_ul2/events.out.tfevents.1680673723.AMC-C02G81B6MD6R.30566.1.v2 b/t5x/training_eval/uniref50_pretrain_ul2/events.out.tfevents.1680673723.AMC-C02G81B6MD6R.30566.1.v2
new file mode 100644
index 0000000000000000000000000000000000000000..cb4997b748ce781895b16041422193548f899d3e
--- /dev/null
+++ b/t5x/training_eval/uniref50_pretrain_ul2/events.out.tfevents.1680673723.AMC-C02G81B6MD6R.30566.1.v2
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9568cfbf5d7d4f8c87353c3ef3b738bff84f388496179cc1002054a8b4835a6a
+size 40